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

use async_trait::async_trait;

use xaynet_core::{
    common::RoundParameters,
    mask::Model,
    SumDict,
    SumParticipantPublicKey,
    UpdateSeedDict,
};

/// A trait used by the [`StateMachine`] to emit notifications upon
/// certain events.
///
/// [`StateMachine`]: crate::StateMachine
pub trait Notify {
    /// Emit a notification when a new round of federated learning
    /// starts
    fn new_round(&mut self) {}
    /// Emit a notification when the participant has been selected for
    /// the sum task
    fn sum(&mut self) {}
    /// Emit a notification when the participant has been selected for
    /// the update task
    fn update(&mut self) {}
    /// Emit a notification when the participant is not selected for
    /// any task and is waiting for another round to start
    fn idle(&mut self) {}
    /// Emit a notification when the participant should populate the
    /// model store (see [`ModelStore`]).
    fn load_model(&mut self) {}
}

/// A trait used by the [`StateMachine`] to load the model trained by
/// the participant, when it has been selected for the update task.
///
/// [`StateMachine`]: crate::StateMachine
#[async_trait]
pub trait ModelStore {
    type Error: std::error::Error;
    type Model: AsRef<Model> + Send;

    /// Attempt to load the model. If the model is not yet available,
    /// `Ok(None)` should be returned.
    async fn load_model(&mut self) -> Result<Option<Self::Model>, Self::Error>;
}

/// A trait used by the [`StateMachine`] to communicate with the
/// Xaynet coordinator.
///
/// [`StateMachine`]: crate::StateMachine
#[async_trait]
pub trait XaynetClient {
    type Error: std::error::Error;

    /// Retrieve the current round parameters
    async fn get_round_params(&mut self) -> Result<RoundParameters, Self::Error>;

    /// Retrieve the current sum dictionary, if available.
    async fn get_sums(&mut self) -> Result<Option<SumDict>, Self::Error>;

    /// Retrieve the current seed dictionary for the given sum
    /// participant, if available.
    async fn get_seeds(
        &mut self,
        pk: SumParticipantPublicKey,
    ) -> Result<Option<UpdateSeedDict>, Self::Error>;

    /// Retrieve the current global model, if available.
    async fn get_model(&mut self) -> Result<Option<Model>, Self::Error>;

    /// Send an encrypted and signed PET message to the coordinator.
    async fn send_message(&mut self, msg: Vec<u8>) -> Result<(), Self::Error>;
}