Tiny add filter, support duplications, add visualizations, fix error and robustness for dump comparator (#16262)

python/sglang/srt/debug_utils/dump_loader.py

@@ -47,16 +47,19 @@ def read_meta(directory):
 
     rows = []
     for p in directory.glob("*.pt"):
-        full_kwargs = {}
-        for kv in p.stem.split("___"):
-            k, v = kv.split("=")
-            full_kwargs[k] = v
-        rows.append(
-            {
-                "filename": str(p.name),
-                **full_kwargs,
-            }
-        )
+        try:
+            full_kwargs = {}
+            for kv in p.stem.split("___"):
+                k, v = kv.split("=")
+                full_kwargs[k] = v
+            rows.append(
+                {
+                    "filename": str(p.name),
+                    **full_kwargs,
+                }
+            )
+        except Exception as e:
+            print(f"[DumpLoader] skip loading {p} due to error {e}")
 
     df = pl.DataFrame(rows)
     df = df.with_columns(
@@ -64,6 +67,17 @@ def read_meta(directory):
         pl.col("rank").cast(int),
         pl.col("dump_index").cast(int),
     )
+    df = _add_duplicate_index(df)
+    df = df.sort("rank", "dump_index")
+    return df
+
+
+def _add_duplicate_index(df: pl.DataFrame) -> pl.DataFrame:
+    group_cols = [c for c in df.columns if c not in ["filename", "dump_index"]]
+    df = df.sort(group_cols + ["dump_index"])
+    df = df.with_columns(
+        pl.cum_count("dump_index").over(group_cols).sub(1).alias("duplicate_index")
+    )
     return df