py/inferscale/api.py

from __future__ import annotations

from .agentic import (
    adaptive_alpha_sweep,
    adaptive_tiering_study,
    agent_affinity_sweep,
    agent_memory_budget_sweep,
    compare_agent_memory_policies,
    compare_agent_policies,
    run_agent_session_simulation,
    ttl_retention_sweep,
)
from .optimizer import capacity_search, compare_schedulers, compare_topologies, design_space_search
from .profiles import ACCELERATORS, MODELS
from .research import STUDIES, paired_study, robustness_study
from .simulator import SCHEDULERS, run_simulation


def metadata() -> dict:
    return {
        "models": list(MODELS.keys()),
        "accelerators": list(ACCELERATORS.keys()),
        "schedulers": sorted(SCHEDULERS),
        "topologies": ["colocated", "disaggregated_pd"],
        "research_studies": STUDIES,
        "profile_type": "analytical-reference",
        "agentic_modes": [
            "session_simulation", "policy_compare", "ttl_sweep",
            "memory_policy_compare", "memory_budget_sweep", "affinity_sweep",
            "predictive_tiering", "adaptive_alpha_sweep",
        ],
    }


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 == "paired_study":
        config = payload.get("config", payload)
        return paired_study(
            config,
            study=str(payload.get("study", "prefix_cache")),
            repetitions=int(payload.get("repetitions", 12)),
            bootstrap_samples=int(payload.get("bootstrap_samples", 500)),
        )
    if action == "agent_simulate":
        return run_agent_session_simulation(payload.get("config", payload))
    if action == "agent_compare":
        return compare_agent_policies(payload.get("config", payload))
    if action == "agent_ttl_sweep":
        return ttl_retention_sweep(payload.get("config", payload), payload.get("ttl_values"))
    if action == "agent_memory_compare":
        return compare_agent_memory_policies(payload.get("config", payload))
    if action == "agent_memory_sweep":
        return agent_memory_budget_sweep(payload.get("config", payload), payload.get("budget_multipliers"))
    if action == "agent_affinity_sweep":
        return agent_affinity_sweep(payload.get("config", payload), payload.get("slack_values_ms"))
    if action == "agent_predictive_tiering":
        return adaptive_tiering_study(
            payload.get("config", payload),
            horizon_s=float(payload.get("horizon_s", 120.0)),
            shift_fraction=float(payload.get("shift_fraction", 0.55)),
            shift_multiplier=float(payload.get("shift_multiplier", 2.5)),
            alpha=float(payload.get("alpha", 0.30)),
        )
    if action == "agent_adaptive_alpha_sweep":
        return adaptive_alpha_sweep(
            payload.get("config", payload),
            payload.get("alpha_values"),
            horizon_s=float(payload.get("horizon_s", 120.0)),
            shift_fraction=float(payload.get("shift_fraction", 0.55)),
            shift_multiplier=float(payload.get("shift_multiplier", 2.5)),
        )
    if action == "robustness_study":
        config = payload.get("config", payload)
        return robustness_study(
            config,
            study=str(payload.get("study", "pd_vs_colocated")),
            samples=int(payload.get("samples", 32)),
            uncertainty=float(payload.get("uncertainty", 0.20)),
        )
    if action == "metadata":
        return metadata()
    raise ValueError(f"Unknown action: {action}")