burn-remote 0.21.0

Backend router decorator over the network.
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134

use core::marker::PhantomData;
use std::sync::Arc;

use crate::shared::{ConnectionId, TaskResponse, TensorRemote};

use super::processor::{Processor, ProcessorTask};
use burn_backend::TensorData;
use burn_communication::{
    Protocol,
    data_service::{TensorDataService, TensorTransferId},
};
use burn_ir::{BackendIr, OperationIr, TensorId, TensorIr};
use burn_router::Runner;
use burn_std::DType;
use tokio::sync::mpsc::{Receiver, Sender};

/// A stream makes sure all operations registered are executed in the order they were sent to the
/// server, potentially waiting to reconstruct consistency.
#[derive(Clone)]
pub struct Stream<B, P>
where
    B: BackendIr,
    P: Protocol,
{
    compute_sender: Sender<ProcessorTask>,
    writer_sender: Sender<Receiver<TaskResponse>>,
    _p: PhantomData<B>,
    _n: PhantomData<P>,
}

impl<B, P> Stream<B, P>
where
    B: BackendIr,
    P: Protocol,
{
    pub async fn new(
        runner: Runner<B>,
        writer_sender: Sender<Receiver<TaskResponse>>,
        data_service: Arc<TensorDataService<B, P>>,
    ) -> Self {
        let sender = Processor::<B, P>::start(runner, data_service).await;

        Self {
            compute_sender: sender,
            writer_sender,
            _p: PhantomData,
            _n: PhantomData,
        }
    }

    pub async fn register_operation(&self, op: Box<OperationIr>) {
        self.compute_sender
            .send(ProcessorTask::RegisterOperation(op))
            .await
            .unwrap();
    }

    pub async fn register_tensor(&self, tensor_id: TensorId, data: TensorData) {
        self.compute_sender
            .send(ProcessorTask::RegisterTensor(tensor_id, data))
            .await
            .unwrap();
    }

    pub async fn register_tensor_remote(&self, tensor: TensorRemote, new_id: TensorId) {
        self.compute_sender
            .send(ProcessorTask::RegisterTensorRemote(tensor, new_id))
            .await
            .unwrap();
    }

    pub async fn expose_tensor_remote(
        &self,
        tensor: TensorIr,
        count: u32,
        transfer_id: TensorTransferId,
    ) {
        self.compute_sender
            .send(ProcessorTask::ExposeTensorRemote {
                tensor,
                count,
                transfer_id,
            })
            .await
            .unwrap();
    }

    pub async fn read_tensor(&self, id: ConnectionId, desc: TensorIr) {
        let (callback_sender, callback_rec) = tokio::sync::mpsc::channel(1);

        self.compute_sender
            .send(ProcessorTask::ReadTensor(id, desc, callback_sender))
            .await
            .unwrap();

        self.writer_sender.send(callback_rec).await.unwrap();
    }

    pub async fn sync(&self, id: ConnectionId) {
        let (callback_sender, callback_rec) = tokio::sync::mpsc::channel(1);

        self.compute_sender
            .send(ProcessorTask::Sync(id, callback_sender))
            .await
            .unwrap();

        self.writer_sender.send(callback_rec).await.unwrap();
    }

    pub async fn close(&self) {
        self.compute_sender
            .send(ProcessorTask::Close)
            .await
            .unwrap();
    }

    pub async fn seed(&self, seed: u64) {
        self.compute_sender
            .send(ProcessorTask::Seed(seed))
            .await
            .unwrap();
    }

    pub async fn dtype_usage(&self, id: ConnectionId, dtype: DType) {
        let (callback_sender, callback_rec) = tokio::sync::mpsc::channel(1);

        self.compute_sender
            .send(ProcessorTask::DTypeUsage(id, dtype, callback_sender))
            .await
            .unwrap();

        self.writer_sender.send(callback_rec).await.unwrap();
    }
}