Burn-remote to_device function

crates/burn-remote/Cargo.toml

@@ -44,6 +44,9 @@ axum = { version = "0.8.3", features = ["ws"], optional = true }
 tracing-core = { workspace = true, optional = true }
 tracing-subscriber = { workspace = true, optional = true }
 
+[dev-dependencies]
+burn-ndarray = { path = "../burn-ndarray", version = "0.18.0" }
+
 [package.metadata.docs.rs]
 features = ["doc"]
 rustdoc-args = ["--cfg", "docsrs"]

crates/burn-remote/src/client/channel.rs

@@ -1,5 +1,5 @@
 use burn_ir::TensorIr;
-use burn_router::{RouterTensor, RunnerChannel, TensorHandle};
+use burn_router::{RouterTensor, RunnerChannel, RunnerClient, TensorHandle};
 
 use super::{
     WsClient,
@@ -29,16 +29,18 @@ impl RunnerChannel for WsChannel {
         WsClient::init(device.clone())
     }
 
-    fn get_tensor_handle(_tensor: &TensorIr, _client: &Self::Client) -> TensorHandle<Self::Bridge> {
-        panic!("Unsupported")
+    fn get_tensor_handle(tensor: &TensorIr, client: &Self::Client) -> TensorHandle<Self::Bridge> {
+        client.runtime.block_on(client.read_tensor(tensor.clone()))
     }
 
     fn register_tensor(
-        _client: &Self::Client,
-        _handle: TensorHandle<Self::Bridge>,
+        client: &Self::Client,
+        handle: TensorHandle<Self::Bridge>,
         _shape: Vec<usize>,
         _dtype: burn_tensor::DType,
     ) -> RouterTensor<Self::Client> {
-        panic!("Unsupported")
+        let router_tensor = client.register_tensor_data(handle);
+
+        router_tensor
     }
 }

crates/burn-remote/src/client/runner.rs

@@ -4,7 +4,10 @@ use burn_tensor::{
     DType, TensorData,
     backend::{DeviceId, DeviceOps},
 };
-use std::sync::Arc;
+use std::{
+    hash::{DefaultHasher, Hash, Hasher},
+    sync::Arc,
+};
 
 use crate::shared::{ComputeTask, TaskResponseContent};
 
@@ -92,6 +95,8 @@ impl RunnerClient for WsClient {
 /// The device contains the connection information of the server.
 pub struct WsDevice {
     pub(crate) address: Arc<String>,
+    // Unique ID generated from hash of the address
+    pub(crate) id: u32,
 }
 
 impl WsDevice {
@@ -106,8 +111,13 @@ impl WsDevice {
             address += url;
         };
 
+        let mut hasher = DefaultHasher::new();
+        address.hash(&mut hasher);
+        let id = hasher.finish() as u32;
+
         Self {
             address: Arc::new(address),
+            id,
         }
     }
 }
@@ -119,17 +129,15 @@ impl Default for WsDevice {
             Err(_) => String::from("ws://127.0.0.1:3000"),
         };
 
-        Self {
-            address: Arc::new(address),
-        }
+        Self::new(&address)
     }
 }
 
 impl DeviceOps for WsDevice {
     fn id(&self) -> DeviceId {
         DeviceId {
             type_id: 0,
-            index_id: 0,
+            index_id: self.id,
         }
     }
 }

crates/burn-remote/src/lib.rs

@@ -35,3 +35,51 @@ mod __client {
 }
 #[cfg(feature = "client")]
 pub use __client::*;
+
+#[cfg(all(test, feature = "client", feature = "server"))]
+mod tests {
+    use crate::RemoteBackend;
+    use burn_ndarray::NdArray;
+    use burn_tensor::{Distribution, Tensor};
+
+    #[test]
+    pub fn test_to_device_over_websocket() {
+        let rt = tokio::runtime::Builder::new_multi_thread()
+            .enable_io()
+            .build()
+            .unwrap();
+
+        rt.spawn(crate::server::start_async::<NdArray>(
+            Default::default(),
+            3000,
+        ));
+        rt.spawn(crate::server::start_async::<NdArray>(
+            Default::default(),
+            3010,
+        ));
+
+        let remote_device_1 = super::RemoteDevice::new("ws://localhost:3000");
+        let remote_device_2 = super::RemoteDevice::new("ws://localhost:3010");
+
+        // Some random input
+        let input_shape = [1, 28, 28];
+        let input = Tensor::<RemoteBackend, 3>::random(
+            input_shape,
+            Distribution::Default,
+            &remote_device_1,
+        );
+        let numbers_expected: Vec<f32> = input.to_data().to_vec().unwrap();
+
+        // Move tensor to device 2
+        let input = input.to_device(&remote_device_2);
+        let numbers: Vec<f32> = input.to_data().to_vec().unwrap();
+        assert_eq!(numbers, numbers_expected);
+
+        // Move tensor back to device 1
+        let input = input.to_device(&remote_device_1);
+        let numbers: Vec<f32> = input.to_data().to_vec().unwrap();
+        assert_eq!(numbers, numbers_expected);
+
+        rt.shutdown_background();
+    }
+}

crates/burn-remote/src/server/base.rs

@@ -38,7 +38,9 @@ impl<B: BackendIr> WsServer<B> {
                 }
                 true
             }));
-        registry().with(layer).init();
+
+        // If we start multiple servers in the same process, this will fail, it's ok
+        let _ = registry().with(layer).try_init();
 
         let address = format!("0.0.0.0:{port}");
         log::info!("Start server {address} on device {device:?}");
@@ -173,13 +175,17 @@ impl<B: BackendIr> WsServer<B> {
             };
         }
 
-        log::info!("Closing connection");
+        log::info!("Closing session {:?}", session_id);
         self.state.close(session_id);
     }
 }
 
+pub async fn start_async<B: BackendIr>(device: Device<B>, port: u16) {
+    WsServer::<B>::start(device, port).await;
+}
+
 #[tokio::main]
 /// Start the server on the given port and [device](Device).
 pub async fn start<B: BackendIr>(device: Device<B>, port: u16) {
-    WsServer::<B>::start(device, port).await;
+    start_async::<B>(device, port).await
 }

crates/burn-remote/src/server/mod.rs

@@ -4,4 +4,4 @@ pub(crate) mod stream;
 
 mod base;
 
-pub use base::start;
+pub use base::{start, start_async};

crates/burn-remote/src/shared/mod.rs

@@ -1,2 +1,4 @@
 mod task;
+
+#[allow(unused_imports)]
 pub(crate) use task::*;

crates/burn/Cargo.toml

@@ -23,6 +23,7 @@ default = [
     "burn-tch?/default",
     "burn-wgpu?/default",
     "burn-router?/default",
+    "burn-remote?/default",
     "burn-cuda?/default",
     "burn-autodiff?/default",
     "burn-rocm?/default",
@@ -131,3 +132,4 @@ burn-remote = { path = "../burn-remote", version = "0.18.0", default-features =
 burn-router = { path = "../burn-router", version = "0.18.0", default-features = false, optional = true }
 burn-tch = { path = "../burn-tch", version = "0.18.0", optional = true }
 burn-wgpu = { path = "../burn-wgpu", version = "0.18.0", optional = true, default-features = false }
+burn-ir = { path = "../burn-ir", version = "0.18.0", default-features = false }

crates/burn/src/backend.rs

@@ -59,3 +59,5 @@ pub use burn_router::Router;
 
 #[cfg(feature = "router")]
 pub use burn_router as router;
+
+pub use burn_ir::BackendIr;