Mark everything in CompileId as optional, add compiled_autograd_id

ezyang · Dec 9, 2024 · a6a444d · a6a444d
1 parent 43c718e
commit a6a444d
Show file tree

Hide file tree

Showing 8 changed files with 6,470 additions and 2,500 deletions.
diff --git a/src/lib.rs b/src/lib.rs
@@ -381,10 +381,22 @@ pub fn parse_path(path: &PathBuf, config: ParseConfig) -> anyhow::Result<ParseOu
                 .map_or(
                     format!("unknown_{lineno}"),
                     |CompileId {
+                         compiled_autograd_id,
                          frame_id,
                          frame_compile_id,
                          attempt,
-                     }| { format!("{frame_id}_{frame_compile_id}_{attempt}") },
+                     }| {
+                        let frame_id_str = frame_id.map_or("-".to_string(), |v| v.to_string());
+                        let frame_compile_id_str =
+                            frame_compile_id.map_or("-".to_string(), |v| v.to_string());
+                        let attempt_str = attempt.map_or("-".to_string(), |v| v.to_string());
+
+                        if let Some(ca_id) = compiled_autograd_id {
+                            format!("{ca_id}_{frame_id_str}_{frame_compile_id_str}_{attempt_str}")
+                        } else {
+                            format!("{frame_id_str}_{frame_compile_id_str}_{attempt_str}")
+                        }
+                    },
                 )
                 .into();
             let parser: Box<dyn StructuredLogParser> =
@@ -450,7 +462,10 @@ 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;
+                if let Some(_frame_id) = c.frame_compile_id {
+                    // data migration for old logs that don't have attempt
+                    c.attempt = Some(0);
+                }
             }
             metrics_index.entry(cid).or_default().push(m.clone());
         }

diff --git a/src/parsers.rs b/src/parsers.rs
@@ -59,10 +59,22 @@ fn simple_file_output(
         .map_or(
             format!("unknown_{lineno}"),
             |CompileId {
+                 compiled_autograd_id,
                  frame_id,
                  frame_compile_id,
                  attempt,
-             }| { format!("{frame_id}_{frame_compile_id}_{attempt}") },
+             }| {
+                let frame_id_str = frame_id.map_or("-".to_string(), |v| v.to_string());
+                let frame_compile_id_str =
+                    frame_compile_id.map_or("-".to_string(), |v| v.to_string());
+                let attempt_str = attempt.map_or("-".to_string(), |v| v.to_string());
+
+                if let Some(ca_id) = compiled_autograd_id {
+                    format!("{ca_id}_{frame_id_str}_{frame_compile_id_str}_{attempt_str}")
+                } else {
+                    format!("{frame_id_str}_{frame_compile_id_str}_{attempt_str}")
+                }
+            },
         )
         .into();
     let subdir = PathBuf::from(compile_id_dir);
@@ -380,7 +392,10 @@ 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;
+                if let Some(_frame_id) = c.frame_compile_id {
+                    // data migration for old logs that don't have attempt
+                    c.attempt = Some(0);
+                }
             }
             let stack_html = self
                 .stack_index

diff --git a/src/types.rs b/src/types.rs
@@ -128,16 +128,27 @@ impl StackTrieNode {
 
 #[derive(Eq, PartialEq, Hash, Deserialize, Serialize, Debug, Clone)]
 pub struct CompileId {
-    pub frame_id: u32,
-    pub frame_compile_id: u32,
-    pub attempt: u32,
+    pub compiled_autograd_id: Option<u32>,
+    pub frame_id: Option<u32>,
+    pub frame_compile_id: Option<u32>,
+    pub attempt: Option<u32>,
 }
 
 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)?;
+        write!(f, "[")?;
+        if let Some(compiled_autograd_id) = self.compiled_autograd_id {
+            write!(f, "{}/", compiled_autograd_id)?;
+        }
+        let frame_id = self.frame_id.map_or("-".to_string(), |v| v.to_string());
+        let frame_compile_id = self
+            .frame_compile_id
+            .map_or("-".to_string(), |v| v.to_string());
+        write!(f, "{}/{}", frame_id, frame_compile_id)?;
+        if let Some(attempt) = self.attempt {
+            if attempt != 0 {
+                write!(f, "_{}", attempt)?;
+            }
         }
         write!(f, "]")
     }