diff --git a/src/cli.rs b/src/cli.rs index 3a4b243..0f9e33e 100644 --- a/src/cli.rs +++ b/src/cli.rs @@ -6,6 +6,24 @@ use std::path::PathBuf; use tlparse::{parse_path, ParseConfig}; +// Main output filename used by both single rank and multi-rank processing +const MAIN_OUTPUT_FILENAME: &str = "index.html"; + +// Helper function to setup output directory (handles overwrite logic) +fn setup_output_directory(out_path: &PathBuf, overwrite: bool) -> anyhow::Result<()> { + if out_path.exists() { + if !overwrite { + bail!( + "Directory {} already exists, use -o OUTDIR to write to another location or pass --overwrite to overwrite the old contents", + out_path.display() + ); + } + fs::remove_dir_all(&out_path)?; + } + fs::create_dir(&out_path)?; + Ok(()) +} + #[derive(Parser)] #[command(author, version, about, long_about = None)] #[command(propagate_version = true)] @@ -47,12 +65,20 @@ pub struct Cli { /// For inductor provenance tracking highlighter #[arg(short, long)] inductor_provenance: bool, + /// Parse all ranks and generate a single unified HTML page + #[arg(long)] + all_ranks_html: bool, } fn main() -> anyhow::Result<()> { let cli = Cli::parse(); + + if cli.all_ranks_html { + return handle_all_ranks(cli); + } + let path = if cli.latest { - let input_path = cli.path; + let input_path = &cli.path; // Path should be a directory if !input_path.is_dir() { bail!( @@ -61,7 +87,7 @@ fn main() -> anyhow::Result<()> { ); } - let last_modified_file = std::fs::read_dir(&input_path) + let last_modified_file = std::fs::read_dir(input_path) .with_context(|| format!("Couldn't access directory {}", input_path.display()))? .flatten() .filter(|f| f.metadata().unwrap().is_file()) @@ -72,45 +98,191 @@ fn main() -> anyhow::Result<()> { }; last_modified_file.path() } else { - cli.path + cli.path.clone() }; - let out_path = cli.out; + let out_path = cli.out.clone(); + setup_output_directory(&out_path, cli.overwrite)?; - if out_path.exists() { - if !cli.overwrite { - bail!( - "Directory {} already exists, use -o OUTDIR to write to another location or pass --overwrite to overwrite the old contents", - out_path.display() - ); - } - fs::remove_dir_all(&out_path)?; + // Use handle_one_rank for single rank processing (don't create directory since it already exists) + handle_one_rank(&path, &out_path, &cli, false)?; + + if !cli.no_browser { + opener::open(out_path.join("index.html"))?; + } + Ok(()) +} + +// Helper function to handle parsing and writing output for a single rank +// Returns the relative path to the main output file within the rank directory +fn handle_one_rank( + rank_path: &PathBuf, + rank_out_dir: &PathBuf, + cli: &Cli, + create_output_dir: bool, +) -> anyhow::Result { + if create_output_dir { + fs::create_dir(rank_out_dir)?; } - fs::create_dir(&out_path)?; let config = ParseConfig { strict: cli.strict, strict_compile_id: cli.strict_compile_id, custom_parsers: Vec::new(), - custom_header_html: cli.custom_header_html, + custom_header_html: cli.custom_header_html.clone(), verbose: cli.verbose, plain_text: cli.plain_text, export: cli.export, inductor_provenance: cli.inductor_provenance, + all_ranks: false, }; - let output = parse_path(&path, config)?; + let output = parse_path(rank_path, config)?; + + let mut main_output_path = None; - for (filename, path) in output { - let out_file = out_path.join(filename); + // Write output files to output directory + for (filename, content) in output { + let out_file = rank_out_dir.join(&filename); if let Some(dir) = out_file.parent() { fs::create_dir_all(dir)?; } - fs::write(out_file, path)?; + fs::write(out_file, content)?; + + // Track the main output file (typically index.html) + if filename.file_name().and_then(|name| name.to_str()) == Some(MAIN_OUTPUT_FILENAME) { + main_output_path = Some(filename); + } } - if !cli.no_browser { - opener::open(out_path.join("index.html"))?; + Ok(main_output_path.unwrap_or_else(|| PathBuf::from(MAIN_OUTPUT_FILENAME))) +} + +// handle_all_ranks function with placeholder landing page +fn handle_all_ranks(cli: Cli) -> anyhow::Result<()> { + let input_path = &cli.path; + + if !input_path.is_dir() { + bail!( + "Input path {} must be a directory when using --all-ranks-html", + input_path.display() + ); + } + + let out_path = &cli.out; + setup_output_directory(out_path, cli.overwrite)?; + + // Find all rank log files in the directory + let rank_files: Vec<_> = std::fs::read_dir(input_path) + .with_context(|| format!("Couldn't access directory {}", input_path.display()))? + .flatten() + .filter(|entry| { + let path = entry.path(); + if !path.is_file() { + return false; + } + + let Some(filename) = path.file_name().and_then(|name| name.to_str()) else { + return false; + }; + + // Only support PyTorch TORCH_TRACE files: dedicated_log_torch_trace_rank_0_hash.log + if !filename.starts_with("dedicated_log_torch_trace_rank_") + || !filename.ends_with(".log") + { + return false; + } + + // Extract rank number from the pattern + let after_prefix = &filename[31..]; // Remove "dedicated_log_torch_trace_rank_" + if let Some(underscore_pos) = after_prefix.find('_') { + let rank_part = &after_prefix[..underscore_pos]; + return !rank_part.is_empty() && rank_part.chars().all(|c| c.is_ascii_digit()); + } + + false + }) + .collect(); + + if rank_files.is_empty() { + bail!( + "No rank log files found in directory {}", + input_path.display() + ); } + + let mut rank_links = Vec::new(); + + // Process each rank file + for rank_file in rank_files { + let rank_path = rank_file.path(); + let rank_name = rank_path + .file_stem() + .and_then(|name| name.to_str()) + .unwrap_or("unknown"); + + // Extract rank number from PyTorch TORCH_TRACE filename + let rank_num = + if let Some(after_prefix) = rank_name.strip_prefix("dedicated_log_torch_trace_rank_") { + if let Some(underscore_pos) = after_prefix.find('_') { + let rank_part = &after_prefix[..underscore_pos]; + if rank_part.is_empty() || !rank_part.chars().all(|c| c.is_ascii_digit()) { + bail!( + "Could not extract rank number from TORCH_TRACE filename: {}", + rank_name + ); + } + rank_part.to_string() + } else { + bail!("Invalid TORCH_TRACE filename format: {}", rank_name); + } + } else { + bail!( + "Filename does not match PyTorch TORCH_TRACE pattern: {}", + rank_name + ); + }; + + println!( + "Processing rank {} from file: {}", + rank_num, + rank_path.display() + ); + + // Create subdirectory for this rank and handle parsing + let rank_out_dir = out_path.join(format!("rank_{rank_num}")); + let main_output_path = handle_one_rank(&rank_path, &rank_out_dir, &cli, true)?; + + // Add link to this rank's page using the actual output path from handle_one_rank + let rank_link = format!("rank_{}/{}", rank_num, main_output_path.display()); + rank_links.push((rank_num.clone(), rank_link)); + } + + // Sort rank links by rank number + rank_links.sort_by(|a, b| { + let a_num: i32 = + a.0.parse() + .expect(&format!("Failed to parse rank number from '{}'", a.0)); + let b_num: i32 = + b.0.parse() + .expect(&format!("Failed to parse rank number from '{}'", b.0)); + a_num.cmp(&b_num) + }); + + // Core logic complete - no HTML generation yet + // TODO - Add landing page HTML generation using template system + + println!( + "Generated multi-rank report with {} ranks", + rank_links.len() + ); + println!("Individual rank reports available in:"); + for (rank_num, _) in &rank_links { + println!(" - rank_{}/index.html", rank_num); + } + + // No browser opening since no landing page yet + // TODO - Generate landing page and open browser + Ok(()) } diff --git a/src/lib.rs b/src/lib.rs index ff7110e..ff11115 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -37,6 +37,7 @@ pub struct ParseConfig { pub plain_text: bool, pub export: bool, pub inductor_provenance: bool, + pub all_ranks: bool, } impl Default for ParseConfig { @@ -50,6 +51,7 @@ impl Default for ParseConfig { plain_text: false, export: false, inductor_provenance: false, + all_ranks: false, } } }