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

use super::{RemoteClient, runner::RemoteDevice};
use crate::shared::{ConnectionId, SessionId, Task, TaskResponse, TaskResponseContent};
use burn_communication::{CommunicationChannel, Message, ProtocolClient};
use std::{collections::HashMap, marker::PhantomData, sync::Arc};

pub type CallbackSender = async_channel::Sender<TaskResponseContent>;

#[derive(Debug)]
pub enum ClientRequest {
    WithSyncCallback(Task, CallbackSender),
    WithoutCallback(Task),
}

pub(crate) struct ClientWorker<C: ProtocolClient> {
    requests: HashMap<ConnectionId, CallbackSender>,
    _p: PhantomData<C>,
}

impl<C: ProtocolClient> ClientWorker<C> {
    async fn on_response(&mut self, response: TaskResponse) {
        match self.requests.remove(&response.id) {
            Some(request) => {
                request.send(response.content).await.unwrap();
            }
            None => {
                panic!("Can't ignore message from the server.");
            }
        }
    }

    fn register_callback(&mut self, id: ConnectionId, callback: CallbackSender) {
        self.requests.insert(id, callback);
    }
}

impl<C: ProtocolClient> ClientWorker<C> {
    pub fn start(device: RemoteDevice) -> RemoteClient {
        let runtime = Arc::new(
            tokio::runtime::Builder::new_multi_thread()
                .enable_io()
                .build()
                .unwrap(),
        );

        let (sender, rec) = async_channel::bounded(10);

        let session_id = SessionId::new();
        let address = device.address.clone();

        #[allow(deprecated)]
        runtime.spawn(async move {
            log::info!("Connecting to {} ...", address.clone());
            let mut stream_request = C::connect(address.clone(), "request")
                .await
                .expect("Server to be accessible");
            let mut stream_response = C::connect(address, "response")
                .await
                .expect("Server to be accessible");

            let state = Arc::new(tokio::sync::Mutex::new(ClientWorker::<C>::default()));

            // Init the connection.
            let bytes: bytes::Bytes = rmp_serde::to_vec(&Task::Init(session_id))
                .expect("Can serialize tasks to bytes.")
                .into();
            stream_request
                .send(Message::new(bytes.clone()))
                .await
                .expect("Can send the message over the comms channel.");
            stream_response
                .send(Message::new(bytes))
                .await
                .expect("Can send the message on the websocket.");

            // Async worker loading callbacks from the server.
            let state_ws = state.clone();
            tokio::spawn(async move {
                while let Ok(msg) = stream_response.recv().await {
                    let msg = match msg {
                        Some(msg) => msg,
                        None => {
                            log::warn!("Closed connection");
                            return;
                        }
                    };

                    let response: TaskResponse = rmp_serde::from_slice(&msg.data)
                        .expect("Can deserialize messages from the websocket.");
                    let mut state = state_ws.lock().await;
                    state.on_response(response).await;
                }
            });

            // Channel async worker sending operations to the server.
            tokio::spawn(async move {
                while let Ok(req) = rec.recv().await {
                    let task = match req {
                        ClientRequest::WithSyncCallback(task, callback) => {
                            if let Task::Compute(_content, id) = &task {
                                let mut state = state.lock().await;
                                state.register_callback(*id, callback);
                            }
                            task
                        }
                        ClientRequest::WithoutCallback(task) => task,
                    };
                    let bytes = rmp_serde::to_vec(&task)
                        .expect("Can serialize tasks to bytes.")
                        .into();
                    stream_request
                        .send(Message::new(bytes))
                        .await
                        .expect("Can send the message on the websocket.");
                }
            });
        });

        RemoteClient::new(device, sender, runtime, session_id)
    }
}

impl<C: ProtocolClient> Default for ClientWorker<C> {
    fn default() -> Self {
        Self {
            requests: Default::default(),
            _p: PhantomData,
        }
    }
}