1use std::ffi::OsString;
8use std::path::PathBuf;
9
10use anyhow::Result;
11use clap::{Args, CommandFactory, FromArgMatches, Parser, Subcommand, ValueEnum};
12
13use super::trace_cli::{self, QueryLabelFilter, QueryOptions, TraceQueryKind, QUERY_MAX_LIMIT};
14
15#[derive(Parser, Debug)]
17#[command(
18 name = "candle-graph",
19 version,
20 about = "Import and visualize Candle execution traces",
21 long_about = "Capability-qualified evidence and atomic bundles from candle-graph/trace/10 runs."
22)]
23pub struct Cli {
24 #[command(subcommand)]
25 command: Command,
26}
27
28impl Cli {
29 pub fn run(self) -> Result<()> {
31 match self.command {
32 Command::Import(import) => {
33 trace_cli::run_import(&import.trace, import.output.as_deref())
34 }
35 #[cfg(feature = "visualizer")]
36 Command::View(view) => {
37 trace_cli::run_view(&view.trace, &view.output, view.nsight_dir.as_deref())
38 }
39 Command::Summary(summary) => trace_cli::run_summary(
40 &summary.input,
41 summary.output.as_deref(),
42 summary.require_valid,
43 ),
44 Command::Query(query) => {
45 let options = query.options();
46 trace_cli::run_query(
47 &query.input,
48 query.kind.into(),
49 &options,
50 query.output.as_deref(),
51 )
52 }
53 Command::Overview(overview) => {
54 trace_cli::run_overview(&overview.input, overview.output.as_deref())
55 }
56 Command::Compare(compare) => trace_cli::run_compare(
57 &compare.baseline,
58 &compare.candidate,
59 compare.unverified_traces,
60 compare.require_eligible,
61 compare.output.as_deref(),
62 ),
63 Command::Report(report) => trace_cli::run_report(
64 &report.trace,
65 report.nsight_dir.as_deref(),
66 &report.bundle,
67 report.output.as_deref(),
68 ),
69 Command::Verify(verify) => {
70 trace_cli::run_verify(&verify.bundle, verify.semantic, verify.output.as_deref())
71 }
72 Command::Protocol(protocol) => trace_cli::run_protocol(protocol.output.as_deref()),
73 Command::CampaignStatus(status) => {
74 trace_cli::run_campaign_status(&status.manifest, status.output.as_deref())
75 }
76 Command::Series(series) => trace_cli::run_series(
77 series.manifest.as_deref(),
78 &series.bundle,
79 series.label_prefix.as_deref(),
80 series.output.as_deref(),
81 ),
82 }
83 }
84
85 pub fn parse_as_cargo_subcommand() -> Self {
93 let mut argv: Vec<OsString> = std::env::args_os().collect();
94 if argv
95 .get(1)
96 .is_some_and(|argument| argument == "candle-graph")
97 {
98 argv.remove(1);
99 }
100 let command = Self::command()
101 .name("cargo-candle-graph")
102 .bin_name("cargo candle-graph")
103 .about("Import and analyze candle-graph execution trace files");
104 let matches = command.get_matches_from(argv);
105 match Self::from_arg_matches(&matches) {
106 Ok(cli) => cli,
107 Err(error) => error.exit(),
108 }
109 }
110}
111
112pub fn command_catalog() -> Vec<serde_json::Value> {
114 Cli::command()
115 .get_subcommands()
116 .map(|subcommand| {
117 serde_json::json!({
118 "name": subcommand.get_name(),
119 "about": subcommand.get_about().map(|about| about.to_string()),
120 })
121 })
122 .collect()
123}
124
125#[derive(Subcommand, Debug)]
126enum Command {
127 Import(ImportArgs),
129 #[cfg(feature = "visualizer")]
131 View(ViewArgs),
132 Summary(SummaryArgs),
134 Query(QueryArgs),
136 Overview(OverviewArgs),
138 Compare(CompareArgs),
140 Report(ReportArgs),
142 Verify(VerifyArgs),
144 Protocol(ProtocolArgs),
146 CampaignStatus(CampaignStatusArgs),
148 Series(SeriesArgs),
150}
151
152#[derive(Args, Debug)]
153struct ImportArgs {
154 #[arg(value_name = "INPUT")]
156 trace: PathBuf,
157 #[arg(long, short, value_name = "FILE")]
158 output: Option<PathBuf>,
159}
160
161#[cfg(feature = "visualizer")]
162#[derive(Args, Debug)]
163struct ViewArgs {
164 #[arg(value_name = "TRACE")]
166 trace: PathBuf,
167 #[arg(long, value_name = "FILE")]
168 output: PathBuf,
169 #[arg(long, value_name = "DIR")]
171 nsight_dir: Option<PathBuf>,
172}
173
174#[derive(Args, Debug)]
175struct SummaryArgs {
176 #[arg(value_name = "INPUT")]
178 input: PathBuf,
179 #[arg(long)]
181 require_valid: bool,
182 #[arg(long, short, value_name = "FILE")]
183 output: Option<PathBuf>,
184}
185
186#[derive(Args, Debug)]
187struct QueryArgs {
188 #[arg(value_name = "INPUT")]
190 input: PathBuf,
191 #[arg(long, value_enum)]
192 kind: CliTraceQueryKind,
193 #[arg(long, value_name = "S", conflicts_with = "label_prefix")]
195 label: Option<String>,
196 #[arg(long, value_name = "S")]
198 label_prefix: Option<String>,
199 #[arg(
201 long,
202 value_name = "N",
203 value_parser = parse_query_limit,
204 conflicts_with = "all"
205 )]
206 limit: Option<usize>,
207 #[arg(long, value_name = "N", conflicts_with = "all")]
209 offset: Option<usize>,
210 #[arg(long, conflicts_with_all = ["limit", "offset"])]
212 all: bool,
213 #[arg(long, short, value_name = "FILE")]
214 output: Option<PathBuf>,
215}
216
217impl QueryArgs {
218 fn filter(&self) -> Option<QueryLabelFilter> {
219 self.label
220 .clone()
221 .map(QueryLabelFilter::Exact)
222 .or_else(|| self.label_prefix.clone().map(QueryLabelFilter::Prefix))
223 }
224
225 fn options(&self) -> QueryOptions {
226 QueryOptions {
227 filter: self.filter(),
228 limit: self.limit,
229 offset: self.offset,
230 all: self.all,
231 }
232 }
233}
234
235fn parse_query_limit(value: &str) -> std::result::Result<usize, String> {
236 let limit = value
237 .parse::<usize>()
238 .map_err(|_| format!("limit must be an integer in 1..={QUERY_MAX_LIMIT}"))?;
239 if !(1..=QUERY_MAX_LIMIT).contains(&limit) {
240 return Err(format!("limit must be in 1..={QUERY_MAX_LIMIT}"));
241 }
242 Ok(limit)
243}
244
245#[derive(Args, Debug)]
246struct OverviewArgs {
247 #[arg(value_name = "INPUT")]
249 input: PathBuf,
250 #[arg(long, short, value_name = "FILE")]
251 output: Option<PathBuf>,
252}
253
254#[derive(Args, Debug)]
255struct CompareArgs {
256 #[arg(long, required = true, num_args = 1.., value_name = "BUNDLE")]
258 baseline: Vec<PathBuf>,
259 #[arg(long, required = true, num_args = 1.., value_name = "BUNDLE")]
261 candidate: Vec<PathBuf>,
262 #[arg(long)]
264 unverified_traces: bool,
265 #[arg(long)]
267 require_eligible: bool,
268 #[arg(long, short, value_name = "FILE")]
269 output: Option<PathBuf>,
270}
271
272#[derive(Args, Debug)]
273struct ReportArgs {
274 #[arg(value_name = "TRACE")]
276 trace: PathBuf,
277 #[arg(long, value_name = "DIR")]
278 nsight_dir: Option<PathBuf>,
279 #[arg(long, value_name = "DIR")]
280 bundle: PathBuf,
281 #[arg(long, short, value_name = "FILE")]
283 output: Option<PathBuf>,
284}
285
286#[derive(Args, Debug)]
287struct VerifyArgs {
288 #[arg(value_name = "BUNDLE")]
290 bundle: PathBuf,
291 #[arg(long)]
293 semantic: bool,
294 #[arg(long, short, value_name = "FILE")]
295 output: Option<PathBuf>,
296}
297
298#[derive(Args, Debug)]
299struct ProtocolArgs {
300 #[arg(long, short, value_name = "FILE")]
301 output: Option<PathBuf>,
302}
303
304#[derive(Args, Debug)]
305struct CampaignStatusArgs {
306 #[arg(long, value_name = "FILE")]
308 manifest: PathBuf,
309 #[arg(long, short, value_name = "FILE")]
310 output: Option<PathBuf>,
311}
312
313#[derive(Args, Debug)]
314struct SeriesArgs {
315 #[arg(
317 long,
318 value_name = "FILE",
319 conflicts_with = "bundle",
320 required_unless_present = "bundle"
321 )]
322 manifest: Option<PathBuf>,
323 #[arg(long, value_name = "DIR", num_args = 1..)]
325 bundle: Vec<PathBuf>,
326 #[arg(long, value_name = "P")]
328 label_prefix: Option<String>,
329 #[arg(long, short, value_name = "FILE")]
330 output: Option<PathBuf>,
331}
332
333#[derive(Debug, Clone, Copy, PartialEq, Eq, ValueEnum)]
334enum CliTraceQueryKind {
335 Labels,
336 SlowestHost,
337 SlowestDevice,
338 Activations,
339 Heaviest,
340 Memory,
341 Spans,
342 Tensors,
343 TensorStats,
344 Gradients,
345 Capabilities,
346 GpuStatus,
347 GpuCorrelation,
348 GpuPhases,
349 GpuKernels,
350 GpuAttributionGaps,
351}
352
353impl From<CliTraceQueryKind> for TraceQueryKind {
354 fn from(kind: CliTraceQueryKind) -> Self {
355 match kind {
356 CliTraceQueryKind::Labels => Self::Labels,
357 CliTraceQueryKind::SlowestHost => Self::SlowestHost,
358 CliTraceQueryKind::SlowestDevice => Self::SlowestDevice,
359 CliTraceQueryKind::Activations => Self::Activations,
360 CliTraceQueryKind::Heaviest => Self::Heaviest,
361 CliTraceQueryKind::Memory => Self::Memory,
362 CliTraceQueryKind::Spans => Self::Spans,
363 CliTraceQueryKind::Tensors => Self::Tensors,
364 CliTraceQueryKind::TensorStats => Self::TensorStats,
365 CliTraceQueryKind::Gradients => Self::Gradients,
366 CliTraceQueryKind::Capabilities => Self::Capabilities,
367 CliTraceQueryKind::GpuStatus => Self::GpuStatus,
368 CliTraceQueryKind::GpuCorrelation => Self::GpuCorrelation,
369 CliTraceQueryKind::GpuPhases => Self::GpuPhases,
370 CliTraceQueryKind::GpuKernels => Self::GpuKernels,
371 CliTraceQueryKind::GpuAttributionGaps => Self::GpuAttributionGaps,
372 }
373 }
374}