src/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 .execution import (
    execution_budget_sweep,
    execution_decay_sweep,
    execution_horizon_sweep,
    execution_planning_study,
    execution_prefetch_study,
    execution_threshold_sweep,
    run_execution_learning,
)
from .optimizer import capacity_search, compare_schedulers, compare_topologies, design_space_search
from .consolidation import repeated_seed_policy_study
from .measurements import calibrate_measurements, import_measurements
from .reports import generate_research_report
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",
        ],
        "execution_learning_modes": [
            "single_run", "prefetch_policy_compare", "confidence_threshold_sweep", "transition_decay_sweep",
            "multistep_planning", "forecast_horizon_sweep", "cache_budget_policy_sweep",
        ],
        "consolidation_modes": ["repeated_seed_policy_study", "measurement_import", "heldout_calibration", "markdown_report"],
    }


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 == "execution_learning_run":
        return run_execution_learning(payload.get("config", payload))
    if action == "execution_prefetch_study":
        return execution_prefetch_study(payload.get("config", payload))
    if action == "execution_threshold_sweep":
        return execution_threshold_sweep(payload.get("config", payload), payload.get("thresholds"))
    if action == "execution_decay_sweep":
        return execution_decay_sweep(payload.get("config", payload), payload.get("decay_values"))
    if action == "execution_planning_study":
        return execution_planning_study(payload.get("config", payload))
    if action == "execution_horizon_sweep":
        return execution_horizon_sweep(payload.get("config", payload), payload.get("horizon_values"))
    if action == "execution_budget_sweep":
        return execution_budget_sweep(payload.get("config", payload), payload.get("budget_values"))
    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 == "consolidation_study":
        return repeated_seed_policy_study(
            payload.get("config", payload),
            repetitions=int(payload.get("repetitions", 12)),
            bootstrap_samples=int(payload.get("bootstrap_samples", 600)),
        )
    if action == "measurement_import":
        return import_measurements(
            str(payload.get("content", "")),
            source=str(payload.get("source", "auto")),
            base_config=payload.get("base_config"),
        )
    if action == "measurement_calibrate":
        return calibrate_measurements(
            payload.get("cases", []),
            holdout_fraction=float(payload.get("holdout_fraction", 0.33)),
            seed=int(payload.get("seed", 7)),
        )
    if action == "research_report":
        return {
            "markdown": generate_research_report(payload.get("robust", {}), payload.get("calibration")),
            "filename": "inferscale-research-consolidation.md",
        }
    if action == "metadata":
        return metadata()
    raise ValueError(f"Unknown action: {action}")