From ab7353a04afa9580be55cd0124a00513eddf245e Mon Sep 17 00:00:00 2001
From: "Edward Z. Yang" <ezyang@meta.com>
Date: Mon, 15 Apr 2024 09:15:14 -0400
Subject: [PATCH] Add support for "graph_dump" log (#17)

---
 src/parsers.rs | 33 +++++++++++++++++++++++++++++++++
 src/types.rs   |  7 +++++++
 2 files changed, 40 insertions(+)

diff --git a/src/parsers.rs b/src/parsers.rs
index 47f8135..c104396 100644
--- a/src/parsers.rs
+++ b/src/parsers.rs
@@ -95,6 +95,38 @@ impl StructuredLogParser for SentinelFileParser {
     }
 }
 
+/**
+ * Generic parser for graph_dump entries
+ */
+pub struct GraphDumpParser;
+impl StructuredLogParser for GraphDumpParser {
+    fn name(&self) -> &'static str {
+        "graph_dump" // ToDO: more specific?
+    }
+    fn get_metadata<'e>(&self, e: &'e Envelope) -> Option<Metadata<'e>> {
+        e.graph_dump.as_ref().map(|m| Metadata::GraphDump(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::GraphDump(metadata) = metadata {
+            let filename: PathBuf = {
+                let mut r = OsString::from(&metadata.name);
+                r.push(OsStr::new(".txt"));
+                r.into()
+            };
+            simple_file_output(&filename.to_string_lossy(), lineno, compile_id, payload)
+        } else {
+            Err(anyhow::anyhow!("Expected GraphDump metadata"))
+        }
+    }
+}
+
 // Same as SentinelFileParser, but can log the size of the graph
 pub struct DynamoOutputGraphParser;
 impl StructuredLogParser for DynamoOutputGraphParser {
@@ -272,6 +304,7 @@ pub fn default_parsers<'t>(tt: &'t TinyTemplate<'t>) -> Vec<Box<dyn StructuredLo
         Box::new(SentinelFileParser::new("inductor_post_grad_graph", |e| {
             e.inductor_post_grad_graph.as_ref()
         })),
+        Box::new(GraphDumpParser),
         Box::new(DynamoOutputGraphParser),
         Box::new(DynamoGuardParser { tt }),
         Box::new(InductorOutputCodeParser),
diff --git a/src/types.rs b/src/types.rs
index 12aa476..c720cc3 100644
--- a/src/types.rs
+++ b/src/types.rs
@@ -156,6 +156,11 @@ pub struct OptimizeDdpSplitChildMetadata {
 #[derive(Debug, Deserialize)]
 pub struct EmptyMetadata {}
 
+#[derive(Debug, Deserialize)]
+pub struct GraphDumpMetadata {
+    pub name: String,
+}
+
 #[derive(Debug, Deserialize)]
 pub struct DynamoOutputGraphMetadata {
     _sizes: Option<FxHashMap<String, Vec<SymInt>>>,
@@ -246,6 +251,7 @@ pub struct RestartsAndFailuresContext {
 #[derive(Debug)]
 pub enum Metadata<'e> {
     Empty(&'e EmptyMetadata),
+    GraphDump(&'e GraphDumpMetadata),
     DynamoOutputGraph(&'e DynamoOutputGraphMetadata),
     #[allow(dead_code)]
     DynamoStart(&'e DynamoStartMetadata),
@@ -275,6 +281,7 @@ pub struct Envelope {
     pub inductor_post_grad_graph: Option<EmptyMetadata>,
     pub inductor_output_code: Option<InductorOutputCodeMetadata>,
     pub compilation_metrics: Option<CompilationMetricsMetadata>,
+    pub graph_dump: Option<GraphDumpMetadata>,
 }
 
 #[derive(Debug, Deserialize, Serialize)]