diff --git a/Cargo.lock b/Cargo.lock index 9ab1b48..a692dcf 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -59,6 +59,12 @@ dependencies = [ "windows-sys", ] +[[package]] +name = "anyhow" +version = "1.0.80" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5ad32ce52e4161730f7098c077cd2ed6229b5804ccf99e5366be1ab72a98b4e1" + [[package]] name = "base16ct" version = "0.2.0" @@ -447,6 +453,7 @@ dependencies = [ name = "tlparse" version = "0.3.5" dependencies = [ + "anyhow", "base16ct", "clap", "fxhash", diff --git a/Cargo.toml b/Cargo.toml index ed5d9d7..6a265e4 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -9,6 +9,7 @@ license = "BSD-3-Clause" # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html [dependencies] +anyhow = "1.0.80" base16ct = "0.2.0" clap = { version = "4.5.0", features = ["derive"] } fxhash = "0.2.1" diff --git a/src/main.rs b/src/main.rs index 4ec14cb..e8b4865 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1,3 +1,4 @@ +use anyhow::anyhow; use base16ct; use clap::Parser; use core::hash::BuildHasherDefault; @@ -5,6 +6,7 @@ use fxhash::{FxHashMap, FxHasher}; use html_escape::encode_text; use indexmap::IndexMap; use md5::{Digest, Md5}; +use std::ffi::{OsStr, OsString}; use regex::Regex; use std::fmt::{self, Display, Formatter}; @@ -18,7 +20,6 @@ use tinytemplate::TinyTemplate; use indicatif::{MultiProgress, ProgressBar, ProgressStyle}; use once_cell::sync::Lazy; use serde::{Deserialize, Serialize}; -use std::process::ExitCode; use std::sync::Mutex; use std::time::Instant; @@ -64,8 +65,8 @@ static TEMPLATE_INDEX: &str = r#" struct Cli { path: PathBuf, /// Output directory, defaults to `tl_out` - #[arg(short)] - out: Option, + #[arg(short, default_value = "tl_out")] + out: PathBuf, /// Delete out directory if it already exists #[arg(long)] overwrite: bool, @@ -217,7 +218,7 @@ struct EmptyMetadata {} #[derive(Debug, Deserialize)] struct DynamoOutputGraphMetadata { - sizes: Option>>, + _sizes: Option>>, } #[derive(Debug, Deserialize)] @@ -227,7 +228,7 @@ struct DynamoStartMetadata { #[derive(Debug, Deserialize)] struct InductorOutputCodeMetadata { - filename: Option, + filename: Option, } #[derive(Debug, Deserialize)] @@ -244,7 +245,7 @@ struct Envelope { optimize_ddp_split_graph: Option, optimize_ddp_split_child: Option, compiled_autograd_graph: Option, - dynamo_guards: Option, + _dynamo_guards: Option, aot_forward_graph: Option, aot_backward_graph: Option, aot_joint_graph: Option, @@ -259,10 +260,10 @@ struct IndexContext { stack_trie_html: String, } -fn main() -> ExitCode { +fn main() -> anyhow::Result<()> { let cli = Cli::parse(); let path = cli.path; - let out_path = cli.out.unwrap_or(PathBuf::from("tl_out")); + let out_path = cli.out; if out_path.exists() { if !cli.overwrite { @@ -271,17 +272,17 @@ fn main() -> ExitCode { out_path.display() ); } - fs::remove_dir_all(&out_path).unwrap(); + fs::remove_dir_all(&out_path)?; } - fs::create_dir(&out_path).unwrap(); + fs::create_dir(&out_path)?; - let file = File::open(path).unwrap(); - let metadata = file.metadata().unwrap(); + let file = File::open(path)?; + let metadata = file.metadata()?; let file_size = metadata.len(); let multi = MultiProgress::new(); let pb = multi.add(ProgressBar::new(file_size)); pb.set_style(ProgressStyle::default_bar() - .template("{spinner:.green} [{elapsed_precise}] [{wide_bar:.cyan/blue}] {bytes}/{total_bytes} [{bytes_per_sec}] ({eta})").unwrap() + .template("{spinner:.green} [{elapsed_precise}] [{wide_bar:.cyan/blue}] {bytes}/{total_bytes} [{bytes_per_sec}] ({eta})")? .progress_chars("#>-")); let spinner = multi.add(ProgressBar::new_spinner()); let reader = io::BufReader::new(file); @@ -292,8 +293,7 @@ fn main() -> ExitCode { r"(?\d+)", r"(?[^:]+):(?\d+)\] ", r"(?.)" - )) - .unwrap(); + ))?; let mut stack_trie = StackTrieNode::default(); @@ -377,21 +377,26 @@ fn main() -> ExitCode { } }; - // TODO: borrow only here - let compile_id_dir = e + let compile_id_dir: PathBuf = e .compile_id - .clone() - .map_or(format!("unknown_{}", lineno), |e: CompileId| { - format!("{}_{}_{}", e.frame_id, e.frame_compile_id, e.attempt) - }); + .as_ref() + .map_or( + format!("unknown_{lineno}"), + |CompileId { + frame_id, + frame_compile_id, + attempt, + }| { format!("{frame_id}_{frame_compile_id}_{attempt}") }, + ) + .into(); let subdir = out_path.join(&compile_id_dir); - fs::create_dir_all(&subdir).unwrap(); + fs::create_dir_all(&subdir)?; let mut payload = String::new(); if let Some(expect) = e.has_payload { let mut first = true; - while let Some((payload_lineno, payload_line)) = + while let Some((_payload_lineno, payload_line)) = iter.next_if(|(_, l)| l.starts_with('\t')) { // Careful! Distinguish between missing EOL and not @@ -429,53 +434,55 @@ fn main() -> ExitCode { }; }; - let mut write_dump = |filename: &str, sentinel: Option| { - if let Some(_r) = sentinel { - let f = subdir.join(filename); - fs::write(&f, &payload).unwrap(); - compile_directory.push(Path::new(&compile_id_dir).join(filename)); - } - }; + let mut write_dump = + |filename: &str, sentinel: Option| -> anyhow::Result<()> { + if sentinel.is_some() { + let f = subdir.join(filename); + fs::write(&f, &payload)?; + compile_directory.push(compile_id_dir.join(filename)); + } + Ok(()) + }; - write_dump("optimize_ddp_split_graph.txt", e.optimize_ddp_split_graph); - write_dump("compiled_autograd_graph.txt", e.compiled_autograd_graph); - write_dump("aot_forward_graph.txt", e.aot_forward_graph); - write_dump("aot_backward_graph.txt", e.aot_backward_graph); - write_dump("aot_joint_graph.txt", e.aot_joint_graph); - write_dump("inductor_post_grad_graph.txt", e.inductor_post_grad_graph); + write_dump("optimize_ddp_split_graph.txt", e.optimize_ddp_split_graph)?; + write_dump("compiled_autograd_graph.txt", e.compiled_autograd_graph)?; + write_dump("aot_forward_graph.txt", e.aot_forward_graph)?; + write_dump("aot_backward_graph.txt", e.aot_backward_graph)?; + write_dump("aot_joint_graph.txt", e.aot_joint_graph)?; + write_dump("inductor_post_grad_graph.txt", e.inductor_post_grad_graph)?; - if let Some(_metadata) = e.dynamo_output_graph { + if e.dynamo_output_graph.is_some() { // TODO: dump sizes let filename = "dynamo_output_graph.txt"; let f = subdir.join(&filename); - fs::write(&f, &payload).unwrap(); - compile_directory.push(Path::new(&compile_id_dir).join(filename)); + fs::write(&f, &payload)?; + compile_directory.push(compile_id_dir.join(filename)); } if let Some(metadata) = e.inductor_output_code { - let filename = match metadata.filename { - Some(p) => - // Bah, where's pattern guards when you need 'em - { - match Path::new(&p).file_stem() { - Some(stem) => { - format!("inductor_output_code_{}.txt", stem.to_str().unwrap()) - } - None => "inductor_output_code.txt".to_string(), - } - } - None => "inductor_output_code.txt".to_string(), - }; + let filename = metadata + .filename + .as_ref() + .and_then(|p| Path::file_stem(p)) + .map_or_else( + || PathBuf::from("inductor_output_code.txt"), + |stem| { + let mut r = OsString::from("inductor_output_code_"); + r.push(stem); + r.push(OsStr::new(".txt")); + r.into() + }, + ); let f = subdir.join(&filename); - fs::write(&f, &payload).unwrap(); - compile_directory.push(Path::new(&compile_id_dir).join(filename)); + fs::write(&f, &payload)?; + compile_directory.push(compile_id_dir.join(filename)); } if let Some(metadata) = e.optimize_ddp_split_child { let filename = format!("optimize_ddp_split_child_{}.txt", metadata.name); let f = subdir.join(&filename); - fs::write(&f, &payload).unwrap(); - compile_directory.push(Path::new(&compile_id_dir).join(filename)); + fs::write(&f, &payload)?; + compile_directory.push(compile_id_dir.join(filename)); } } pb.finish_with_message("done"); @@ -485,7 +492,7 @@ fn main() -> ExitCode { let mut tt = TinyTemplate::new(); tt.add_formatter("format_unescaped", tinytemplate::format_unescaped); - tt.add_template("index.html", TEMPLATE_INDEX).unwrap(); + tt.add_template("index.html", TEMPLATE_INDEX)?; let index_context = IndexContext { css: CSS, directory: directory @@ -496,12 +503,11 @@ fn main() -> ExitCode { }; fs::write( out_path.join("index.html"), - tt.render("index.html", &index_context).unwrap(), - ) - .unwrap(); + tt.render("index.html", &index_context)?, + )?; if !cli.no_browser { - opener::open(out_path.join("index.html")).unwrap(); + opener::open(out_path.join("index.html"))?; } // other_rank is included here because you should only have logs from one rank when @@ -510,8 +516,8 @@ fn main() -> ExitCode { && (stats.fail_glog + stats.fail_json + stats.fail_payload_md5 + stats.other_rank > 0) { // Report something went wrong - ExitCode::from(1) - } else { - ExitCode::from(0) + return Err(anyhow!("Something went wrong")); } + + Ok(()) }