feat(sdk, gate, cli): upload protocol poc uploading wasm file for WasmEdge Runtime for single replica mode (#631)

Upload protocol for wasm files and atrifacts for `WasmEdge Runtime` for
single replica mode

#### Motivation and context

- Upload WasmEdge Runtime artifacts during typegraph deploy
- Access and load WasmEdge Runtime artifacts from the local file system
from typegate

#### Migration notes

*No Migrations Needed*

### Checklist

- [x] The change come with new or modified tests
- [ ] Hard-to-understand functions have explanatory comments
- [ ] End-user documentation is updated to reflect the change

---------

Signed-off-by: Estifanos Bireda <[email protected]>
Co-authored-by: afmika <[email protected]>
Co-authored-by: Yohe-Am <[email protected]>

Author: 3 people

.github/workflows/tests.yml

@@ -13,6 +13,7 @@ env:
   DENO_VERSION: "1.41.0"
   RUST_BACKTRACE: "full"
   RUST_LOG: "info,swc_ecma_codegen=off,tracing::span=off"
+  DENO_DIR: deno-dir
 
 jobs:
   changes:
@@ -151,7 +152,7 @@ jobs:
           # temporary fix
           cache-key-prefix: ${{ matrix.os }}
       - shell: bash
-        env: 
+        env:
           WASM_FILE: target/debug/typegraph_core.wasm
         run: |
           python3 -m venv .venv
@@ -184,6 +185,8 @@ jobs:
         include:
           - platform: linux/amd64
             runner: ubuntu-latest
+          # - platform: linux/amd64
+          #   runner: custom-ubuntu-large
           # FIXME: try macos-14 runner once all actions support it
           # docker buildx action broken as of 2024-02-09
 
@@ -290,14 +293,14 @@ jobs:
           path: .venv
           key: ${{ runner.os }}-venv-${{ hashFiles('**/poetry.lock', '.ghjk/lock.json') }}
       # FIXME: the custom runner is missing some ambient items found in
-      # github runner images by default. remove this hack when ghjk handles 
+      # github runner images by default. remove this hack when ghjk handles
       # zstd/bsdtar on it's own
       - run: |
           sudo apt update
           sudo apt install -y --no-install-recommends zstd
       - uses: metatypedev/setup-ghjk@a7bbf22959e3c0f80b8ba9f800b0a9f1ff17fa7b
       - shell: bash
-        env: 
+        env:
           WASM_FILE: target/debug/typegraph_core.wasm
         run: |
           echo pwd in $(pwd)
@@ -345,6 +348,7 @@ jobs:
           deno cache --import-map typegate/import_map.json \
             typegate/src/main.ts \
             typegate/tests/utils/*.ts \
+            typegate/tests/runtimes/wasmedge/*.ts \
             dev/deps.ts \
             dev/utils.ts
           deno --unstable-worker-options --unstable-net coverage ./coverage --lcov > coverage.lcov

docs/workflows/upload_protocol.svg

@@ -10,6 +10,7 @@ pub mod visitor;
 
 pub use types::*;
 
+use std::collections::HashMap;
 use std::path::{Path, PathBuf};
 use std::sync::Arc;
 
@@ -99,6 +100,7 @@ pub struct TypeMeta {
     pub rate: Option<Rate>,
     pub version: String,
     pub random_seed: Option<u32>,
+    pub ref_artifacts: HashMap<String, PathBuf>,
 }
 
 #[cfg_attr(feature = "codegen", derive(JsonSchema))]

libs/common/src/typegraph/mod.rs

@@ -10,6 +10,8 @@ use serde::{Deserialize, Serialize};
 pub struct WasiMatData {
     pub wasm: String,
     pub func: String,
+    pub artifact_hash: String,
+    pub tg_name: Option<String>,
 }
 
 #[cfg_attr(feature = "codegen", derive(JsonSchema))]

libs/common/src/typegraph/runtimes/wasmedge.rs

@@ -10,14 +10,15 @@ use anyhow::bail;
 #[rustfmt::skip]
 use deno_core as deno_core; // necessary for re-exported macros to work
 
-use base64::{engine::general_purpose, Engine as _};
 use wasmedge_sdk::VmBuilder;
 use wasmedge_sdk::{
     config::{ConfigBuilder, HostRegistrationConfigOptions},
     dock::{Param, VmDock},
     Module,
 };
 
+use std::{env, fs};
+
 #[derive(Deserialize)]
 #[serde(crate = "serde")]
 pub struct WasiInput {
@@ -54,9 +55,14 @@ fn param_cast(out: &str, res: &mut Vec<Box<dyn Any + Send + Sync>>) -> Result<St
 pub fn op_wasmedge_wasi(#[serde] input: WasiInput) -> Result<String> {
     // https://github.com/second-state/wasmedge-rustsdk-examples
 
-    let bytes = general_purpose::STANDARD
-        .decode(input.wasm.as_bytes())
-        .unwrap();
+    let wasm_relative_path = PathBuf::from(input.wasm);
+
+    let wasm_absolute_path = match env::current_dir() {
+        Ok(cwd) => cwd.join(wasm_relative_path),
+        Err(e) => return Err(anyhow::anyhow!(e)),
+    };
+
+    let bytes = fs::read(wasm_absolute_path).unwrap();
     let module = Module::from_bytes(None, bytes).unwrap();
 
     let config = ConfigBuilder::default()

typegate/engine/src/runtimes/wasmedge.rs

@@ -7,6 +7,7 @@ import { Resolver, RuntimeInitParams } from "../types.ts";
 import { nativeResult } from "../utils.ts";
 import { ComputeStage } from "../engine/query_engine.ts";
 import { registerRuntime } from "./mod.ts";
+import config from "../config.ts";
 
 @registerRuntime("wasmedge")
 export class WasmEdgeRuntime extends Runtime {
@@ -29,7 +30,7 @@ export class WasmEdgeRuntime extends Runtime {
     _verbose: boolean,
   ): ComputeStage[] {
     const { materializer, argumentTypes, outType } = stage.props;
-    const { wasm, func } = materializer?.data ?? {};
+    const { wasm, func, artifact_hash, tg_name } = materializer?.data ?? {};
     const order = Object.keys(argumentTypes ?? {});
 
     // always wasi
@@ -40,7 +41,8 @@ export class WasmEdgeRuntime extends Runtime {
         await native.wasmedge_wasi(
           {
             func: func as string,
-            wasm: wasm as string,
+            wasm:
+              `${config.tmp_dir}/metatype_artifacts/${tg_name as string}/artifacts/${wasm as string}.${artifact_hash as string}`,
             args: transfert,
             out: outType.type,
           },

typegate/src/runtimes/wasmedge.ts

@@ -0,0 +1,120 @@
+// Copyright Metatype OÜ, licensed under the Elastic License 2.0.
+// SPDX-License-Identifier: Elastic-2.0
+
+import config from "../config.ts";
+import { signJWT, verifyJWT } from "../crypto.ts";
+import { UploadUrlMeta } from "../typegate/mod.ts";
+
+interface TypegraphArtifact {
+  name: string;
+  artifact_hash: string;
+  artifact_size_in_bytes: number;
+}
+
+function createUploadPath(origin: string, typegraphName: string) {
+  const rand_path = crypto.randomUUID();
+  return `${origin}/${typegraphName}/upload-artifacts/artifacts/${rand_path}`;
+}
+
+export async function handleUploadUrl(
+  request: Request,
+  tgName: string,
+  urlCache: Map<string, UploadUrlMeta>,
+) {
+  const url = new URL(request.url);
+  const origin = url.origin;
+  const { name, artifact_hash, artifact_size_in_bytes }: TypegraphArtifact =
+    await request
+      .json();
+  const uploadUrlMeta: UploadUrlMeta = {
+    artifactName: name,
+    artifactHash: artifact_hash,
+    artifactSizeInBytes: artifact_size_in_bytes,
+    urlUsed: false,
+  };
+
+  let uploadUrl = createUploadPath(origin, tgName);
+
+  const expiresIn = 5 * 60; // 5 minutes
+  const payload = {
+    "expiresIn": expiresIn,
+  };
+  const token = await signJWT(payload, expiresIn);
+  uploadUrl = `${uploadUrl}?token=${token}`;
+
+  urlCache.set(uploadUrl, uploadUrlMeta);
+
+  return new Response(JSON.stringify({ uploadUrl: uploadUrl }));
+}
+
+export async function handleArtifactUpload(
+  request: Request,
+  tgName: string,
+  urlCache: Map<string, UploadUrlMeta>,
+) {
+  const url = new URL(request.url);
+  if (request.method !== "PUT") {
+    throw new Error(
+      `${url.pathname} does not support ${request.method} method`,
+    );
+  }
+
+  const uploadMeta = urlCache.get(url.toString());
+
+  if (!uploadMeta) {
+    throw new Error(`Endpoint ${url.toString()} does not exist`);
+  }
+
+  const token = url.searchParams.get("token");
+  try {
+    const _ = await verifyJWT(token!);
+  } catch (e) {
+    throw new Error("Invalid token: " + e.toString());
+  }
+
+  const { artifactName, artifactHash, artifactSizeInBytes, urlUsed }:
+    UploadUrlMeta = uploadMeta!;
+
+  if (urlUsed) {
+    throw new Error(`Endpoint ${url.toString()} is disabled`);
+  }
+
+  const reader = request.body?.getReader()!;
+
+  let artifactData = new Uint8Array();
+  let bytesRead = 0;
+  while (true) {
+    const { done, value } = await reader.read();
+    if (done) {
+      break;
+    }
+    if (value) {
+      bytesRead += value.length;
+      const temp = new Uint8Array(artifactData.length + value.length);
+      temp.set(artifactData);
+      temp.set(value, artifactData.length);
+      artifactData = temp;
+    }
+  }
+
+  if (bytesRead !== artifactSizeInBytes) {
+    throw new Error(
+      `File size does not match ${bytesRead}, ${JSON.stringify(uploadMeta)}`,
+    );
+  }
+
+  // adjust relative to the root path
+  const artifactStorageDir =
+    `${config.tmp_dir}/metatype_artifacts/${tgName}/artifacts`;
+  await Deno.mkdir(artifactStorageDir, { recursive: true });
+  const artifactPath = `${artifactStorageDir}/${artifactName}.${artifactHash}`;
+  await Deno.writeFile(artifactPath, artifactData);
+
+  // mark as the url used once the request completes.
+  uploadMeta.urlUsed = true;
+  urlCache.set(url.toString(), uploadMeta);
+
+  return new Response(JSON.stringify({
+    "success": true,
+  }));
+}

typegate/src/services/artifact_upload_service.ts

@@ -35,6 +35,10 @@ import { MigrationFailure } from "../runtimes/prisma/hooks/run_migrations.ts";
 import introspectionJson from "../typegraphs/introspection.json" with {
   type: "json",
 };
+import {
+  handleArtifactUpload,
+  handleUploadUrl,
+} from "../services/artifact_upload_service.ts";
 
 const INTROSPECTION_JSON_STR = JSON.stringify(introspectionJson);
 
@@ -58,8 +62,16 @@ export type PushResult = {
   response: PushResponse;
 };
 
+export interface UploadUrlMeta {
+  artifactName: string;
+  artifactHash: string;
+  artifactSizeInBytes: number;
+  urlUsed: boolean;
+}
+
 export class Typegate {
   #onPushHooks: PushHandler[] = [];
+  artifactUploadUrlCache: Map<string, UploadUrlMeta> = new Map();
 
   constructor(
     public readonly register: Register,
@@ -113,6 +125,24 @@ export class Typegate {
       }
 
       const [engineName, serviceName] = parsePath(url.pathname);
+
+      // artifact upload handlers
+      if (serviceName === "get-upload-url") {
+        return handleUploadUrl(
+          request,
+          engineName,
+          this.artifactUploadUrlCache,
+        );
+      }
+
+      if (serviceName === "upload-artifacts") {
+        return handleArtifactUpload(
+          request,
+          engineName,
+          this.artifactUploadUrlCache,
+        );
+      }
+
       if (!engineName || ignoreList.has(engineName)) {
         return notFound();
       }

typegate/src/typegate/mod.ts

@@ -507,6 +507,7 @@ export interface TypeMeta {
   rate?: Rate | null;
   version: string;
   random_seed: number | undefined;
+  ref_artifacts: Map<string, string>;
 }
 export interface Queries {
   dynamic: boolean;

typegate/src/typegraph/types.ts

@@ -597,6 +597,7 @@
     },
     "auths": [],
     "rate": null,
+    "ref_artifacts": {},
     "version": "0.0.3",
     "random_seed": null
   }