Skip to main content

vtcode_core/cli/
pods_commands.rs

1//! CLI commands for GPU pod management.
2
3use crate::cli::PodsCommands;
4use crate::pods::{
5    PodGpu, PodManager, PodStartRequest, PodState, PodStatusDetail, PodStatusReport,
6};
7use crate::utils::colors::{bold, cyan, green, underline, yellow};
8use anyhow::{Result, anyhow};
9
10/// Handle GPU pod commands.
11pub async fn handle_pods_command(command: PodsCommands) -> Result<()> {
12    let manager = PodManager::new()?;
13
14    match command {
15        PodsCommands::Start {
16            name,
17            model,
18            pod_name,
19            ssh,
20            gpus,
21            models_path,
22            profile,
23            gpus_count,
24            memory,
25            context,
26        } => {
27            let request = PodStartRequest {
28                pod_name,
29                ssh,
30                gpus: parse_gpu_entries(&gpus)?,
31                models_path,
32                name,
33                model,
34                profile,
35                requested_gpu_count: gpus_count,
36                memory,
37                context,
38            };
39            let result = manager.start_model(request).await?;
40            print_start_result(&result);
41        }
42        PodsCommands::Stop { name } => {
43            let stopped = manager
44                .stop_model(&name)
45                .await?
46                .ok_or_else(|| anyhow!("unknown model '{}'", name))?;
47            println!(
48                "{} Stopped {} (pid {})",
49                green("✓"),
50                cyan(&name),
51                stopped.pid
52            );
53        }
54        PodsCommands::StopAll => {
55            let count = manager.stop_all_models().await?;
56            println!("{} Stopped {} model(s)", green("✓"), count);
57        }
58        PodsCommands::List => {
59            let report = manager.list_models().await?;
60            print_list_report(&report);
61        }
62        PodsCommands::Logs { name } => {
63            manager.stream_logs(&name).await?;
64        }
65        PodsCommands::KnownModels => {
66            let report = manager.known_models().await?;
67            print_known_models(&report, manager.load_state().await?.active_pod.as_ref());
68        }
69    }
70
71    Ok(())
72}
73
74fn parse_gpu_entries(raw: &[String]) -> Result<Vec<PodGpu>> {
75    raw.iter()
76        .map(|entry| {
77            let (id, name) = entry
78                .split_once(':')
79                .or_else(|| entry.split_once('='))
80                .ok_or_else(|| anyhow!("invalid GPU entry '{}'; expected ID:NAME", entry))?;
81            let id = id
82                .trim()
83                .parse::<u32>()
84                .map_err(|_| anyhow!("invalid GPU id '{}'", id.trim()))?;
85            let name = name.trim();
86            if name.is_empty() {
87                return Err(anyhow!("GPU name cannot be empty"));
88            }
89            Ok(PodGpu {
90                id,
91                name: name.to_string(),
92            })
93        })
94        .collect()
95}
96
97fn print_start_result(result: &crate::pods::PodStartResult) {
98    println!("{} Started {}", green("✓"), bold(&result.entry.model));
99    println!("  Pod: {}", cyan(&result.pod.name));
100    println!("  Profile: {}", cyan(&result.entry.profile));
101    println!("  PID: {}", cyan(&result.entry.pid.to_string()));
102    println!("  Port: {}", cyan(&result.entry.port.to_string()));
103    println!(
104        "  GPUs: {}",
105        cyan(
106            &result
107                .entry
108                .gpu_ids
109                .iter()
110                .map(u32::to_string)
111                .collect::<Vec<_>>()
112                .join(", ")
113        )
114    );
115    println!("  Launch: {}", yellow(&result.launch_command));
116}
117
118fn print_list_report(report: &PodStatusReport) {
119    println!("{}", underline(&bold("Active Pod")));
120    println!("Pod: {}", cyan(&report.pod_name));
121    println!();
122
123    if report.entries.is_empty() {
124        println!("{}", yellow("No running models"));
125        return;
126    }
127
128    for entry in &report.entries {
129        println!(
130            "{} {} | model={} | port={} | pid={} | gpus={}",
131            status_symbol(entry.status),
132            bold(&entry.name),
133            entry.model,
134            entry.port,
135            entry.pid,
136            entry
137                .gpu_ids
138                .iter()
139                .map(u32::to_string)
140                .collect::<Vec<_>>()
141                .join(", ")
142        );
143    }
144}
145
146fn print_known_models(report: &crate::pods::KnownModelsReport, active_pod: Option<&PodState>) {
147    println!("{}", underline(&bold("Known Models")));
148    if let Some(pod) = active_pod {
149        println!("Pod: {}", cyan(&pod.name));
150        println!("GPUs: {}", cyan(&pod.gpu_count().to_string()));
151    }
152    println!();
153
154    println!("{}", green("Compatible"));
155    for model in &report.compatible {
156        print_model_detail(model, green("  ✓"));
157    }
158
159    println!();
160    println!("{}", yellow("Incompatible"));
161    for model in &report.incompatible {
162        print_model_detail(model, yellow("  •"));
163    }
164}
165
166fn print_model_detail(model: &PodStatusDetail, prefix: impl std::fmt::Display) {
167    println!(
168        "{} {} ({}, {} GPU{})",
169        prefix,
170        cyan(&model.name),
171        model.model,
172        model.gpu_count,
173        if model.gpu_count == 1 { "" } else { "s" }
174    );
175}
176
177fn status_symbol(status: crate::pods::PodHealth) -> &'static str {
178    match status {
179        crate::pods::PodHealth::Running => "✓",
180        crate::pods::PodHealth::Starting => "…",
181        crate::pods::PodHealth::Crashed => "✗",
182        crate::pods::PodHealth::Dead => "•",
183    }
184}