burn_central_cli/
cli.rs

1use clap::{Parser, Subcommand};
2
3use crate::app_config::Environment;
4use crate::commands;
5use crate::commands::default_command;
6use crate::config::Config;
7use crate::context::CliContext;
8use crate::tools::terminal::Terminal;
9
10#[derive(Parser, Debug)]
11#[clap(author, version, about, long_about = None)]
12pub struct CliArgs {
13    #[command(subcommand)]
14    pub command: Option<Commands>,
15
16    /// Use development environment (localhost:9001) with separate dev credentials
17    #[arg(long, action = clap::ArgAction::SetTrue, hide = true)]
18    pub dev: bool,
19}
20
21#[derive(Subcommand, Debug)]
22pub enum Commands {
23    /// Run a training or inference locally or trigger a remote run.
24    Train(commands::training::TrainingArgs),
25
26    /// Package your project for running on a remote machine.
27    Package(commands::package::PackageArgs),
28    /// Log in to the Burn Central server.
29    Login(commands::login::LoginArgs),
30    /// Initialize a new project or reinitialize an existing one.
31    Init(commands::init::InitArgs),
32    /// Unlink the burn central project from this repository.
33    Unlink,
34    /// Display current user information.
35    Me,
36    /// Display current project information.
37    Project,
38}
39
40pub fn cli_main() {
41    let args = CliArgs::parse();
42
43    let environment = Environment::from_dev_flag(args.dev);
44
45    let config = Config {
46        api_endpoint: if args.dev {
47            "http://localhost:9001/".to_string()
48        } else {
49            "https://heat.tracel.ai/api/".to_string()
50        },
51    };
52
53    let terminal = Terminal::new();
54
55    if args.dev {
56        terminal
57            .print_warning("Running in development mode - using local server and dev credentials");
58    }
59
60    let context = CliContext::new(terminal.clone(), &config, environment).init();
61
62    let cli_res = match args.command {
63        Some(command) => handle_command(command, context),
64        None => default_command(context),
65    };
66
67    if let Err(e) = cli_res {
68        terminal.cancel_finalize(&format!("{e}"));
69    }
70}
71
72fn handle_command(command: Commands, context: CliContext) -> anyhow::Result<()> {
73    match command {
74        Commands::Train(run_args) => commands::training::handle_command(run_args, context),
75        Commands::Package(package_args) => commands::package::handle_command(package_args, context),
76        Commands::Login(login_args) => commands::login::handle_command(login_args, context),
77        Commands::Init(init_args) => commands::init::handle_command(init_args, context),
78        Commands::Unlink => commands::unlink::handle_command(context),
79        Commands::Me => commands::me::handle_command(context),
80        Commands::Project => commands::project::handle_command(context),
81    }
82}