moose 0.2.2

Encrypted learning and data processing framework
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

use moose::computation::Role;
use moose::execution::Identity;
use moose::networking::tcpstream::TcpStreamNetworking;
use moose::prelude::*;
use moose::storage::local::LocalAsyncStorage;
use moose::tokio;
use std::collections::HashMap;
use std::convert::TryFrom;
use std::sync::Arc;
use structopt::StructOpt;
use tracing_subscriber::layer::SubscriberExt;
use tracing_subscriber::util::SubscriberInitExt;

#[derive(Debug, StructOpt, Clone)]
pub struct Opt {
    #[structopt(long)]
    comp: String,

    #[structopt(long)]
    session_id: String,

    #[structopt(long)]
    placement: String,

    #[structopt(long)]
    role_assignment: String,

    #[structopt(long)]
    hosts: String,
}

fn init_tracer() {
    let fmt_layer = Some(tracing_subscriber::fmt::Layer::default());
    tracing_subscriber::registry()
        .with(tracing_subscriber::EnvFilter::from_default_env()) // The tracing formatter defaults to the max log level set by RUST_LOG
        .with(fmt_layer)
        .try_init()
        .unwrap_or_else(|e| println!("Failed to initialize telemetry subscriber: {}", e));
}

#[tokio::main]
async fn main() -> Result<(), Box<dyn std::error::Error>> {
    init_tracer();
    tracing::info!("starting up");
    let opt = Opt::from_args();

    let hosts: HashMap<String, String> = serde_json::from_str(&opt.hosts)?;

    let computation_bytes = std::fs::read(&opt.comp)?;

    let computation = Computation::from_msgpack(computation_bytes)?;

    let storage = Arc::new(LocalAsyncStorage::default());

    let networking = TcpStreamNetworking::new(&opt.placement, hosts).await?;
    let networking = Arc::new(networking);

    let arguments = HashMap::new();

    let session_id = moose::computation::SessionId::try_from(opt.session_id.as_ref())?;

    let role_assignment_map: HashMap<String, String> = serde_json::from_str(&opt.role_assignment)?;
    let role_assignment: HashMap<Role, Identity> = role_assignment_map
        .iter()
        .map(|(key, value)| {
            let role = Role::from(key);
            let identity = Identity::from(value);
            (role, identity)
        })
        .collect();

    let moose_session = moose::execution::AsyncSession::new(
        session_id,
        arguments,
        role_assignment.clone(),
        networking,
        storage,
    );

    let own_identity: Identity = Identity::from(&opt.placement);
    let mut executor = AsyncExecutor::default();

    let outputs_handle = executor.run_computation(
        &computation,
        &role_assignment,
        &own_identity,
        &moose_session,
    )?;

    let moose_session_handle = moose_session.into_handle()?;
    tracing::info!("joining on tasks");
    moose_session_handle.join_on_first_error().await?;

    let mut outputs = HashMap::new();

    tracing::info!("collecting outputs");
    for (output_name, output_future) in outputs_handle {
        tracing::info!("awaiting output: {}", output_name);
        let value = output_future.await.unwrap();
        tracing::info!("got output: {}", output_name);
        outputs.insert(output_name, value);
    }

    tracing::info!("outputs: {:?}", outputs);
    Ok(())
}