Update

[ghstack-poisoned]
ezyang · Jul 11, 2024 · 88c3dca · 88c3dca
1 parent 7887ade
commit 88c3dca
Show file tree

Hide file tree

Showing 4 changed files with 88 additions and 0 deletions.
diff --git a/src/lib.rs b/src/lib.rs
@@ -188,6 +188,10 @@ 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(
+        "bwd_compilation_metrics.html",
+        TEMPLATE_BWD_COMPILATION_METRICS,
+    )?;
     tt.add_template(
         "aot_autograd_backward_compilation_metrics.html",
         TEMPLATE_AOT_AUTOGRAD_BACKWARD_COMPILATION_METRICS,

diff --git a/src/parsers.rs b/src/parsers.rs
@@ -436,6 +436,44 @@ impl StructuredLogParser for AOTAutogradBackwardCompilationMetricsParser<'_> {
     }
 }
 
+pub struct BwdCompilationMetricsParser<'t> {
+    tt: &'t TinyTemplate<'t>,
+}
+impl StructuredLogParser for BwdCompilationMetricsParser<'_> {
+    fn name(&self) -> &'static str {
+        "bwd_compilation_metrics"
+    }
+    fn get_metadata<'e>(&self, e: &'e Envelope) -> Option<Metadata<'e>> {
+        e.bwd_compilation_metrics
+            .as_ref()
+            .map(|m| Metadata::BwdCompilationMetrics(m))
+    }
+    fn parse<'e>(
+        &self,
+        lineno: usize,
+        metrics: Metadata<'e>,
+        _rank: Option<u32>,
+        compile_id: &Option<CompileId>,
+        _payload: &str,
+    ) -> anyhow::Result<ParserResults> {
+        let filename = format!("{}.html", self.name());
+        if let Metadata::BwdCompilationMetrics(m) = metrics {
+            let id = compile_id
+                .clone()
+                .map_or("(unknown) ".to_string(), |c| format!("{cid} ", cid = c));
+            let context = BwdCompilationMetricsContext {
+                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 BwdCompilationMetrics metadata"))
+        }
+    }
+}
+
 pub struct ArtifactParser;
 impl StructuredLogParser for ArtifactParser {
     fn name(&self) -> &'static str {
@@ -506,6 +544,7 @@ pub fn default_parsers<'t>(tt: &'t TinyTemplate<'t>) -> Vec<Box<dyn StructuredLo
         Box::new(InductorOutputCodeParser),
         Box::new(OptimizeDdpSplitChildParser),
         Box::new(AOTAutogradBackwardCompilationMetricsParser { tt }), // TODO: use own tt instances
+        Box::new(BwdCompilationMetricsParser { tt }),                 // TODO: use own tt instances
         Box::new(LinkParser),
         Box::new(ArtifactParser),
     ];

diff --git a/src/templates.rs b/src/templates.rs
@@ -302,3 +302,31 @@ pub static TEMPLATE_AOT_AUTOGRAD_BACKWARD_COMPILATION_METRICS: &str = r#"
 </body>
 </html>
 "#;
+
+pub static TEMPLATE_BWD_COMPILATION_METRICS: &str = r#"
+<html>
+<head>
+    <style>
+    {css}
+    </style>
+    <title>Backward Compilation Metrics</title>
+</head>
+<body>
+    <h1>Backward Compilation Info for {compile_id}</h1>
+    <h2>Compile Time(seconds)</h2>
+    {{ if m.inductor_compile_time_s }}
+    <p>Inductor <abbr title="Total time spent running inductor">[?]</abbr>: {m.inductor_compile_time_s}</div>
+    {{ endif }}
+    {{ if m.code_gen_time_s }}
+    <p>Code Gen Time: {m.code_gen_time_s}</p>
+    {{ endif}}
+    <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>
+"#;
diff --git a/src/types.rs b/src/types.rs
@@ -277,6 +277,14 @@ pub struct CompilationMetricsMetadata {
     pub dynamo_time_before_restart_s: Option<f64>,
 }
 
+#[derive(Debug, Deserialize, Serialize, Clone)]
+pub struct BwdCompilationMetricsMetadata {
+    pub inductor_compile_time_s: Option<f64>,
+    pub code_gen_time_s: Option<f64>,
+    pub fail_type: Option<String>,
+    pub fail_reason: Option<String>,
+}
+
 #[derive(Debug, Deserialize, Serialize)]
 pub struct AOTAutogradBackwardCompilationMetricsMetadata {
     pub start_time: Option<f64>,
@@ -295,6 +303,13 @@ pub struct SymbolicShapeSpecializationMetadata {
     pub user_stack: Option<StackSummary>,
 }
 
+#[derive(Debug, Serialize)]
+pub struct BwdCompilationMetricsContext<'e> {
+    pub m: &'e BwdCompilationMetricsMetadata,
+    pub css: &'static str,
+    pub compile_id: String,
+}
+
 #[derive(Debug, Serialize)]
 pub struct AOTAutogradBackwardCompilationMetricsContext<'e> {
     pub m: &'e AOTAutogradBackwardCompilationMetricsMetadata,
@@ -364,6 +379,7 @@ pub enum Metadata<'e> {
     OptimizeDdpSplitChild(&'e OptimizeDdpSplitChildMetadata),
     CompilationMetrics(&'e CompilationMetricsMetadata),
     AOTAutogradBackwardCompilationMetrics(&'e AOTAutogradBackwardCompilationMetricsMetadata),
+    BwdCompilationMetrics(&'e BwdCompilationMetricsMetadata),
     Artifact(&'e ArtifactMetadata),
 }
 
@@ -390,6 +406,7 @@ pub struct Envelope {
     pub dynamo_cpp_guards_str: Option<EmptyMetadata>,
     pub inductor_output_code: Option<InductorOutputCodeMetadata>,
     pub compilation_metrics: Option<CompilationMetricsMetadata>,
+    pub bwd_compilation_metrics: Option<BwdCompilationMetricsMetadata>,
     pub aot_autograd_backward_compilation_metrics:
         Option<AOTAutogradBackwardCompilationMetricsMetadata>,
     pub graph_dump: Option<GraphDumpMetadata>,