Rewrite some more parsers in new type

src/lib.rs

@@ -1,12 +1,10 @@
 use anyhow::anyhow;
 use fxhash::FxHashMap;
 use md5::{Digest, Md5};
-use std::ffi::{OsStr, OsString};
 
 use regex::Regex;
 use std::fs::File;
 use std::io::{self, BufRead};
-use std::path::Path;
 use std::path::PathBuf;
 use tinytemplate::TinyTemplate;
 use indicatif::{MultiProgress, ProgressBar, ProgressStyle};
@@ -85,7 +83,7 @@ pub fn parse_path(path: &PathBuf, config: ParseConfig) -> anyhow::Result<ParseOu
         })
         .peekable();
 
-    let all_parsers = all_parsers();
+    let all_parsers = all_parsers(&tt);
 
     while let Some((lineno, line)) = iter.next() {
         bytes_read += line.len() as u64;
@@ -141,21 +139,6 @@ pub fn parse_path(path: &PathBuf, config: ParseConfig) -> anyhow::Result<ParseOu
             }
         };
 
-        let compile_id_dir: PathBuf = e
-            .compile_id
-            .as_ref()
-            .map_or(
-                format!("unknown_{lineno}"),
-                |CompileId {
-                     frame_id,
-                     frame_compile_id,
-                     attempt,
-                 }| { format!("{frame_id}_{frame_compile_id}_{attempt}") },
-            )
-            .into();
-
-        let subdir = &compile_id_dir;
-
         let mut payload = String::new();
         if let Some(ref expect) = e.has_payload {
             let mut first = true;
@@ -199,8 +182,9 @@ pub fn parse_path(path: &PathBuf, config: ParseConfig) -> anyhow::Result<ParseOu
             }
         }
 
-
-        // TODO: implement these as StructuredLogParseres
+        // TODO: implement this as StructuredLogParser
+        // This one is hard because it consumes the stack, and
+        // I don't want to clone the stack when passing it as reference
         if let Some(m) = e.dynamo_start {
             if let Some(stack) = m.stack {
                 stack_trie.insert(
@@ -212,47 +196,6 @@ pub fn parse_path(path: &PathBuf, config: ParseConfig) -> anyhow::Result<ParseOu
             };
         };
 
-        if e.dynamo_guards.is_some() {
-            let filename = "dynamo_guards.html";
-            let f = subdir.join(filename);
-            match serde_json::from_str::<Vec<DynamoGuard>>(payload.as_str()) {
-                Ok(guards) => {
-                    let guards_context = DynamoGuardsContext { guards };
-                    output.push((f, tt.render("dynamo_guards.html", &guards_context)?));
-                    compile_directory.push(compile_id_dir.join(filename));
-                }
-                Err(err) => {
-                    eprintln!("Failed to parse guards json: {}", err);
-                    stats.fail_dynamo_guards_json += 1;
-                }
-            }
-        }
-
-        if let Some(metadata) = e.inductor_output_code {
-            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);
-            output.push((f, payload.clone()));
-            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);
-            output.push((f, payload.clone()));
-            compile_directory.push(compile_id_dir.join(filename));
-        }
     }
     pb.finish_with_message("done");
     spinner.finish();

src/parsers.rs

@@ -1,14 +1,19 @@
 use crate::types::*;
 use std::path::PathBuf;
+use std::ffi::{OsStr, OsString};
+use std::path::Path;
+use tinytemplate::TinyTemplate;
 
 /**
  * StructuredLogParser
  * Parses a structured log and returns a vec of file outputs.
  * Implement this trait to add your own analyses.
+ *
+ * 'e is the lifetime of the individual envelope
  */
 pub trait StructuredLogParser {
     // If this returns Some value, the parser will be run on that metadata.
-    // Otherwise, it will be skipped. 'e is the lifetime of the envelope.
+    // Otherwise, it will be skipped.
     fn get_metadata<'e>(&self, e: &'e Envelope) -> Option<Metadata<'e>>;
     fn parse<'e>(&self,
         lineno: usize, // Line number from log
@@ -85,8 +90,86 @@ impl StructuredLogParser for DynamoOutputGraphParser {
     }
 }
 
