Initial Wasm runner implementation (#173)

This PR adds a WASM runner, which can run WASM models compiled using the interface (subject to change #175) defined in ```../carton-runner-wasm/wit/lib.wit```. The existing implementation is still unoptimized, requiring 2 copies per Tensor moved to/from WASM. An example of compiling a compatible model can be found in ```carton-runner-wasm/tests/test_model```.

## Limitations
- Only the ```wasm32-unknown-unknown``` target has been tested to be working.
- Only ```infer``` is supported for now.
- Packing only supports a single ```.wasm``` file and no other artifacts.
- No WebGPU, and probably not for a while.

## Test Coverage
All type conversions from Carton to WASM and vice versa and fully covered. Pack, Load, Infer are covered in pack.rs.

## TODOs
Track in #164

.gitignore

@@ -1,3 +1,5 @@
 target/
 .vscode/
 libtorch/
+.idea/
+*.DS_Store

Cargo.toml

@@ -15,6 +15,7 @@ members = [
     "source/carton-runner-py",
     "source/carton-runner-rust-bert",
     "source/carton-runner-torch",
+    "source/carton-runner-wasm",
     "source/carton-utils-py",
     "source/anywhere",
     "source/fetch-deps",

ci/build.py

@@ -116,6 +116,7 @@ def update_version_numbers():
         run_command(["cargo", "run", RELEASE_FLAG, "--timings", "--target", TARGET, "-p", "carton-runner-py", "--bin", "build_releases", "--", "--output-path", args.runner_release_dir])
         run_command(["cargo", "run", RELEASE_FLAG, "--timings", "--target", TARGET, "-p", "carton-runner-rust-bert", "--bin", "build_rust_bert_releases", "--", "--output-path", args.runner_release_dir])
         run_command(["cargo", "run", RELEASE_FLAG, "--timings", "--target", TARGET, "-p", "carton-runner-torch", "--bin", "build_torch_releases", "--", "--output-path", args.runner_release_dir])
+        run_command(["cargo", "run", RELEASE_FLAG, "--timings", "--target", TARGET, "-p", "carton-runner-wasm", "--bin", "build_wasm_releases", "--", "--output-path", args.runner_release_dir])
 
     # Show sccache stats
     RUSTC_WRAPPER = os.getenv("RUSTC_WRAPPER", "")

source/carton-runner-wasm/Cargo.toml

@@ -0,0 +1,31 @@
+[package]
+name = "carton-runner-wasm"
+version = "0.0.1"
+edition = "2021"
+publish = false
+
+exclude = ["tests/test_model", "carton-wasm-interface"]
+
+[dependencies]
+carton = { path = "../carton" }
+carton-runner-interface = { path = "../carton-runner-interface" }
+color-eyre = "0.6.2"
+lunchbox = { version = "0.1", default-features = false }
+wasmtime = { version = "13.0.0", features = ["component-model"] }
+tokio = "1.32.0"
+ndarray = "0.15.6"
+
+# Used by the `build_releases` binary
+escargot = "0.5.8"
+carton-runner-packager = { path = "../carton-runner-packager" }
+clap = { version = "4.4.6", features = ["derive"] }
+env_logger = "0.10.0"
+log = "0.4.20"
+target-lexicon = "0.12.11"
+serde_json = "1.0.107"
+semver = "1.0.20"
+
+[dev-dependencies]
+escargot = "0.5.8"
+paste = "1.0.14"
+tempfile = "3.8.0"

source/carton-runner-wasm/README.md

@@ -0,0 +1,5 @@
+# WASM Runner (Experimental)
+This runner is capable of running models that implement the interface defined in 'carton-runner-wasm/wit/lib.wit'.
+
+## Disclaimer
+The defined interface is subject to change, and backwards compatability is not guaranteed, while experimental.

source/carton-runner-wasm/src/bin/build_wasm_releases.rs

@@ -0,0 +1,76 @@
+//! Same as torch runner
+
+use std::{path::PathBuf, time::SystemTime};
+
+use clap::Parser;
+
+use carton_runner_interface::slowlog::slowlog;
+use carton_runner_packager::discovery::RunnerInfo;
+
+// TODO: This should be the version of carton-interface-wasm, but it's not done yet.
+const INTERFACE_VERSION: semver::Version = semver::Version::new(0, 0, 1);
+
+#[derive(Parser, Debug)]
+struct Args {
+    #[arg(long)]
+    output_path: PathBuf,
+}
+
+#[tokio::main]
+async fn main() {
+    // Logging (for long running downloads)
+    env_logger::Builder::from_env(env_logger::Env::default().default_filter_or("info")).init();
+
+    // Parse args
+    let args = Args::parse();
+
+    log::info!("Starting runner build...");
+    let mut sl = slowlog("Building runner", 5).await.without_progress();
+    // Build the runner
+    let runner_path = escargot::CargoBuild::new()
+        .package("carton-runner-wasm")
+        .bin("carton-runner-wasm")
+        .current_release()
+        .current_target()
+        .arg("--timings")
+        .run()
+        .unwrap()
+        .path()
+        .display()
+        .to_string();
+
+    sl.done();
+    log::info!("Runner Path: {}", runner_path);
+
+    let package = carton_runner_packager::package(
+        RunnerInfo {
+            runner_name: "wasm".to_string(),
+            framework_version: INTERFACE_VERSION,
+            runner_compat_version: 1,
+            runner_interface_version: 1,
+            runner_release_date: SystemTime::now().into(),
+            runner_path,
+            platform: target_lexicon::HOST.to_string(),
+        },
+        vec![],
+    )
+    .await;
+
+    // Write the zip file to our output dir
+    tokio::fs::write(
+        &args
+            .output_path
+            .join(format!("{}.zip", package.get_data_sha256())),
+        package.get_data(),
+    )
+    .await
+    .unwrap();
+
+    // Write the package config so it can be loaded when the runner zip files will be uploaded
+    tokio::fs::write(
+        &args.output_path.join(format!("{}.json", package.get_id())),
+        serde_json::to_string_pretty(&package).unwrap(),
+    )
+    .await
+    .unwrap();
+}

source/carton-runner-wasm/src/component.rs

@@ -0,0 +1,11 @@
+wasmtime::component::bindgen!({
+    world: "model",
+    path: "./wit",
+});
+
+use crate::component::carton_wasm::lib::types::Host;
+pub(crate) use carton_wasm::lib::types::{Dtype, TensorNumeric, TensorString};
+
+pub(crate) struct HostImpl;
+
+impl Host for HostImpl {}

source/carton-runner-wasm/src/lib.rs

@@ -0,0 +1,58 @@
+use std::collections::HashMap;
+
+use color_eyre::eyre::{eyre, Result};
+use wasmtime::component::{Component, Linker};
+use wasmtime::{Engine, Store};
+
+use carton_runner_interface::types::Tensor as CartonTensor;
+
+use crate::component::{HostImpl, Model, Tensor};
+
+mod component;
+mod types;
+
+pub struct WASMModelInstance {
+    store: Store<HostImpl>,
+    model: Model,
+}
+
+impl WASMModelInstance {
+    pub fn from_bytes(engine: &Engine, bytes: &[u8]) -> Result<Self> {
+        /*
+        see https://docs.wasmtime.dev/api/wasmtime/component/macro.bindgen.html
+        Some of the names may be confusing, here is the general idea from my
+        understanding:
+        - HostImpl is the host side implementation of what a interface imports
+          since our current interface does not import anything, this is an empty
+          struct
+        - Model is the loaded and linked interface, i.e. the API we expect the
+          user to implement. (Non stateful)
+          TODO: rename to ModelInterface
+         */
+        let comp = Component::from_binary(&engine, bytes).unwrap();
+        let mut linker = Linker::<HostImpl>::new(&engine);
+        Model::add_to_linker(&mut linker, |state: &mut HostImpl| state).unwrap();
+        let mut store = Store::new(&engine, HostImpl);
+        let (model, _) = Model::instantiate(&mut store, &comp, &linker).unwrap();
+        Ok(Self { store, model })
+    }
+
+    pub fn infer(
+        &mut self,
+        inputs: HashMap<String, CartonTensor>,
+    ) -> Result<HashMap<String, CartonTensor>> {
+        let inputs = inputs
+            .into_iter()
+            .map(|(k, v)| Ok((k, v.try_into()?)))
+            .collect::<Result<Vec<(String, Tensor)>>>()?;
+        let outputs = self
+            .model
+            .call_infer(&mut self.store, inputs.as_ref())
+            .map_err(|e| eyre!(e))?;
+        let mut ret = HashMap::new();
+        for (k, v) in outputs.into_iter() {
+            ret.insert(k, v.into());
+        }
+        Ok(ret)
+    }
+}

source/carton-runner-wasm/src/main.rs

@@ -0,0 +1,75 @@
+use color_eyre::eyre::{eyre, Result};
+use lunchbox::{path::Path, types::WritableFileSystem, ReadableFileSystem};
+use wasmtime::{Config, Engine};
+
+use carton_runner_interface::server::{init_runner, RequestData, ResponseData};
+use carton_runner_wasm::WASMModelInstance;
+
+fn new_engine() -> Result<Engine> {
+    let mut config = Config::new();
+    config.wasm_component_model(true);
+    Engine::new(&config).map_err(|e| eyre!(e))
+}
+
+#[tokio::main]
+async fn main() {
+    color_eyre::install().unwrap();
+    let mut server = init_runner().await;
+    let engine = new_engine().unwrap();
+    let mut model: Option<WASMModelInstance> = None;
+
+    while let Some(req) = server.get_next_request().await {
+        let req_id = req.id;
+        match req.data {
+            RequestData::Load { fs, .. } => {
+                let fs = server.get_readonly_filesystem(fs).await.unwrap();
+                let bin = &fs.read("model.wasm").await.unwrap();
+                model = Some(
+                    WASMModelInstance::from_bytes(&engine, bin)
+                        .expect("Failed to initialize WASM model"),
+                );
+                server
+                    .send_response_for_request(req_id, ResponseData::Load)
+                    .await
+                    .unwrap();
+            }
+            RequestData::Pack {
+                input_path,
+                temp_folder,
+                fs,
+            } => {
+                let fs = server.get_writable_filesystem(fs).await.unwrap();
+                fs.symlink(input_path, Path::new(&temp_folder).join("model.wasm"))
+                    .await
+                    .unwrap();
+                server
+                    .send_response_for_request(
+                        req_id,
+                        ResponseData::Pack {
+                            output_path: temp_folder,
+                        },
+                    )
+                    .await
+                    .unwrap();
+            }
+            RequestData::Seal { .. } => {
+                todo!()
+            }
+            RequestData::InferWithTensors { tensors, .. } => {
+                let result = model.as_mut().map(|m| m.infer(tensors)).unwrap();
+                server
+                    .send_response_for_request(
+                        req_id,
+                        ResponseData::Infer {
+                            tensors: result.unwrap(),
+                        },
+                    )
+                    .await
+                    .unwrap();
+            }
+            RequestData::InferWithHandle { .. } => {
+                todo!()
+            }
+        }
+    }
+}