scripts/import_measurements.py

from __future__ import annotations

import argparse
import json
from pathlib import Path

from inferscale.measurements import import_measurements
from inferscale.models import SimulationConfig


def main() -> None:
    parser = argparse.ArgumentParser(description="Normalize vLLM/SGLang serving benchmark artifacts for InferScale.")
    parser.add_argument("artifact", type=Path)
    parser.add_argument("--source", choices=["auto", "vllm", "sglang", "case_bundle"], default="auto")
    parser.add_argument("--base-config", type=Path, help="Optional JSON SimulationConfig used when artifact metadata is incomplete.")
    parser.add_argument("--output", type=Path, default=Path("measurement_cases.json"))
    args = parser.parse_args()

    base = SimulationConfig().to_dict()
    if args.base_config:
        base.update(json.loads(args.base_config.read_text()))
    result = import_measurements(args.artifact.read_text(), args.source, base)
    args.output.write_text(json.dumps(result, indent=2) + "\n")
    print(f"Normalized {result['case_count']} cases -> {args.output}")


if __name__ == "__main__":
    main()