ruCCL 0.21.1

Ruda collective communication algorithms and orchestration.
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
135
136
137
138

use std::fmt::Debug;
use std::sync::Arc;
use tokio::sync::Mutex;

use crate::global::{
    orchestrator::state::GlobalCollectiveState,
    shared::{CollectiveMessage, GlobalCollectiveError},
};
use ruda_communication::{
    CommunicationChannel, Message, ProtocolServer, util::os_shutdown_signal, websocket::WsServer,
};

/// The global collective state manages collective operations on the global level
#[derive(Clone)]
pub(crate) struct GlobalOrchestrator {
    state: Arc<Mutex<GlobalCollectiveState>>,
}

impl GlobalOrchestrator {
    /// Starts the comms server with two routes: "/request" and "/response"
    pub(crate) async fn start<F, S: ProtocolServer + Debug>(
        shutdown_signal: F,
        comms_server: S,
    ) -> Result<(), GlobalCollectiveError>
    where
        F: Future<Output = ()> + Send + 'static,
    {
        let state = GlobalCollectiveState::new();
        let server = Self {
            state: Arc::new(tokio::sync::Mutex::new(state)),
        };

        comms_server
            .route("/response", {
                let server = server.clone();
                async move |socket| {
                    if let Err(err) = server.handle_socket_response::<S>(socket).await {
                        log::error!("[Response Handler] Error: {err:?}")
                    }
                }
            })
            .route("/request", {
                let server = server.clone();
                async move |socket| {
                    if let Err(err) = server.handle_socket_request::<S>(socket).await {
                        log::error!("[Request Handler] Error: {err:?}")
                    }
                }
            })
            .serve(shutdown_signal)
            .await
            .map_err(|err| GlobalCollectiveError::Server(format!("{err:?}")))?;

        Ok(())
    }

    async fn handle_socket_response<S: ProtocolServer>(
        self,
        mut stream: S::Channel,
    ) -> Result<(), GlobalCollectiveError> {
        log::info!("[Response Handler] On new connection.");

        let msg = stream
            .recv()
            .await
            .map_err(|err| GlobalCollectiveError::Server(format!("{err:?}")))?;
        let Some(msg) = msg else {
            log::warn!("Response socket closed early!");
            return Ok(());
        };

        let msg = rmp_serde::from_slice::<CollectiveMessage>(&msg.data)
            .map_err(|_| GlobalCollectiveError::InvalidMessage)?;

        let CollectiveMessage::Init(id) = msg else {
            return Err(GlobalCollectiveError::FirstMsgNotInit);
        };

        let mut receiver = {
            let mut state = self.state.lock().await;
            state.get_session_responder(id)
        };

        while let Some(response) = receiver.recv().await {
            let bytes = rmp_serde::to_vec(&response).unwrap();

            stream.send(Message::new(bytes.into())).await?;
        }

        log::info!("[Response Handler] Closing connection.");
        Ok(())
    }

    async fn handle_socket_request<S: ProtocolServer>(
        self,
        mut stream: S::Channel,
    ) -> Result<(), GlobalCollectiveError> {
        log::info!("[Request Handler] On new connection.");

        let mut session_id = None;

        loop {
            let packet = stream.recv().await?;
            let Some(msg) = packet else {
                log::info!("Peer closed the connection");
                break;
            };

            let mut state = self.state.lock().await;

            let msg = rmp_serde::from_slice::<CollectiveMessage>(&msg.data)
                .map_err(|_| GlobalCollectiveError::InvalidMessage)?;
            match msg {
                CollectiveMessage::Init(id) => {
                    state.init_session(id);
                    session_id = Some(id);
                }
                CollectiveMessage::Request(request_id, remote_request) => {
                    let session_id = session_id.ok_or(GlobalCollectiveError::FirstMsgNotInit)?;
                    state
                        .process_request(session_id, request_id, remote_request)
                        .await;
                }
            }
        }

        Ok(())
    }
}

/// Start a global orchestrator with WebSocket on the given port
pub async fn start_global_orchestrator(port: u16) {
    let server = WsServer::new(port);
    let res = GlobalOrchestrator::start(os_shutdown_signal(), server).await;
    if let Err(err) = res {
        log::error!("Global Collective Orchestrator error: {err:?}");
    }
}