Serialize vLLM attention profile results
This commit is contained in:
@@ -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__":
|
||||||
|
|||||||
Reference in New Issue
Block a user