diff --git a/src/lib.rs b/src/lib.rs index 56ddfe6..63295c3 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -61,7 +61,7 @@ pub fn parse_path(path: &PathBuf, config: ParseConfig) -> anyhow::Result> = None; - let mut directory: FxIndexMap, Vec> = FxIndexMap::default(); + let mut directory: FxIndexMap, Vec<(PathBuf, i32)>> = FxIndexMap::default(); // Store results in an output Vec let mut output: Vec<(PathBuf, String)> = Vec::new(); @@ -72,6 +72,12 @@ pub fn parse_path(path: &PathBuf, config: ParseConfig) -> anyhow::Result anyhow::Result { for (filename, out) in results { output.push((filename.clone(), out)); - compile_directory.push(filename); + compile_directory.push((filename, output_count)); + output_count += 1; } } Err(err) => match parser.name() { diff --git a/src/parsers.rs b/src/parsers.rs index c104396..04ac786 100644 --- a/src/parsers.rs +++ b/src/parsers.rs @@ -282,6 +282,46 @@ impl StructuredLogParser for CompilationMetricsParser<'_> { } } +pub struct AOTAutogradBackwardCompilationMetricsParser<'t> { + tt: &'t TinyTemplate<'t>, +} +impl StructuredLogParser for AOTAutogradBackwardCompilationMetricsParser<'_> { + fn name(&self) -> &'static str { + "aot_autograd_backward_compilation_metrics" + } + fn get_metadata<'e>(&self, e: &'e Envelope) -> Option> { + e.aot_autograd_backward_compilation_metrics + .as_ref() + .map(|m| Metadata::AOTAutogradBackwardCompilationMetrics(m)) + } + fn parse<'e>( + &self, + lineno: usize, + metrics: Metadata<'e>, + _rank: Option, + compile_id: &Option, + _payload: &str, + ) -> anyhow::Result { + let filename = format!("{}.html", self.name()); + if let Metadata::AOTAutogradBackwardCompilationMetrics(m) = metrics { + let id = compile_id + .clone() + .map_or("(unknown) ".to_string(), |c| format!("{cid} ", cid = c)); + let context = AOTAutogradBackwardCompilationMetricsContext { + css: crate::CSS, + m: &m, + compile_id: id, + }; + let output = self.tt.render(&filename, &context)?; + simple_file_output(&filename, lineno, compile_id, &output) + } else { + Err(anyhow::anyhow!( + "Expected AOTAutogradBackwardCompilationMetrics metadata" + )) + } + } +} + // Register your parser here pub fn default_parsers<'t>(tt: &'t TinyTemplate<'t>) -> Vec> { // We need to use Box wrappers here because vecs in Rust need to have known size @@ -310,6 +350,7 @@ pub fn default_parsers<'t>(tt: &'t TinyTemplate<'t>) -> Vec{compile_directory.0}
    - {{ for path in compile_directory.1 }} -
  • {path}
  • + {{ for path_idx in compile_directory.1 }} +
  • {path_idx.0} ({path_idx.1})
  • {{ endfor }}
@@ -216,3 +216,24 @@ pub static TEMPLATE_COMPILATION_METRICS: &str = r#" "#; + +pub static TEMPLATE_AOT_AUTOGRAD_BACKWARD_COMPILATION_METRICS: &str = r#" + + + + AOT Autograd Backward Compilation Metrics + + +

Compilation Info for {compile_id}

+

Failures

+ {{ if m.fail_type }} +

Failure Exception:

{m.fail_type}

+

Failure Reason:

{m.fail_reason}

+ {{ else }} +

No failures!

+ {{ endif }} + + +"#; diff --git a/src/types.rs b/src/types.rs index aa888ef..83b6cd1 100644 --- a/src/types.rs +++ b/src/types.rs @@ -205,6 +205,21 @@ pub struct CompilationMetricsMetadata { pub dynamo_time_before_restart_s: Option, } +#[derive(Debug, Deserialize, Serialize)] +pub struct AOTAutogradBackwardCompilationMetricsMetadata { + pub start_time: Option, + pub elapsed_time: Option, // technically redundant with envelope + pub fail_type: Option, + pub fail_reason: Option, +} + +#[derive(Debug, Serialize)] +pub struct AOTAutogradBackwardCompilationMetricsContext<'e> { + pub m: &'e AOTAutogradBackwardCompilationMetricsMetadata, + pub css: &'static str, + pub compile_id: String, +} + #[derive(Debug, Serialize)] pub struct CompilationMetricsContext<'e> { pub m: &'e CompilationMetricsMetadata, @@ -262,6 +277,7 @@ pub enum Metadata<'e> { InductorOutputCode(&'e InductorOutputCodeMetadata), OptimizeDdpSplitChild(&'e OptimizeDdpSplitChildMetadata), CompilationMetrics(&'e CompilationMetricsMetadata), + AOTAutogradBackwardCompilationMetrics(&'e AOTAutogradBackwardCompilationMetricsMetadata), } #[derive(Debug, Deserialize)] @@ -286,6 +302,8 @@ pub struct Envelope { pub inductor_post_grad_graph: Option, pub inductor_output_code: Option, pub compilation_metrics: Option, + pub aot_autograd_backward_compilation_metrics: + Option, pub graph_dump: Option, } @@ -304,7 +322,7 @@ pub struct DynamoGuardsContext { #[derive(Debug, Serialize)] pub struct IndexContext { pub css: &'static str, - pub directory: Vec<(String, Vec)>, + pub directory: Vec<(String, Vec<(PathBuf, i32)>)>, pub stack_trie_html: String, pub unknown_stack_trie_html: String, pub has_unknown_stack_trie: bool,