py/inferscale/api.py

from __future__ import annotations

from .optimizer import capacity_search, compare_schedulers, compare_topologies, design_space_search
from .profiles import ACCELERATORS, MODELS
from .simulator import SCHEDULERS, run_simulation


def metadata() -> dict:
    return {
        "version": "0.3.0",
        "models": list(MODELS.keys()),
        "accelerators": list(ACCELERATORS.keys()),
        "schedulers": sorted(SCHEDULERS),
        "topologies": ["colocated", "disaggregated_pd"],
        "profile_type": "analytical-reference",
    }


def execute(action: str, payload: dict) -> dict:
    if action == "simulate":
        return run_simulation(payload)
    if action == "capacity":
        config = payload.get("config", payload)
        return capacity_search(
            config,
            min_rate=float(payload.get("min_rate", 0.25)),
            max_rate=float(payload.get("max_rate", 32.0)),
            iterations=int(payload.get("iterations", 8)),
            repetitions=int(payload.get("repetitions", 2)),
            headroom=float(payload.get("headroom", 0.20)),
        )
    if action == "compare":
        config = payload.get("config", payload)
        return compare_schedulers(config, payload.get("schedulers"))
    if action == "topology_compare":
        config = payload.get("config", payload)
        return compare_topologies(config)
    if action == "design_space":
        config = payload.get("config", payload)
        return design_space_search(config, bool(payload.get("include_disaggregated", True)))
    if action == "metadata":
        return metadata()
    raise ValueError(f"Unknown action: {action}")