Serialize vLLM attention profile results

This commit is contained in:
2026-07-16 21:17:27 +08:00
parent 2f04971c90
commit b3e437acb0

View File

@@ -46,6 +46,15 @@ def git_head(repo: Path) -> str:
).strip() ).strip()
def json_default(value: object) -> object:
if isinstance(value, (Path, torch.dtype, torch.device)):
return str(value)
item = getattr(value, "item", None)
if callable(item):
return item()
raise TypeError(f"cannot serialize {type(value).__name__}")
def main() -> None: def main() -> None:
args = parse_args() args = parse_args()
if vllm.__version__ != VLLM_VERSION: if vllm.__version__ != VLLM_VERSION:
@@ -161,6 +170,7 @@ def main() -> None:
"error": result.error, "error": result.error,
}, },
sort_keys=True, sort_keys=True,
default=json_default,
), ),
flush=True, flush=True,
) )
@@ -183,7 +193,9 @@ def main() -> None:
"rows": rows, "rows": rows,
} }
args.output.parent.mkdir(parents=True, exist_ok=True) args.output.parent.mkdir(parents=True, exist_ok=True)
args.output.write_text(json.dumps(payload, indent=2, sort_keys=True) + "\n") args.output.write_text(
json.dumps(payload, indent=2, sort_keys=True, default=json_default) + "\n"
)
if __name__ == "__main__": if __name__ == "__main__":