Add an Allocatable trait to the runner interface (#179)

This PR adds an `Allocatable` trait to the runner interface. This
generally makes it cleaner to write generic functions over Carton types
without needing to reference implementation details in trait bounds
(e.g. `where InlineAllocator: TypedAlloc<T, Output =
InlineTensorStorage>`).

Note: this PR does touch files inside `do_not_modify`, but it does so in
a way that does not affect the wire protocol.

Author: VivekPanyam

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

@@ -31,8 +31,17 @@ pub trait AsPtr<T> {
     fn as_mut_ptr(&mut self) -> *mut T;
 }
 
-pub trait TypedAlloc<T> {
-    type Output: AsPtr<T>;
+pub trait Allocator {
+    type Output;
+}
 
+pub trait TypedAlloc<T>: Allocator
+where
+    Self::Output: AsPtr<T>,
+{
     fn alloc(&self, numel: usize) -> Self::Output;
 }
+
+pub trait AllocatableBy<A: Allocator>: Sized {
+    fn alloc(allocator: &A, numel: usize) -> A::Output;
+}

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

@@ -23,7 +23,7 @@ use once_cell::sync::Lazy;
 use serde::{Deserialize, Serialize};
 
 use super::{
-    alloc::{AsPtr, NumericTensorType, TypedAlloc},
+    alloc::{AllocatableBy, Allocator, AsPtr, NumericTensorType, TypedAlloc},
     alloc_pool::{PoolAllocator, PoolItem},
     storage::TensorStorage,
 };
@@ -79,12 +79,14 @@ impl<T> AsPtr<T> for InlineTensorStorage {
     }
 }
 
+impl Allocator for InlineAllocator {
+    type Output = InlineTensorStorage;
+}
+
 for_each_numeric_carton_type! {
     $(
         /// We're using a macro here instead of a generic impl because rust gives misleading error messages otherwise.
         impl TypedAlloc<$RustType> for InlineAllocator {
-            type Output = InlineTensorStorage;
-
             fn alloc(&self, numel: usize) -> Self::Output {
                 // We need to convert to size_bytes since we always use a Vec<u8>
                 let size_bytes = numel * std::mem::size_of::<$RustType>();
@@ -101,8 +103,6 @@ for_each_numeric_carton_type! {
 }
 
 impl TypedAlloc<String> for InlineAllocator {
-    type Output = InlineTensorStorage;
-
     fn alloc(&self, numel: usize) -> Self::Output {
         let out = if !self.use_pool {
             vec![String::default(); numel].into()
@@ -115,10 +115,8 @@ impl TypedAlloc<String> for InlineAllocator {
 }
 
 // Copy the data
-impl<T: NumericTensorType + Default + Copy> From<ndarray::ArrayViewD<'_, T>>
-    for TensorStorage<T, InlineTensorStorage>
-where
-    InlineAllocator: TypedAlloc<T, Output = InlineTensorStorage>,
+impl<T: NumericTensorType + Default + Copy + AllocatableBy<InlineAllocator>>
+    From<ndarray::ArrayViewD<'_, T>> for TensorStorage<T, InlineTensorStorage>
 {
     fn from(view: ndarray::ArrayViewD<'_, T>) -> Self {
         // Alloc a tensor
@@ -152,17 +150,14 @@ impl From<ndarray::ArrayViewD<'_, String>> for TensorStorage<String, InlineTenso
 
 // Allocates a contiguous tensor with a shape and type
 #[cfg(feature = "benchmark")]
-pub fn alloc_tensor_no_pool<T: Default + Clone>(
+pub fn alloc_tensor_no_pool<T: Default + Clone + AllocatableBy<InlineAllocator>>(
     shape: Vec<u64>,
-) -> TensorStorage<T, InlineTensorStorage>
-where
-    InlineAllocator: TypedAlloc<T, Output = InlineTensorStorage>,
-{
+) -> TensorStorage<T, InlineTensorStorage> {
     static POOL_ALLOCATOR: Lazy<InlineAllocator> = Lazy::new(|| InlineAllocator::without_pool());
 
     let numel = shape.iter().product::<u64>().max(1) as usize;
 
-    let data = <InlineAllocator as TypedAlloc<T>>::alloc(&POOL_ALLOCATOR, numel);
+    let data = T::alloc(&POOL_ALLOCATOR, numel);
 
     TensorStorage {
         data,
@@ -172,15 +167,14 @@ where
     }
 }
 
-pub fn alloc_tensor<T: Default + Clone>(shape: Vec<u64>) -> TensorStorage<T, InlineTensorStorage>
-where
-    InlineAllocator: TypedAlloc<T, Output = InlineTensorStorage>,
-{
+pub fn alloc_tensor<T: Default + Clone + AllocatableBy<InlineAllocator>>(
+    shape: Vec<u64>,
+) -> TensorStorage<T, InlineTensorStorage> {
     static POOL_ALLOCATOR: Lazy<InlineAllocator> = Lazy::new(|| InlineAllocator::new());
 
     let numel = shape.iter().product::<u64>().max(1) as usize;
 
-    let data = <InlineAllocator as TypedAlloc<T>>::alloc(&POOL_ALLOCATOR, numel);
+    let data = T::alloc(&POOL_ALLOCATOR, numel);
 
     TensorStorage {
         data,
@@ -190,11 +184,17 @@ where
     }
 }
 
-impl<T: Default + Clone> TensorStorage<T, InlineTensorStorage>
+impl<T: Default + Clone + AllocatableBy<InlineAllocator>> TensorStorage<T, InlineTensorStorage> {
+    pub fn new(shape: Vec<u64>) -> TensorStorage<T, InlineTensorStorage> {
+        alloc_tensor(shape)
+    }
+}
+
+impl<T> AllocatableBy<InlineAllocator> for T
 where
     InlineAllocator: TypedAlloc<T, Output = InlineTensorStorage>,
 {
-    pub fn new(shape: Vec<u64>) -> TensorStorage<T, InlineTensorStorage> {
-        alloc_tensor(shape)
+    fn alloc(allocator: &InlineAllocator, numel: usize) -> InlineTensorStorage {
+        <InlineAllocator as TypedAlloc<T>>::alloc(allocator, numel)
     }
 }

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

@@ -25,7 +25,7 @@ use dashmap::DashMap;
 use once_cell::sync::Lazy;
 
 use super::{
-    alloc::{AsPtr, NumericTensorType, TypedAlloc},
+    alloc::{AsPtr, NumericTensorType, TypedAlloc, Allocator, AllocatableBy},
     alloc_pool::{AllocItem, PoolAllocator, PoolItem},
     storage::TensorStorage,
 };
@@ -260,12 +260,14 @@ impl SHMAllocator {
     }
 }
 
+impl Allocator for SHMAllocator {
+    type Output = SHMTensorStorage;
+}
+
 for_each_numeric_carton_type! {
     $(
         /// We're using a macro here instead of a generic impl because rust gives misleading error messages otherwise.
         impl TypedAlloc<$RustType> for SHMAllocator {
-            type Output = SHMTensorStorage;
-
             fn alloc(&self, numel: usize) -> Self::Output {
                 // We need to convert to size_bytes
                 let size_bytes = numel * std::mem::size_of::<$RustType>();
@@ -282,8 +284,6 @@ for_each_numeric_carton_type! {
 }
 
 impl TypedAlloc<String> for SHMAllocator {
-    type Output = SHMTensorStorage;
-
     fn alloc(&self, numel: usize) -> Self::Output {
         let out = if !self.use_pool {
             vec![String::default(); numel].into()
@@ -384,17 +384,15 @@ impl From<ndarray::ArrayViewD<'_, String>> for TensorStorage<String, SHMTensorSt
 
 // Allocates a contiguous tensor with a shape and type
 #[cfg(feature = "benchmark")]
-pub fn alloc_tensor_no_pool<T: Default + Clone>(
+pub fn alloc_tensor_no_pool<T: Default + Clone + AllocatableBy<SHMAllocator>>(
     shape: Vec<u64>,
 ) -> TensorStorage<T, SHMTensorStorage>
-where
-    SHMAllocator: TypedAlloc<T, Output = SHMTensorStorage>,
 {
     static POOL_ALLOCATOR: Lazy<SHMAllocator> = Lazy::new(|| SHMAllocator::without_pool());
 
     let numel = shape.iter().product::<u64>().max(1) as usize;
 
-    let data = <SHMAllocator as TypedAlloc<T>>::alloc(&POOL_ALLOCATOR, numel);
+    let data = T::alloc(&POOL_ALLOCATOR, numel);
 
     TensorStorage {
         data,
@@ -404,15 +402,13 @@ where
     }
 }
 
-pub fn alloc_tensor<T: Default + Clone>(shape: Vec<u64>) -> TensorStorage<T, SHMTensorStorage>
-where
-    SHMAllocator: TypedAlloc<T, Output = SHMTensorStorage>,
+pub fn alloc_tensor<T: Default + Clone + AllocatableBy<SHMAllocator>>(shape: Vec<u64>) -> TensorStorage<T, SHMTensorStorage>
 {
     static POOL_ALLOCATOR: Lazy<SHMAllocator> = Lazy::new(|| SHMAllocator::new());
 
     let numel = shape.iter().product::<u64>().max(1) as usize;
 
-    let data = <SHMAllocator as TypedAlloc<T>>::alloc(&POOL_ALLOCATOR, numel);
+    let data = T::alloc(&POOL_ALLOCATOR, numel);
 
     TensorStorage {
         data,
@@ -421,3 +417,13 @@ where
         pd: PhantomData,
     }
 }
+
+
+impl<T> AllocatableBy<SHMAllocator> for T
+where
+    SHMAllocator: TypedAlloc<T, Output = SHMTensorStorage>,
+{
+    fn alloc(allocator: &SHMAllocator, numel: usize) -> SHMTensorStorage {
+        <SHMAllocator as TypedAlloc<T>>::alloc(allocator, numel)
+    }
+}

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

@@ -12,11 +12,15 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-pub use carton_macros::for_each_carton_type;
+pub use carton_macros::{for_each_carton_type, for_each_numeric_carton_type};
 use serde::{Deserialize, Serialize};
 use std::collections::HashMap;
 
-use super::{alloc_inline::InlineTensorStorage, comms::Comms};
+use super::{
+    alloc::AllocatableBy,
+    alloc_inline::{InlineAllocator, InlineTensorStorage},
+    comms::Comms,
+};
 
 #[derive(Debug, Serialize, Deserialize)]
 pub(crate) struct RPCRequest {
@@ -247,6 +251,9 @@ for_each_carton_type! {
 
 pub type TensorStorage<T> = super::storage::TensorStorage<T, InlineTensorStorage>;
 
+pub trait Allocatable: AllocatableBy<InlineAllocator> {}
+impl<T> Allocatable for T where T: AllocatableBy<InlineAllocator> {}
+
 for_each_carton_type! {
     $(
         impl From<TensorStorage<$RustType>> for Tensor {

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

@@ -17,12 +17,8 @@ use std::{collections::HashMap, sync::Arc};
 use crate::{
     client::Client,
     do_not_modify::comms::OwnedComms,
-    do_not_modify::{
-        alloc::TypedAlloc,
-        alloc_inline::{InlineAllocator, InlineTensorStorage},
-        types::{Device, RPCRequestData, RPCResponseData, SealHandle, Tensor},
-    },
-    types::{Handle, RunnerOpt, TensorStorage},
+    do_not_modify::types::{Device, RPCRequestData, RPCResponseData, SealHandle, Tensor},
+    types::{Allocatable, Handle, RunnerOpt, TensorStorage},
 };
 
 use futures::Stream;
@@ -299,9 +295,11 @@ impl Runner {
         }
     }
 
-    pub fn alloc_tensor<T: Clone + Default>(&self, shape: Vec<u64>) -> Result<Tensor, String>
+    pub fn alloc_tensor<T: Clone + Default + Allocatable>(
+        &self,
+        shape: Vec<u64>,
+    ) -> Result<Tensor, String>
     where
-        InlineAllocator: TypedAlloc<T, Output = InlineTensorStorage>,
         Tensor: From<TensorStorage<T>>,
     {
         Ok(TensorStorage::new(shape).into())

source/carton-runner-wasm/Cargo.toml

@@ -7,7 +7,6 @@ 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 }
@@ -28,4 +27,5 @@ semver = "1.0.20"
 [dev-dependencies]
 escargot = "0.5.8"
 paste = "1.0.14"
-tempfile = "3.8.0"
+tempfile = "3.8.0"
+carton = { path = "../carton" }

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

@@ -1,8 +1,9 @@
 use color_eyre::eyre::{ensure, eyre};
 use color_eyre::{Report, Result};
 
-use carton::types::for_each_numeric_carton_type;
-use carton_runner_interface::types::{Tensor as CartonTensor, TensorStorage as CartonStorage};
+use carton_runner_interface::types::{
+    for_each_numeric_carton_type, Tensor as CartonTensor, TensorStorage as CartonStorage,
+};
 
 use crate::component::{Dtype, Tensor as WasmTensor, TensorNumeric, TensorString};