Initial Wasm runner implementation

.gitignore

@@ -1,3 +1,4 @@
 target/
 .vscode/
 libtorch/
+.idea/

Cargo.toml

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

source/carton-runner-interface/src/do_not_modify/storage.rs

@@ -60,4 +60,24 @@ where
         let data = self.data.as_mut_ptr();
         unsafe { ndarray::ArrayViewMut::from_shape_ptr(self.get_shape(), data) }
     }
+
+    pub unsafe fn into_bytes(self) -> (Vec<u8>, Vec<u64>) {
+        // TODO: unsafe because not guaranteed to be numeric, maybe move elsewhere.
+        let len = self.strides.iter().product() as usize * std::mem::size_of::<T>();
+        let buffer = unsafe { Vec::from_raw_parts(
+            self.data.as_ptr() as *mut u8,
+            len,
+            len,
+        ) };
+        (buffer, self.shape)
+    }
+
+    pub fn into_vec(self) -> (Vec<T>, Vec<u64>) {
+        let len = self.strides.iter().product() as usize;
+        let buffer = unsafe { Vec::from_raw_parts(
+            self.data.as_ptr() as *mut T,
+            len,
+            len,
+        ) };
+    }
 }

source/carton-runner-wasm/Cargo.toml

@@ -0,0 +1,23 @@
+[package]
+name = "carton-runner-wasm"
+version = "0.0.1"
+edition = "2021"
+publish = false
+
+# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
+
+exclude = ["tests/basic_model"]
+
+[dependencies]
+carton = { path = "../carton" }
+carton-runner-interface = { path = "../carton-runner-interface" }
+color-eyre = "0.6.2"
+lunchbox = { version = "0.1", features = ["serde"], default-features = false }
+wasmtime = { version = "13.0.0", features = ["component-model"] }
+tokio = "1.32.0"
+serde = "1.0.188"
+serde_json = "1.0.107"
+ndarray = "0.15.6"
+
+[dev-dependencies]
+tempfile = "3.8.0"

source/carton-runner-wasm/build_interface.sh

@@ -0,0 +1,3 @@
+cargo build -p carton-interface-wasm --target wasm32-unknown-unknown
+wasm-tools component new ./target/wasm32-unknown-unknown/debug/carton_interface_wasm.wasm \
+    -o carton-component.wasm

source/carton-runner-wasm/carton-interface-wasm/Cargo.toml

@@ -0,0 +1,10 @@
+[package]
+name = "carton-interface-wasm"
+version = "0.0.1"
+edition = "2021"
+
+[dependencies]
+wit-bindgen = "0.12.0"
+
+[lib]
+crate-type = ["cdylib"]

source/carton-runner-wasm/carton-interface-wasm/src/candle.rs

@@ -0,0 +1,3 @@
+use crate::component::{Tensor, TensorNumeric, TensorString, Dtype};
+
+

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

@@ -0,0 +1,8 @@
+wit_bindgen::generate!({
+	world: "model",
+	exports: {
+		world: Model
+	}
+});
+
+pub use carton_wasm::lib::types::*;

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

@@ -0,0 +1,4 @@
+pub use component::{Tensor, TensorNumeric, TensorString, Dtype};
+
+pub mod candle;
+mod component;

source/carton-runner-wasm/carton-interface-wasm/wit/lib.wit

@@ -0,0 +1,34 @@
+package carton-wasm:lib
+
+interface types {
+    enum dtype {
+        f32,
+        f64,
+        i8,
+        i16,
+        i32,
+        i64,
+        ui8,
+        ui16,
+        ui32,
+        ui64,
+    }
+    record tensor-numeric {
+        buffer: list<u8>,
+        dtype: dtype,
+        shape: list<u64>,
+    }
+    record tensor-string {
+        buffer: list<string>,
+        shape: list<u64>,
+    }
+    variant tensor {
+        numeric(tensor-numeric),
+        str(tensor-string),
+    }
+}
+
+world model {
+    use types.{tensor, tensor-numeric, tensor-string, dtype};
+    export infer: func(in: list<tuple<string, tensor>>) -> result<list<tuple<string, tensor>>, string>;
+}

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

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

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