+pub struct DynamoGuardParser<'t> {
+    tt: &'t TinyTemplate<'t>,
+}
+impl StructuredLogParser for DynamoGuardParser<'_> {
+    fn get_metadata<'e>(&self, e: &'e Envelope) -> Option<Metadata<'e>> {
+        e.dynamo_guards.as_ref().map(|m| Metadata::Empty(m))
+    }
+    fn parse<'e>(&self,
+            lineno: usize,
+            _metadata: Metadata<'e>,
+            _rank: Option<u32>,
+            compile_id: &Option<CompileId>,
+            payload: &str
+    ) -> anyhow::Result<ParseOutput> {
+        let filename = "dynamo_guards.html";
+        let guards = serde_json::from_str::<Vec<DynamoGuard>>(payload)?;
+        let guards_context = DynamoGuardsContext { guards };
+        let output = self.tt.render(filename, &guards_context)?;
+        simple_file_output(filename, lineno, compile_id, &output)
+    }
+}
+
+pub struct InductorOutputCodeParser;
+impl StructuredLogParser for InductorOutputCodeParser {
+    fn get_metadata<'e>(&self, e: &'e Envelope) -> Option<Metadata<'e>> {
+        e.inductor_output_code.as_ref().map(|m| Metadata::InductorOutputCode(m))
+    }
+
+    fn parse<'e>(&self,
+        lineno: usize,
+        metadata: Metadata<'e>,
+        _rank: Option<u32>,
+        compile_id: &Option<CompileId>,
+        payload: &str
+    ) -> anyhow::Result<ParseOutput> {
+        if let Metadata::InductorOutputCode(metadata) = metadata {
+            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()
+                    },
+            );
+            simple_file_output(&filename.to_string_lossy(), lineno, compile_id, payload)
+        } else {
+            Err(anyhow::anyhow!("Expected InductorOutputCode metadata"))
+        }
+    }
+}
+
+pub struct OptimizeDdpSplitChildParser;
+impl StructuredLogParser for OptimizeDdpSplitChildParser {
+    fn get_metadata<'e>(&self, e: &'e Envelope) -> Option<Metadata<'e>> {
+        e.optimize_ddp_split_child.as_ref().map(|m| Metadata::OptimizeDdpSplitChild(m))
+    }
+
+    fn parse<'e>(&self,
+        lineno: usize,
+        metadata: Metadata<'e>,
+        _rank: Option<u32>,
+        compile_id: &Option<CompileId>,
+        payload: &str
+    ) -> anyhow::Result<ParseOutput> {
+        if let Metadata::OptimizeDdpSplitChild(m) = metadata {
+            let filename = format!("optimize_ddp_split_child_{}.txt", m.name);
+            simple_file_output(&filename, lineno, compile_id, payload)
+        } else {
+            Err(anyhow::anyhow!("Expected OptimizeDdpSplitChild metadata"))
+        }
+    }
+}
+
 // Register your parser here
-pub fn all_parsers() -> Vec<Box<dyn StructuredLogParser>> {
+pub fn all_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
     let result : Vec<Box<dyn StructuredLogParser>> = vec![
         Box::new(SentinelFileParser::new("optimize_ddp_split_graph.txt", |e| e.optimize_ddp_split_graph.as_ref())),
@@ -95,7 +178,10 @@ pub fn all_parsers() -> Vec<Box<dyn StructuredLogParser>> {
         Box::new(SentinelFileParser::new("aot_backward_graph.txt", |e| e.aot_backward_graph.as_ref())),
         Box::new(SentinelFileParser::new("aot_joint_graph.txt", |e| e.aot_joint_graph.as_ref())),
         Box::new(SentinelFileParser::new("inductor_post_grad_graph.txt", |e| e.inductor_post_grad_graph.as_ref())),
-        Box::new(DynamoOutputGraphParser {}),
+        Box::new(DynamoOutputGraphParser),
+        Box::new(DynamoGuardParser { tt }),
+        Box::new(InductorOutputCodeParser),
+        Box::new(OptimizeDdpSplitChildParser),
     ];
     result
 }

src/types.rs

@@ -141,18 +141,18 @@ impl fmt::Display for FrameSummary {
 
 pub type StackSummary = Vec<FrameSummary>;
 
-#[derive(Debug, Deserialize)]
-pub struct OptimizeDdpSplitChildMetadata {
-    pub name: String,
-}
-
 #[derive(Debug, Deserialize)]
 #[serde(untagged)]
 pub enum SymInt {
     Int(i64),
     Symbol(String),
 }
 
+#[derive(Debug, Deserialize)]
+pub struct OptimizeDdpSplitChildMetadata {
+    pub name: String,
+}
+
 #[derive(Debug, Deserialize)]
 pub struct EmptyMetadata {}
 
@@ -176,8 +176,8 @@ pub enum Metadata<'e> {
     DynamoOutputGraph(&'e DynamoOutputGraphMetadata),
     #[allow(dead_code)]
     DynamoStart(&'e DynamoStartMetadata),
-    #[allow(dead_code)]
     InductorOutputCode(&'e InductorOutputCodeMetadata),
+    OptimizeDdpSplitChild(&'e OptimizeDdpSplitChildMetadata),
 }
 
 #[derive(Debug, Deserialize)]