burn_remote/client/

runner.rs

use burn_common::future::DynFut;
use burn_ir::TensorIr;
use burn_router::{MultiBackendBridge, RouterTensor, RunnerClient, get_client};
use burn_tensor::{
    DType, TensorData,
    backend::{DeviceId, DeviceOps},
};
use std::{
    hash::{DefaultHasher, Hash, Hasher},
    sync::Arc,
};

use crate::shared::{ComputeTask, TaskResponseContent, TensorRemote};

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.
//
// The overhead is minimal, since we only wait for the task to be sent to the async
// channel, but not sent to the websocket server and even less processed by the server.
impl RunnerClient for WsClient {
    type Device = WsDevice;

    fn register(&self, op: burn_ir::OperationIr) {
        self.sender
            .send(ComputeTask::RegisterOperation(Box::new(op)));
    }

    fn read_tensor(&self, tensor: burn_ir::TensorIr) -> DynFut<TensorData> {
        // Important for ordering to call the creation of the future sync.
        let fut = self.sender.send_callback(ComputeTask::ReadTensor(tensor));

        Box::pin(async move {
            match fut.await {
                TaskResponseContent::ReadTensor(data) => data,
                _ => panic!("Invalid message type"),
            }
        })
    }

    fn register_tensor_data(&self, data: TensorData) -> RouterTensor<Self> {
        let id = self.sender.new_tensor_id();
        let shape = data.shape.clone();
        let dtype = data.dtype;

        self.sender.send(ComputeTask::RegisterTensor(id, data));

        RouterTensor::new(id, shape, dtype, self.clone())
    }

    fn register_empty_tensor(
        &self,
        shape: Vec<usize>,
        dtype: burn_tensor::DType,
    ) -> RouterTensor<Self> {
        let id = self.sender.new_tensor_id();

        RouterTensor::new(id, shape, dtype, self.clone())
    }

    fn register_float_tensor(
        &self,
        shape: Vec<usize>,
        _dtype: burn_tensor::FloatDType,
    ) -> RouterTensor<Self> {
        self.register_empty_tensor(shape, DType::F32)
    }

    fn device(&self) -> Self::Device {
        self.device.clone()
    }

    fn sync(&self) {
        // Important for ordering to call the creation of the future sync.
        let fut = self.sender.send_callback(ComputeTask::SyncBackend);

        let runtime = self.runtime.clone();

        match runtime.block_on(fut) {
            TaskResponseContent::SyncBackend => {}
            _ => panic!("Invalid message type"),
        };
    }

    fn seed(&self, _seed: u64) {
        // TODO
    }
}

#[derive(Clone, PartialEq, Eq, Debug)]
/// 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 {
    /// Create a device from an url.
    pub fn new(url: &str) -> Self {
        let mut address = String::new();

        if !url.starts_with("ws://") {
            address += "ws://";
            address += url;
        } else {
            address += url;
        };

        let mut hasher = DefaultHasher::new();
        address.hash(&mut hasher);
        let id = hasher.finish() as u32;

        Self {
            address: Arc::new(address),
            id,
        }
    }
}

impl Default for WsDevice {
    fn default() -> Self {
        let address = match std::env::var("BURN_REMOTE_ADDRESS") {
            Ok(address) => address,
            Err(_) => String::from("ws://127.0.0.1:3000"),
        };

        Self::new(&address)
    }
}

impl DeviceOps for WsDevice {
    fn id(&self) -> DeviceId {
        DeviceId {
            type_id: 0,
            index_id: self.id,
        }
    }
}

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 = RemoteTensorHandle;
    type Device = WsDevice;

    fn change_backend_float(
        tensor: Self::TensorHandle,
        _shape: burn_tensor::Shape,
        target_device: &Self::Device,
    ) -> Self::TensorHandle {
        tensor.change_backend(target_device)
    }

    fn change_backend_int(
        tensor: Self::TensorHandle,
        _shape: burn_tensor::Shape,
        target_device: &Self::Device,
    ) -> Self::TensorHandle {
        tensor.change_backend(target_device)
    }

    fn change_backend_bool(
        tensor: Self::TensorHandle,
        _shape: burn_tensor::Shape,
        target_device: &Self::Device,
    ) -> Self::TensorHandle {
        tensor.change_backend(target_device)
    }
}