Introduce CompileId variants and CompiledAutogradInitiated CompileId …

…variant
ezyang · Dec 6, 2024 · 8ff043f · 8ff043f
1 parent d9962b1
commit 8ff043f
Show file tree

Hide file tree

Showing 9 changed files with 7,096 additions and 2,568 deletions.
diff --git a/src/lib.rs b/src/lib.rs
@@ -380,11 +380,19 @@ pub fn parse_path(path: &PathBuf, config: ParseConfig) -> anyhow::Result<ParseOu
                 .as_ref()
                 .map_or(
                     format!("unknown_{lineno}"),
-                    |CompileId {
-                         frame_id,
-                         frame_compile_id,
-                         attempt,
-                     }| { format!("{frame_id}_{frame_compile_id}_{attempt}") },
+                    |c | match c {
+                        CompileId::UserInitiated(d) => {
+                            format!("-_{}_{}_{}", d.frame_id, d.frame_compile_id, d.attempt)
+                        }
+                        CompileId::CompiledAutogradInitiated { compiled_autograd_id, dynamo_id } => {
+
+                            if let Some(d) = dynamo_id {
+                                format!("{}_{}_{}_{}", compiled_autograd_id, d.frame_id, d.frame_compile_id, d.attempt)
+                            } else {
+                                format!("{}_-_-_-", compiled_autograd_id)
+                            }
+                        },
+                    }
                 )
                 .into();
             let parser: Box<dyn StructuredLogParser> =
@@ -450,7 +458,15 @@ pub fn parse_path(path: &PathBuf, config: ParseConfig) -> anyhow::Result<ParseOu
             }
             let mut cid = e.compile_id.clone();
             if let Some(c) = cid.as_mut() {
-                c.attempt = 0;
+                match c {
+                    CompileId::UserInitiated(d) => {
+                        // Is this for data migration?
+                        d.attempt = 0;
+                    },
+                    CompileId::CompiledAutogradInitiated { compiled_autograd_id: _, dynamo_id: _ } => {
+                        // DynamoId should already have attempt set
+                    },
+                }
             }
             metrics_index.entry(cid).or_default().push(m.clone());
         }

diff --git a/src/parsers.rs b/src/parsers.rs
@@ -58,11 +58,18 @@ fn simple_file_output(
         .as_ref()
         .map_or(
             format!("unknown_{lineno}"),
-            |CompileId {
-                 frame_id,
-                 frame_compile_id,
-                 attempt,
-             }| { format!("{frame_id}_{frame_compile_id}_{attempt}") },
+            |c | match c {
+                CompileId::UserInitiated(d) => {
+                    format!("-_{}_{}_{}", d.frame_id, d.frame_compile_id, d.attempt)
+                }
+                CompileId::CompiledAutogradInitiated { compiled_autograd_id, dynamo_id } => {
+                    if let Some(d) = dynamo_id {
+                        format!("{}_{}_{}_{}", compiled_autograd_id, d.frame_id, d.frame_compile_id, d.attempt)
+                    } else {
+                        format!("{}_-_-_-", compiled_autograd_id)
+                    }
+                }
+            }
         )
         .into();
     let subdir = PathBuf::from(compile_id_dir);
@@ -380,7 +387,15 @@ impl StructuredLogParser for CompilationMetricsParser<'_> {
                 .map_or("(unknown) ".to_string(), |c| format!("{cid} ", cid = c));
             let mut cid = compile_id.clone();
             if let Some(c) = cid.as_mut() {
-                c.attempt = 0;
+                match c {
+                    CompileId::UserInitiated(d) => {
+                        // Is this for data migration?
+                        d.attempt = 0;
+                    }
+                    CompileId::CompiledAutogradInitiated { compiled_autograd_id: _, dynamo_id: _} => {
+                        // DynamoId should already have attempt set
+                    }
+                }
             }
             let stack_html = self
                 .stack_index

diff --git a/src/types.rs b/src/types.rs
@@ -127,19 +127,54 @@ impl StackTrieNode {
 }
 
 #[derive(Eq, PartialEq, Hash, Deserialize, Serialize, Debug, Clone)]
-pub struct CompileId {
+pub struct DynamoId {
     pub frame_id: u32,
     pub frame_compile_id: u32,
     pub attempt: u32,
 }
 
+#[derive(Eq, PartialEq, Hash, Deserialize, Serialize, Debug, Clone)]
+#[serde(untagged)]
+pub enum CompileId {
+    // NOTE: serde(untagged) will match in the order
+    // this enum is defined 
+
+    // When compiled autograd calls torch.compile
+    CompiledAutogradInitiated {
+        compiled_autograd_id: u32,
+        #[serde(flatten)]
+        dynamo_id: Option<DynamoId>,
+    },
+    // When user calls torch.compile
+    UserInitiated(DynamoId)
+}
+
+fn _format_dynamo_id(d: &DynamoId) -> String {
+    if d.attempt != 0 {
+        format!("{}/{}_{}", d.frame_id, d.frame_compile_id, d.attempt)
+    } else {
+        format!("{}/{}", d.frame_id, d.frame_compile_id)
+    }
+}
+
+fn format_dynamo_id(dynamo_id: &Option<DynamoId>) -> String {
+    if let Some(d) = dynamo_id {
+        _format_dynamo_id(d)
+    } else {
+        format!("-/-")
+    }
+}
+
 impl fmt::Display for CompileId {
     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
-        write!(f, "[{}/{}", self.frame_id, self.frame_compile_id)?;
-        if self.attempt != 0 {
-            write!(f, "_{}", self.attempt)?;
+        match self {
+            CompileId::UserInitiated(d) => {
+                write!(f, "[-/{}]", _format_dynamo_id(d))
+            },
+            CompileId::CompiledAutogradInitiated { compiled_autograd_id, dynamo_id } => {
+                write!(f, "[{}/{}]", compiled_autograd_id, format_dynamo_id(dynamo_id))
+            },
         }
-        write!(f, "]")
     }
 }