Support dims annotation and enhance dump loader in dumper (#19276)

python/sglang/srt/debug_utils/dump_loader.py

@@ -1,11 +1,60 @@
 import functools
 import os
+from dataclasses import dataclass
 from pathlib import Path
-from typing import Any, Dict
+from typing import Any, Dict, Tuple
 
 import polars as pl
 import torch
 
+_TYPED_FIELDS: list[tuple[str, type]] = [("rank", int)]
+
+
+def parse_meta_from_filename(path: Path) -> Dict[str, Any]:
+    stem = Path(path).stem
+    result: Dict[str, Any] = {}
+    for kv in stem.split("___"):
+        if "=" in kv:
+            k, v = kv.split("=", 1)
+            result[k] = v
+    for field, converter in _TYPED_FIELDS:
+        if field in result:
+            result[field] = converter(result[field])
+    return result
+
+
+@dataclass
+class ValueWithMeta:
+    value: Any
+    meta: Dict[str, Any]
+
+    @staticmethod
+    def load(path: Path) -> "ValueWithMeta":
+        path = Path(path)
+        meta_from_filename = parse_meta_from_filename(path)
+
+        try:
+            raw = torch.load(path, weights_only=False, map_location="cpu")
+        except Exception as e:
+            print(f"Skip load {path} since error {e}")
+            return ValueWithMeta(
+                value=None, meta={**meta_from_filename, "filename": path.name}
+            )
+
+        value, meta_from_embedded = _unwrap_dict_format(raw)
+        return ValueWithMeta(
+            value=value,
+            meta={**meta_from_filename, **meta_from_embedded, "filename": path.name},
+        )
+
+
+def _unwrap_dict_format(obj: Any) -> Tuple[Any, Dict[str, Any]]:
+    if isinstance(obj, dict) and "value" in obj:
+        meta = obj.get("meta", {})
+        assert isinstance(meta, dict), f"Expected meta to be dict, got {type(meta)}"
+        return obj["value"], meta
+    return obj, {}
+
 
 class DumpLoader:
     def __init__(self):
@@ -50,10 +99,7 @@ def read_meta(directory):
     rows = []
     for p in directory.glob("*.pt"):
         try:
-            full_kwargs = {}
-            for kv in p.stem.split("___"):
-                k, v = kv.split("=")
-                full_kwargs[k] = v
+            full_kwargs = parse_meta_from_filename(p)
             rows.append(
                 {
                     "filename": str(p.name),
@@ -83,26 +129,27 @@ def _add_duplicate_index(df: pl.DataFrame) -> pl.DataFrame:
     return df
 
 
-def find_row(df, conditions: Dict[str, Any]):
-    df_sub = df.filter(
-        functools.reduce(
-            lambda a, b: a & b,
-            [
-                (
-                    pl.col(col)
-                    == _cast_to_polars_dtype(conditions[col], df.schema[col])
-                    if conditions[col] is not None
-                    else pl.col(col).is_null()
-                )
-                for col in conditions.keys()
-                if col in df.columns
-            ],
+def filter_rows(df: pl.DataFrame, conditions: Dict[str, Any]) -> list[dict]:
+    filter_exprs = [
+        (
+            pl.col(col) == _cast_to_polars_dtype(conditions[col], df.schema[col])
+            if conditions[col] is not None
+            else pl.col(col).is_null()
         )
-    )
-    if len(df_sub) > 1:
-        print(f"find_row find ambiguous results: {df_sub=}")
+        for col in conditions
+        if col in df.columns
+    ]
+    if not filter_exprs:
+        return []
+    return df.filter(functools.reduce(lambda a, b: a & b, filter_exprs)).to_dicts()
+
+
+def find_row(df: pl.DataFrame, conditions: Dict[str, Any]):
+    rows = filter_rows(df, conditions)
+    if len(rows) > 1:
+        print(f"find_row find ambiguous results: {rows=}")
         return None
-    return df_sub.to_dicts()[0] if len(df_sub) > 0 else None
+    return rows[0] if rows else None
 
 
 def _cast_to_polars_dtype(value, target_dtype):