Refactor CompileDirectoryEntry, print size

Fixes https://github.com/pytorch/pytorch/issues/134301 Signed-off-by: Edward Z. Yang <[email protected]> ghstack-source-id: 29c22b4215dc648459d78da4bd2991391b112b75 ghstack-comment-id: 2324862517 Pull Request resolved: #64
pytorch · Sep 2, 2024 · 973e34f · 973e34f
1 parent 6c6e3c3
commit 973e34f
Show file tree

Hide file tree

Showing 6 changed files with 51 additions and 20 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/Cargo.toml b/Cargo.toml
@@ -32,3 +32,4 @@ regex = "1.9.2"
 serde = { version = "1.0.185", features = ["serde_derive"] }
 serde_json = "1.0.100"
 tinytemplate = "1.1.0"
+human_bytes = "0.4.3"
diff --git a/src/lib.rs b/src/lib.rs
@@ -3,6 +3,7 @@ use fxhash::{FxHashMap, FxHashSet};
 use md5::{Digest, Md5};
 use std::ffi::{OsStr, OsString};
 
+use human_bytes::human_bytes;
 use indicatif::{MultiProgress, ProgressBar, ProgressStyle};
 use regex::Regex;
 use std::cell::RefCell;
@@ -82,7 +83,7 @@ fn run_parser<'t>(
     payload: &str,
     output_count: &mut i32,
     output: &mut Vec<(PathBuf, String)>,
-    compile_directory: &mut Vec<(String, String, i32)>,
+    compile_directory: &mut Vec<CompileDirectoryEntry>,
     multi: &MultiProgress,
     stats: &mut Stats,
 ) {
@@ -106,27 +107,36 @@ fn run_parser<'t>(
                             } else {
                                 raw_filename
                             };
+                            let size = human_bytes(out.len() as f64);
                             output.push((filename.clone(), out));
                             let filename_str = format!("{}", filename.to_string_lossy());
-                            compile_directory.push((
-                                filename_str.clone(),
-                                filename_str,
-                                *output_count,
-                            ));
+                            compile_directory.push(CompileDirectoryEntry {
+                                filename: filename_str.clone(),
+                                display_filename: filename_str,
+                                seq_nr: *output_count,
+                                size: size,
+                            });
                             *output_count += 1;
                         }
                         ParserOutput::GlobalFile(filename, out) => {
+                            let size = human_bytes(out.len() as f64);
                             output.push((filename.clone(), out));
                             let filename_str = format!("{}", filename.to_string_lossy());
-                            compile_directory.push((
-                                filename_str.clone(),
-                                filename_str,
-                                *output_count,
-                            ));
+                            compile_directory.push(CompileDirectoryEntry {
+                                filename: filename_str.clone(),
+                                display_filename: filename_str,
+                                seq_nr: *output_count,
+                                size: size,
+                            });
                             *output_count += 1;
                         }
                         ParserOutput::Link(name, url) => {
-                            compile_directory.push((url, name, *output_count));
+                            compile_directory.push(CompileDirectoryEntry {
+                                filename: url,
+                                display_filename: name,
+                                seq_nr: *output_count,
+                                size: "".to_string(),
+                            });
                             *output_count += 1;
                         }
                     }
@@ -191,7 +201,7 @@ pub fn parse_path(path: &PathBuf, config: ParseConfig) -> anyhow::Result<ParseOu
     // Each entry is a compile id => (link, rendered name, output number)
     // For files, link and rendered name are the same
     // For links, you can specify a custom name for the link
-    let mut directory: FxIndexMap<Option<CompileId>, Vec<(String, String, i32)>> =
+    let mut directory: FxIndexMap<Option<CompileId>, Vec<CompileDirectoryEntry>> =
         FxIndexMap::default();
 
     let mut metrics_index: CompilationMetricsIndex = FxIndexMap::default();

diff --git a/src/parsers.rs b/src/parsers.rs
@@ -349,7 +349,7 @@ pub struct CompilationMetricsParser<'t> {
     pub tt: &'t TinyTemplate<'t>,
     pub stack_index: &'t RefCell<StackIndex>,
     pub symbolic_shape_specialization_index: &'t RefCell<SymbolicShapeSpecializationIndex>,
-    pub output_files: &'t Vec<(String, String, i32)>,
+    pub output_files: &'t Vec<CompileDirectoryEntry>,
     pub compile_id_dir: &'t PathBuf,
 }
 impl StructuredLogParser for CompilationMetricsParser<'_> {
@@ -416,11 +416,16 @@ impl StructuredLogParser for CompilationMetricsParser<'_> {
                 let new_str: String = parts[1..].join("");
                 new_str
             };
-            let output_files: Vec<(String, String, i32)> = self
+            let output_files: Vec<CompileDirectoryEntry> = self
                 .output_files
                 .iter()
-                .map(|(url, name, number)| {
-                    return (remove_prefix(url), remove_prefix(name), number.clone());
+                .map(|e| {
+                    return CompileDirectoryEntry {
+                        filename: remove_prefix(&e.filename),
+                        display_filename: remove_prefix(&e.display_filename),
+                        seq_nr: e.seq_nr.clone(),
+                        size: e.size.clone(),
+                    };
                 })
                 .collect();
             let context = CompilationMetricsContext {

diff --git a/src/templates.rs b/src/templates.rs
@@ -138,7 +138,7 @@ Build products below:
     <li><a id="{compile_directory.0}">{compile_directory.0}</a>
     <ul>
         {{ for path_idx in compile_directory.1 }}
-            <li><a href="{path_idx.0}">{path_idx.1}</a> ({path_idx.2})</li>
+            <li><a href="{path_idx.filename}">{path_idx.display_filename}</a> ({path_idx.seq_nr}) {path_idx.size}</li>
         {{ endfor }}
     </ul>
     </li>

diff --git a/src/types.rs b/src/types.rs
@@ -351,7 +351,7 @@ pub struct CompilationMetricsContext<'e> {
     pub compile_id: String,
     pub stack_html: String,
     pub symbolic_shape_specializations: Vec<SymbolicShapeSpecializationContext>,
-    pub output_files: &'e Vec<(String, String, i32)>,
+    pub output_files: &'e Vec<CompileDirectoryEntry>,
     pub compile_id_dir: &'e PathBuf,
     pub mini_stack_html: String,
 }
@@ -556,7 +556,7 @@ pub struct DynamoGuardsContext {
 pub struct IndexContext {
     pub css: &'static str,
     pub javascript: &'static str,
-    pub directory: Vec<(String, Vec<(String, String, i32)>)>,
+    pub directory: Vec<(String, Vec<CompileDirectoryEntry>)>,
     pub stack_trie_html: String,
     pub unknown_stack_trie_html: String,
     pub has_unknown_stack_trie: bool,
@@ -573,3 +573,11 @@ pub struct SymbolicShapeSpecializationContext {
     pub user_stack_html: String,
     pub stack_html: String,
 }
+
+#[derive(Debug, Serialize, Clone)]
+pub struct CompileDirectoryEntry {
+    pub filename: String,
+    pub display_filename: String,
+    pub seq_nr: i32,
+    pub size: String,
+}