heat-sdk-cli 0.1.0

Tracel Heat SDK CLI.
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

use clap::Parser;
use heat_sdk::{
    client::{HeatClient, HeatClientConfig, HeatCredentials},
    schemas::{HeatCodeMetadata, ProjectPath, RegisteredHeatFunction},
};
use quote::ToTokens;

use crate::{context::HeatCliContext, generation::backend::BackendType};

/// Run a training remotely.
/// Not yet supported.
#[derive(Parser, Debug)]
pub struct RemoteTrainingRunArgs {
    /// The training functions to run
    #[clap(short = 'f', long="functions", value_delimiter = ' ', num_args = 1.., required = true, help = "<required> The training functions to run. Annotate a training function with #[heat(training)] to register it.")]
    functions: Vec<String>,
    /// Backend to use
    #[clap(short = 'b', long = "backends", value_delimiter = ' ', num_args = 1.., required = true, help = "<required> Backends to use for training.")]
    backends: Vec<BackendType>,
    /// Config files paths
    #[clap(short = 'c', long = "configs", value_delimiter = ' ', num_args = 1.., required = true, help = "<required> Config files paths.")]
    configs: Vec<String>,
    /// The Heat project ID
    // todo: support project name and creating a project if it doesn't exist
    #[clap(
        short = 'p',
        long = "project",
        required = true,
        help = "<required> The Heat project ID."
    )]
    project_path: String,
    /// The Heat API key
    #[clap(
        short = 'k',
        long = "key",
        required = true,
        help = "<required> The Heat API key."
    )]
    key: String,
    /// The Heat API endpoint
    #[clap(
        short = 'e',
        long = "endpoint",
        help = "The Heat API endpoint.",
        default_value = "http://127.0.0.1:9001"
    )]
    pub heat_endpoint: String,
    /// The runner group name
    #[clap(
        short = 'r',
        long = "runner",
        help = "The runner group name.",
        required = true
    )]
    pub runner: String,
}

fn create_heat_client(api_key: &str, url: &str, project_path: &str) -> HeatClient {
    let creds = HeatCredentials::new(api_key.to_owned());
    let client_config = HeatClientConfig::builder(
        creds,
        ProjectPath::try_from(project_path.to_string()).expect("Project path should be valid."),
    )
    .with_endpoint(url)
    .with_num_retries(10)
    .build();
    HeatClient::create(client_config)
        .expect("Should connect to the Heat server and create a client")
}

pub(crate) fn handle_command(
    args: RemoteTrainingRunArgs,
    context: HeatCliContext,
) -> anyhow::Result<()> {
    let heat_client = create_heat_client(&args.key, &args.heat_endpoint, &args.project_path);

    let crates = crate::util::cargo::package::package(
        &context.get_artifacts_dir_path(),
        context.package_name(),
    )?;

    let flags = crate::registry::get_flags();

    let mut registered_functions = Vec::<RegisteredHeatFunction>::new();
    for flag in flags {
        // function token stream to readable string
        let itemfn = syn_serde::json::from_slice::<syn::ItemFn>(flag.token_stream)
            .expect("Should be able to parse token stream.");
        let syn_tree: syn::File =
            syn::parse2(itemfn.into_token_stream()).expect("Should be able to parse token stream.");
        let code_str = prettyplease::unparse(&syn_tree);
        registered_functions.push(RegisteredHeatFunction {
            mod_path: flag.mod_path.to_string(),
            fn_name: flag.fn_name.to_string(),
            proc_type: flag.proc_type.to_string(),
            code: code_str,
        });
    }

    let heat_metadata = HeatCodeMetadata {
        functions: registered_functions,
    };

    let project_version =
        heat_client.upload_new_project_version(context.package_name(), heat_metadata, crates)?;

    heat_client.start_remote_job(
        args.runner,
        project_version,
        format!(
            "run local training --functions {} --backends {} --configs {} --project {} --key {}",
            args.functions.join(" "),
            args.backends
                .into_iter()
                .map(|backend| backend.to_string())
                .collect::<Vec<_>>()
                .join(" "),
            args.configs.join(" "),
            args.project_path,
            args.key
        ),
    )?;

    Ok(())
}