Add parsing for aot_autograd_backward_compilation_metrics

src/lib.rs

@@ -61,7 +61,7 @@ pub fn parse_path(path: &PathBuf, config: ParseConfig) -> anyhow::Result<ParseOu
 
     let mut expected_rank: Option<Option<u32>> = None;
 
-    let mut directory: FxIndexMap<Option<CompileId>, Vec<PathBuf>> = FxIndexMap::default();
+    let mut directory: FxIndexMap<Option<CompileId>, Vec<(PathBuf, i32)>> = FxIndexMap::default();
 
     // Store results in an output Vec<PathBuf, String>
     let mut output: Vec<(PathBuf, String)> = Vec::new();
@@ -72,6 +72,12 @@ pub fn parse_path(path: &PathBuf, config: ParseConfig) -> anyhow::Result<ParseOu
     tt.add_template("failures_and_restarts.html", TEMPLATE_FAILURES_AND_RESTARTS)?;
     tt.add_template("dynamo_guards.html", TEMPLATE_DYNAMO_GUARDS)?;
     tt.add_template("compilation_metrics.html", TEMPLATE_COMPILATION_METRICS)?;
+    tt.add_template(
+        "aot_autograd_backward_compilation_metrics.html",
+        TEMPLATE_AOT_AUTOGRAD_BACKWARD_COMPILATION_METRICS,
+    )?;
+
+    let mut output_count = 0;
 
     let mut breaks = RestartsAndFailuresContext {
         css: TEMPLATE_FAILURES_CSS,
@@ -187,7 +193,8 @@ pub fn parse_path(path: &PathBuf, config: ParseConfig) -> anyhow::Result<ParseOu
                     Ok(results) => {
                         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() {

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<Metadata<'e>> {
+        e.aot_autograd_backward_compilation_metrics
+            .as_ref()
+            .map(|m| Metadata::AOTAutogradBackwardCompilationMetrics(m))
+    }
+    fn parse<'e>(
+        &self,
+        lineno: usize,
+        metrics: Metadata<'e>,
+        _rank: Option<u32>,
+        compile_id: &Option<CompileId>,
+        _payload: &str,
+    ) -> anyhow::Result<ParseOutput> {
+        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<Box<dyn StructuredLogParser + 't>> {
     // 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<Box<dyn StructuredLo
         Box::new(InductorOutputCodeParser),
         Box::new(OptimizeDdpSplitChildParser),
         Box::new(CompilationMetricsParser { tt }), // TODO: use own tt instances
+        Box::new(AOTAutogradBackwardCompilationMetricsParser { tt }), // TODO: use own tt instances
     ];
 
     result

src/templates.rs

@@ -85,8 +85,8 @@ Build products below:
 {{ for compile_directory in directory }}
     <li><a id="{compile_directory.0}">{compile_directory.0}</a>
     <ul>
-        {{ for path in compile_directory.1 }}
-            <li><a href="{path}">{path}</a></li>
+        {{ for path_idx in compile_directory.1 }}
+            <li><a href="{path_idx.0}">{path_idx.0}</a> ({path_idx.1})</li>
         {{ endfor }}
     </ul>
     </li>
@@ -216,3 +216,24 @@ pub static TEMPLATE_COMPILATION_METRICS: &str = r#"
 </body>
 </html>
 "#;
+
+pub static TEMPLATE_AOT_AUTOGRAD_BACKWARD_COMPILATION_METRICS: &str = r#"
+<html>
+<head>
+    <style>
+    {css}
+    </style>
+    <title>AOT Autograd Backward Compilation Metrics</title>
+</head>
+<body>
+    <h1>Compilation Info for {compile_id}</h1>
+    <h2>Failures</h2>
+    {{ if m.fail_type }}
+    <p>Failure Exception: <pre>{m.fail_type}</pre></p>
+    <p>Failure Reason: <pre>{m.fail_reason}</pre></p>
+    {{ else }}
+    <p> No failures! </p>
+    {{ endif }}
+</body>
+</html>
+"#;

src/types.rs

@@ -205,6 +205,21 @@ pub struct CompilationMetricsMetadata {
     pub dynamo_time_before_restart_s: Option<f64>,
 }
 
+#[derive(Debug, Deserialize, Serialize)]
+pub struct AOTAutogradBackwardCompilationMetricsMetadata {
+    pub start_time: Option<f64>,
+    pub elapsed_time: Option<f64>, // technically redundant with envelope
+    pub fail_type: Option<String>,
+    pub fail_reason: Option<String>,
+}
+
+#[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<EmptyMetadata>,
     pub inductor_output_code: Option<InductorOutputCodeMetadata>,
     pub compilation_metrics: Option<CompilationMetricsMetadata>,
+    pub aot_autograd_backward_compilation_metrics:
+        Option<AOTAutogradBackwardCompilationMetricsMetadata>,
     pub graph_dump: Option<GraphDumpMetadata>,
 }
 
@@ -304,7 +322,7 @@ pub struct DynamoGuardsContext {
 #[derive(Debug, Serialize)]
 pub struct IndexContext {
     pub css: &'static str,
-    pub directory: Vec<(String, Vec<PathBuf>)>,
+    pub directory: Vec<(String, Vec<(PathBuf, i32)>)>,
     pub stack_trie_html: String,
     pub unknown_stack_trie_html: String,
     pub has_unknown_stack_trie: bool,