Lazy tensor downloading in burn-remote (#3276)

Author: Cielbird

Cargo.lock

@@ -15,10 +15,10 @@ version.workspace = true
 workspace = true
 
 [features]
-default = []
+default = ["client", "server"]
 doc = []
 client = ["tokio-tungstenite", "async-channel", "tokio/sync"]
-server = ["axum", "tracing-core", "tracing-subscriber"]
+server = ["tokio-tungstenite", "async-channel", "tokio/sync", "axum", "tracing-core", "tracing-subscriber"]
 
 
 [dependencies]
@@ -27,6 +27,8 @@ burn-tensor = { path = "../burn-tensor", version = "0.18.0", default-features =
 burn-common = { path = "../burn-common", version = "0.18.0", default-features = true }
 burn-router = { path = "../burn-router", version = "0.18.0", default-features = true }
 
+bytes = { version = "1.0" }
+
 # Basic dependencies
 derive-new = { workspace = true }
 log = { workspace = true }

crates/burn-remote/Cargo.toml

@@ -1,9 +1,11 @@
 use burn_ir::TensorIr;
-use burn_router::{RouterTensor, RunnerChannel, RunnerClient, TensorHandle};
+use burn_router::{RouterTensor, RunnerChannel, get_client};
+
+use crate::shared::{ComputeTask, TensorRemote};
 
 use super::{
     WsClient,
-    runner::{WsBridge, WsDevice},
+    runner::{RemoteTensorHandle, WsBridge, WsDevice},
 };
 
 /// A local channel with direct connection to the backend runner clients.
@@ -29,17 +31,47 @@ impl RunnerChannel for WsChannel {
         WsClient::init(device.clone())
     }
 
-    fn get_tensor_handle(tensor: &TensorIr, client: &Self::Client) -> TensorHandle<Self::Bridge> {
-        client.runtime.block_on(client.read_tensor(tensor.clone()))
+    fn get_tensor_handle(tensor: &TensorIr, client: &Self::Client) -> RemoteTensorHandle {
+        RemoteTensorHandle {
+            client: client.clone(),
+            tensor: tensor.clone(),
+        }
     }
 
     fn register_tensor(
         client: &Self::Client,
-        handle: TensorHandle<Self::Bridge>,
-        _shape: Vec<usize>,
-        _dtype: burn_tensor::DType,
+        handle: RemoteTensorHandle,
+        shape: Vec<usize>,
+        dtype: burn_tensor::DType,
+    ) -> RouterTensor<Self::Client> {
+        let remote_tensor = TensorRemote {
+            id: handle.tensor.id,
+            address: client.device.address.to_string(),
+        };
+        let new_id = client.sender.new_tensor_id();
+        client
+            .sender
+            .send(ComputeTask::RegisterTensorRemote(remote_tensor, new_id));
+
+        RouterTensor::new(handle.tensor.id, shape, dtype, client.clone())
+    }
+
+    fn change_client_backend(
+        tensor: RouterTensor<Self::Client>,
+        target_device: &Self::Device, // target device
     ) -> RouterTensor<Self::Client> {
-        let router_tensor = client.register_tensor_data(handle);
+        // Get tensor handle from current client
+        let original_client = tensor.client.clone();
+        let desc = tensor.into_ir();
+        let handle = Self::get_tensor_handle(&desc, &original_client);
+
+        let handle = handle.change_backend(target_device);
+
+        let id = handle.tensor.id;
+
+        let target_client = get_client::<Self>(target_device);
+        let router_tensor: RouterTensor<WsClient> =
+            RouterTensor::new(id, handle.tensor.shape, handle.tensor.dtype, target_client);
 
         router_tensor
     }

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

@@ -1,5 +1,6 @@
 use burn_common::future::DynFut;
-use burn_router::{MultiBackendBridge, RouterTensor, RunnerClient};
+use burn_ir::TensorIr;
+use burn_router::{MultiBackendBridge, RouterTensor, RunnerClient, get_client};
 use burn_tensor::{
     DType, TensorData,
     backend::{DeviceId, DeviceOps},
@@ -9,9 +10,9 @@ use std::{
     sync::Arc,
 };
 
-use crate::shared::{ComputeTask, TaskResponseContent};
+use crate::shared::{ComputeTask, TaskResponseContent, TensorRemote};
 
-use super::WsClient;
+use super::{WsChannel, WsClient};
 
 // It is very important to block on any request made with the sender, since ordering is crucial
 // when registering operation or creating tensors.
@@ -45,7 +46,7 @@ impl RunnerClient for WsClient {
 
         self.sender.send(ComputeTask::RegisterTensor(id, data));
 
-        RouterTensor::new(Arc::new(id), shape, dtype, self.clone())
+        RouterTensor::new(id, shape, dtype, self.clone())
     }
 
     fn register_empty_tensor(
@@ -55,7 +56,7 @@ impl RunnerClient for WsClient {
     ) -> RouterTensor<Self> {
         let id = self.sender.new_tensor_id();
 
-        RouterTensor::new(Arc::new(id), shape, dtype, self.clone())
+        RouterTensor::new(id, shape, dtype, self.clone())
     }
 
     fn register_float_tensor(
@@ -140,31 +141,67 @@ impl DeviceOps for WsDevice {
 
 pub struct WsBridge;
 
+pub struct RemoteTensorHandle {
+    pub(crate) client: WsClient,
+    pub(crate) tensor: TensorIr,
+}
+
+impl RemoteTensorHandle {
+    /// Changes the backend of the tensor via a WebSocket.
+    /// We ask the original server to expose the tensor, then ask the target server to fetch
+    /// the tensor. The target server will open a new websocket connection to the original server
+    /// to download the data.
+    /// This way the client never sees the tensor's data, and we avoid a bottleneck.
+    pub(crate) fn change_backend(mut self, target_device: &WsDevice) -> Self {
+        self.client.sender.send(ComputeTask::ExposeTensorRemote {
+            tensor: self.tensor.clone(),
+            count: 1,
+        });
+
+        let target_client: WsClient = get_client::<WsChannel>(target_device);
+
+        let new_id = target_client.sender.new_tensor_id();
+
+        let remote_tensor = TensorRemote {
+            id: self.tensor.id,
+            address: self.client.device.address.to_string(),
+        };
+        target_client
+            .sender
+            .send(ComputeTask::RegisterTensorRemote(remote_tensor, new_id));
+
+        self.tensor.id = new_id;
+        self.client = target_client;
+
+        self
+    }
+}
+
 impl MultiBackendBridge for WsBridge {
-    type TensorHandle = TensorData;
+    type TensorHandle = RemoteTensorHandle;
     type Device = WsDevice;
 
     fn change_backend_float(
         tensor: Self::TensorHandle,
         _shape: burn_tensor::Shape,
-        _target_device: &Self::Device,
+        target_device: &Self::Device,
     ) -> Self::TensorHandle {
-        tensor
+        tensor.change_backend(target_device)
     }
 
     fn change_backend_int(
         tensor: Self::TensorHandle,
         _shape: burn_tensor::Shape,
-        _target_device: &Self::Device,
+        target_device: &Self::Device,
     ) -> Self::TensorHandle {
-        tensor
+        tensor.change_backend(target_device)
     }
 
     fn change_backend_bool(
         tensor: Self::TensorHandle,
         _shape: burn_tensor::Shape,
-        _target_device: &Self::Device,
+        target_device: &Self::Device,
     ) -> Self::TensorHandle {
-        tensor
+        tensor.change_backend(target_device)
     }
 }