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()
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:
args = parse_args()
if vllm.__version__ != VLLM_VERSION:
@@ -161,6 +170,7 @@ def main() -> None:
"error": result.error,
},
sort_keys=True,
default=json_default,
),
flush=True,
)
@@ -183,7 +193,9 @@ def main() -> None:
"rows": rows,
}
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__":