mod cli;
mod commands;
mod models;
use crate::cli::{
Cli, Commands, CreateCommands, DeleteCommands, DeploymentCommands, GetCommands, StopCommands,
TrainCommands, WorkCommands,
};
use clap::Parser;
use orign::db;
use std::error::Error;
use tracing_subscriber;
#[tokio::main]
async fn main() -> Result<(), Box<dyn Error>> {
tracing_subscriber::fmt::init();
let cli = Cli::parse();
match cli.command {
Commands::Serve { host, port } => {
commands::serve_cmd::execute(host, port).await?;
}
Commands::Chat {
model,
msg,
image,
framework,
adapter,
} => {
commands::chat_cmd::execute(model, msg, image, adapter, framework).await?;
}
Commands::Trigger { buffer_name } => {
commands::trigger_cmd::trigger_buffer(buffer_name).await?;
}
Commands::Create { command } => match command {
CreateCommands::Buffer { command } => {
commands::create_cmd::create_replay_buffer(command).await?;
}
},
Commands::Get { command } => match command {
GetCommands::Deployments { id } => {
commands::get_cmd::get_deployment(id).await?;
}
GetCommands::Trainings { id } => {
commands::get_cmd::get_training(id).await?;
}
GetCommands::Buffers { name } => {
commands::get_cmd::get_buffer(name).await?;
}
GetCommands::Models {} => {
commands::get_cmd::get_models().await?;
}
GetCommands::Datasets {} => {
commands::get_cmd::get_datasets().await?;
}
GetCommands::Adapters {} => {
commands::get_cmd::get_adapters().await?;
}
},
Commands::Delete { command } => match command {
DeleteCommands::Deployment { id } => {
commands::delete_cmd::delete_deployment(id).await?;
}
DeleteCommands::Buffer { name } => {
commands::delete_cmd::delete_buffer(name).await?;
}
DeleteCommands::Adapter { name } => {
commands::delete_cmd::delete_adapter(name).await?;
}
},
Commands::Stop { command } => match command {
StopCommands::Training { id } => {
commands::train_cmd::stop_training(id).await?;
}
},
Commands::Send {
buffer_name,
file,
train,
} => {
commands::send_cmd::send_buffer(buffer_name, file, train).await?;
}
Commands::Prepare {
dataset_type,
url,
split_ratio,
base_path,
image_path,
num_workers,
} => {
commands::prepare_cmd::execute_prepare(
dataset_type,
url,
split_ratio,
base_path,
image_path,
num_workers,
)
.await?;
}
Commands::Train { command } => match command {
TrainCommands::Swift(args) => {
commands::train_cmd::run_swift_training(args).await?;
}
},
Commands::Report {
watch_dir,
save_dir,
debounce_secs,
} => {
commands::report_cmd::execute_report(watch_dir, save_dir, debounce_secs).await?;
}
Commands::Work { command } => match command {
WorkCommands::ReportTrainings { k8s_namespace } => {
let db = db::init_db().await?;
println!("Database connected successfully.");
commands::work_cmd::execute_report_jobs(&db, &k8s_namespace).await?;
}
},
Commands::Login => {
commands::login_cmd::execute().await?;
}
}
Ok(())
}