@@ -0,0 +1,46 @@
+use std::collections::HashMap;
+
+use color_eyre::eyre::{eyre, Result};
+use wasmtime::{Engine, Store};
+use wasmtime::component::*;
+
+use carton_runner_interface::types::Tensor as CartonTensor;
+
+use crate::component::{DummyState, Model, Tensor};
+
+mod component;
+mod types;
+
+pub struct WASMModelInstance {
+    instance: Instance,
+    store: Store<DummyState>,
+    model: Model,
+}
+
+impl WASMModelInstance {
+    pub fn from_bytes(engine: &Engine, bytes: &[u8]) -> Result<Self> {
+        let comp = Component::from_binary(&engine, bytes).unwrap();
+        let mut linker = Linker::<DummyState>::new(&engine);
+        Model::add_to_linker(&mut linker, |state: &mut DummyState| state).unwrap();
+        let mut store = Store::new(&engine, DummyState);
+        let (model, instance) = Model::instantiate(&mut store, &comp, &linker).unwrap();
+        Ok(Self {
+            instance,
+            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,77 @@
+use color_eyre::eyre::{eyre, Result};
+use lunchbox::{path::Path, ReadableFileSystem, types::WritableFileSystem};
+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, runner_opts, .. } => {
+				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))
+                    .await
+                    .unwrap();
+				server
+                    .send_response_for_request(
+                        req_id,
+                        ResponseData::Pack {
+                            output_path: temp_folder,
+                        },
+                    )
+                    .await
+                    .unwrap();
+			}
+			RequestData::Seal { tensors } => {
+				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 { handle, .. } => {
+				todo!()
+			}
+		}
+	}
+}

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

@@ -0,0 +1,90 @@
+use color_eyre::{Report, Result};
+use color_eyre::eyre::eyre;
+use carton_runner_interface::types::{
+    Tensor as CartonTensor,
+	TensorStorage as CartonStorage,
+};
+
+use crate::component::{
+	Tensor as WasmTensor,
+	TensorNumeric,
+	TensorString,
+	Dtype
+};
+
+impl Into<CartonTensor> for WasmTensor {
+	fn into(self) -> CartonTensor {
+		todo!()
+	}
+}
+
+impl TryFrom<CartonTensor> for WasmTensor {
+	type Error = Report;
+
+	fn try_from(value: CartonTensor) -> Result<Self> {
+		Ok(match value {
+			CartonTensor::Float(t) => WasmTensor::Numeric(t.into()),
+			CartonTensor::Double(t) => WasmTensor::Numeric(t.into()),
+			CartonTensor::I8(t) => WasmTensor::Numeric(t.into()),
+			CartonTensor::I16(t) => WasmTensor::Numeric(t.into()),
+			CartonTensor::I32(t) => WasmTensor::Numeric(t.into()),
+			CartonTensor::I64(t) => WasmTensor::Numeric(t.into()),
+			CartonTensor::U8(t) => WasmTensor::Numeric(t.into()),
+			CartonTensor::U16(t) => WasmTensor::Numeric(t.into()),
+			CartonTensor::U32(t) => WasmTensor::Numeric(t.into()),
+			CartonTensor::U64(t) => WasmTensor::Numeric(t.into()),
+			CartonTensor::String(t) => WasmTensor::String(t.into()),
+			CartonTensor::NestedTensor(_) => return Err(eyre!("Nested tensors are not supported")),
+		})
+	}
+}
+
+impl<T: DTypeOf> From<CartonStorage<T>> for TensorNumeric {
+	fn from(value: CartonStorage<T>) -> Self {
+		let (buffer, shape) = unsafe { value.into_bytes() };
+		TensorNumeric {
+			buffer,
+			dtype: T::dtype(),
+			shape,
+		}
+	}
+}
+
+impl From<CartonStorage<String>> for TensorString {
+	fn from(value: CartonStorage<String>) -> Self {
+		let (buffer, shape) = value.into_vec();
+		TensorString {
+			buffer,
+			shape,
+		}
+	}
+}
+
+trait DTypeOf {
+	fn dtype() -> Dtype;
+}
+
+macro_rules! for_each_numeric {
+    ($macro_name:ident) => {
+        $macro_name!(f32, Dtype::F32);
+        $macro_name!(f64, Dtype::F64);
+        $macro_name!(i32, Dtype::I32);
+        $macro_name!(i64, Dtype::I64);
+        $macro_name!(u8, Dtype::Ui8);
+        $macro_name!(u16, Dtype::Ui16);
+        $macro_name!(u32, Dtype::Ui32);
+        $macro_name!(u64, Dtype::Ui64);
+    };
+}
+
+macro_rules! implement_dtypeof {
+    ($type:ty, $enum_variant:expr) => {
+        impl DTypeOf for $type {
+            fn dtype() -> Dtype {
+                $enum_variant
+            }
+        }
+    };
+}
+
+for_each_numeric!(implement_dtypeof);

source/carton-runner-wasm/tests/basic_model/Cargo.toml

@@ -0,0 +1,13 @@
+[package]
+name = "basic_model"
+version = "0.1.0"
+edition = "2021"
+
+# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
+[lib]
+crate-type = ["cdylib"]
+
+[dependencies]
+wit-bindgen = { git = "https://github.com/bytecodealliance/wit-bindgen", version = "0.12.0" }
+
+[workspace]

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

@@ -0,0 +1,14 @@
+wit_bindgen::generate!({
+	world: "infer",
+	exports: {
+		world: Model
+	}
+});
+
+struct Model;
+
+impl Guest for Model {
+	fn infer() {
+		todo!()
+	}
+}