diff --git a/README.md b/README.md index 2e43290be..2bd24ab03 100644 --- a/README.md +++ b/README.md @@ -264,6 +264,7 @@ Use the focused SDK documentation instead of treating CLI internals as public APIs: - [Estimator/FPE Python and Rust SDK](docs/core-api.md) +- [FPM self-service: onboard a model on target hardware](docs/fpm-self-service.md) - [AIC-compatible modeled-power contract (semantics only)](docs/power-model.md) - [Self-benchmarking and FPM onboarding guide](python/aisimulate/docs/fpm/self-benchmarking-and-onboarding.md) - [Replay SDK and artifact contract](crates/core/src/replay/README.md) diff --git a/docs/fpm-self-service.md b/docs/fpm-self-service.md new file mode 100644 index 000000000..71d5fc9f0 --- /dev/null +++ b/docs/fpm-self-service.md @@ -0,0 +1,145 @@ + + +# FPM self-service + +`aisimulate onboard` guides onboarding a new model for FPM simulation on your designated hardware platform. It records the model, runtime, target GPU system and allocation, plans one pure tensor-parallel worker, and produces ordinary `predict` and `recommend` configurations that read your collected FPM data. It builds on AISimulate's existing per-worker FPM support and packaged collector. + +Planning works before the model has an AISimulate model class or measured FPM timings. A valid request records your choices; model metadata and execution readiness, runtime compatibility, and data readiness remain **unchecked**, and accuracy is **not assessed**. This setup does not run target preflight checks, provision GPU resources, or establish measured accuracy. + +In this revision, ordinary FPM `predict` and `recommend` still construct a registered analytical model for resource accounting and SOL-based timing transfer. A separate decoupling change will add a route that uses model/resource metadata and direct interpolation of measured timings without an op-level model class. The existing registered-model/SOL route will remain available. See [Choose the model execution route](#choose-the-model-execution-route) before handing off a new architecture. + +The Inkling pilot is planned for NVIDIA GB200 with TP, DEP, and TEP configurations through the class-independent route. This guided planner currently generates pure-TP configurations only. The checkpoint, runtime, allocation, and strategy degrees must be pinned before that pilot; this guide does not establish Inkling readiness or GB200 accuracy. + +## Create the request + +Use an environment installed from this checkout; see [development setup](../DEVELOPMENT.md). Guided setup requires a terminal and starts only when explicitly requested: + +```bash +aisimulate onboard init --interactive --output support-request.yaml +``` + +Enter the actual model identifier or checkpoint path, pinned model revision, dense/MoE kind, pinned vLLM version, GPU system, allocation, and interconnect. Then choose a TP size and pilot workload. Supplied options skip their prompts. Enter accepts displayed defaults; invalid values can be corrected; Ctrl-C or end-of-input cancels without saving. Existing files require `--overwrite`. + +Both guided and scripted setup use onboarding. `--profile onboarding` is an optional spelling of the same behavior. For automation, supply the identity flags directly: + +```bash +aisimulate onboard init \ + --model /models/your-pinned-checkpoint \ + --model-revision YOUR_IMMUTABLE_REVISION \ + --model-kind dense \ + --framework-version YOUR_PINNED_VLLM_VERSION \ + --gpu h200_sxm --gpu-count 4 --interconnect nvswitch \ + --tensor-parallel 2 \ + --output support-request.yaml +``` + +Replace the model and runtime placeholders with the actual inputs. The example does not identify an Inkling checkpoint or claim that four H200s can run your model. Framework support currently selects vLLM. Optional tokenizer, chat-template, and AISimulate revisions are recorded only when supplied. + +The default pilot uses 1,024 input tokens, 128 output tokens, concurrency 1, four requests, a 16,384-token context limit, TTFT target 1,000 ms, and TPOT target 100 ms. Scripted setup defaults to TP1 unless `--tensor-parallel` is supplied. These are planning defaults, not measured model capacity or latency. Change them with the corresponding flags shown by `aisimulate onboard init --help`. + +One node is the default; GPUs per node then equals `--gpu-count`. For multiple nodes, supply both `--node-count` and `--gpus-per-node`; their product must equal the total allocation. Each pure-TP worker must fit on one node. Advanced flags include `--request-count`, `--max-candidates`, `--objective`, `--seed`, and `--sm`. + +## Plan, preview, and explicitly execute + +```bash +aisimulate onboard plan \ + --config support-request.yaml --output-dir ./aisimulate-support + +aisimulate onboard collect-fpm \ + --config ./aisimulate-support/request.yaml \ + --output-dir ./aisimulate-support +``` + +The first command saves the request, `support-plan.json`, `commands.json`, `predict/pilot.yaml`, `recommend/pilot.yaml`, and a local `systems/` directory. The second prints the collector command without launching it. Generated command vectors and printed next commands use absolute output paths and preserve spaces or shell punctuation. Use a separate output directory for each request. `--overwrite` can repair missing generated files for the identical saved request, including an interrupted write before `support-plan.json` was created. Repair requires an intact matching `request.yaml`, rejects changed inputs, and preserves existing collected data. + +If a software update changes generated guidance, recreate the plan in a new output directory: repair compares generated files byte for byte and does not migrate older plans. Guidance changes do not change the saved-request identity checks used by collection. + +The search uses one selected TP size. By default it evaluates a single worker, so recommendation is not a broad deployment search. `--max-candidates 2` additionally considers the largest count of identical workers that fits the allocation, when that differs from one worker. Each choice gets an independent recommendation config pinned to that replica count with a one-trial budget. The single worker keeps `recommend/pilot.yaml`; the second choice uses `recommend/replicas-N.yaml`, where `N` is its replica count. The plan reports the actual candidate count and lists both config and result paths. Dense collection uses the `tp` preset; MoE uses `pure_tp`; the selected TP size remains exact. + +Before execution, prepare the real checkpoint and the pinned runtime using the existing [FPM collection guide](../python/aisimulate/docs/fpm/self-benchmarking-and-onboarding.md). The packaged collector invokes a Generator-resolved Dynamo/vLLM deployment and needs the corresponding GPU resources, deployment configuration, permissions, and model access. Invoking its command locally does not create that environment. `commands.json` publishes the guarded `aisimulate onboard collect-fpm --execute` command for collection, alongside a read-only collector planning command. + +After those prerequisites are ready, explicitly launch collection from that environment: + +```bash +aisimulate onboard collect-fpm \ + --config ./aisimulate-support/request.yaml \ + --output-dir ./aisimulate-support --execute +``` + +Execution requires a matching saved plan. The request records model and runtime revisions; this setup does not download a pinned checkpoint or check the installed runtime before execution. Keep the actual checkpoint and runtime consistent with the request before collecting or predicting. Successful formal collection and resume verify that the published data's pod-reported runtime version exactly matches `framework_version`, including suffixes such as `+cu128`. A mismatch exits 1 and preserves collection artifacts. Use the declared runtime or create a new request and plan for the observed version; generated configs are never silently retargeted. Diagnostic smoke runs do not publish or verify formal data. + +Set deployment options directly on `onboard collect-fpm`: `--dynamo-version VERSION`, `--image IMAGE`, `--namespace NAME`, `--model-cache NAME[:MOUNT[:SUBPATH]]`, `--transport nvlink|ib|efa`, and `--image-pull-secret NAME`. The mount, when supplied, is an absolute container path. Prefer an immutable image digest. Supply the same options when previewing, executing, and resuming; deployment settings are part of the collector's frozen-plan identity, so changed settings require a new output directory. Arbitrary collector arguments and engine overrides are not accepted by this command. + +For a diagnostic run, add `--execute --smoke`; `--limit N` also requires `--smoke`. Diagnostic smoke and limited runs do not publish formal FPM data. Existing campaign data, raw artifacts, or checkpoints require explicit `--resume` and a readable matching collector checkpoint; otherwise choose a new output directory. A custom `--checkpoint-dir`, if needed, must remain inside the plan's `fpm-checkpoint/` directory. Selecting an empty checkpoint directory does not allow reuse of existing campaign artifacts. Smoke and formal campaigns have separate checkpoints and artifact directories, so an existing smoke run does not prevent the first formal run, or vice versa. The collector verifies the resumed checkpoint's frozen-plan identity. + +Planning and collection reject concurrent onboarding operations. The persistent `.support.lock` file uses an OS advisory lock; ownership is released when the process exits, including after an abrupt termination. Leave the file in place. Request validation, saved onboarding-plan checks, and collector input resolution exit 2. Failures after collector execution starts, including a frozen checkpoint identity mismatch, exit 1 with a concise message. Interruption exits 130. + +The collector narrows initial prefill sampling with the pilot's input-token and concurrency bounds. Decode uses the collector's existing profile; a four-request synthetic pilot does not imply four timing samples or a short decode campaign. Inspect the generated command and collector plan before committing GPU time. Successful formal collection publishes the FPM Parquet file and metadata pair into the plan's local systems data directory; diagnostic success alone does not provide that pair. + +## Run the generated ordinary configurations + +After the selected model execution route is available and formal data collection is complete: + +```bash +aisimulate predict \ + --config ./aisimulate-support/predict/pilot.yaml \ + --output-dir ./aisimulate-support/predict-results/pilot + +aisimulate recommend \ + --config ./aisimulate-support/recommend/pilot.yaml \ + --output-dir ./aisimulate-support/recommend-results/pilot +``` + +With two candidates, run both replica counts. This loop derives the candidate names from the validated saved request and constructs fixed `aisimulate recommend` commands: + +```bash +python3 - <<'PY' +import shlex +import subprocess +from pathlib import Path + +from aisimulate.support.plan import check_plan +from aisimulate.support.schema import SupportRequest + +root = Path("./aisimulate-support").resolve() +request = SupportRequest.from_yaml(root / "request.yaml") +check_plan(request, root) +max_replicas = request.identity.node_count * (request.identity.gpus_per_node // request.search.tensor_parallel) +replicas = list(dict.fromkeys((1, max_replicas)))[:request.search.max_candidates] +statuses = [] +for count in replicas: + name = "pilot" if count == 1 else f"replicas-{count}" + command = [ + "aisimulate", "recommend", + "--config", str(root / f"recommend/{name}.yaml"), + "--output-dir", str(root / f"recommend-results/{name}"), + "--format", "json", + ] + result = subprocess.run(command, check=False) + statuses.append(result.returncode) + print(f"exit {result.returncode}: {shlex.join(command)}", flush=True) +raise SystemExit(1 if any(statuses) else 0) +PY +``` + +The loop reports each command's exit status and attempts all commands, including when the pilot finds no feasible candidate. It exits 1 after all attempts if any command returned a nonzero status, otherwise 0. Successful outputs remain usable even when the loop exits 1; inspect each result before comparing candidates. + +Run either the individual recommendation command or the loop against fresh result directories. The loop also handles a one-candidate plan and does not execute entries from `commands.json`. Results remain separate under `recommend-results/pilot` and, when present, `recommend-results/replicas-N`; compare their objective and latency results for the same workload. This plan does not produce a combined ranking or search additional TP sizes or scheduler settings. + +The generated configurations select `engine.workers.aggregated.timing.estimation_mode: fpm_interpolation` with `fallback_policy: deny` and set `engine.systems_paths` to a list containing the plan's absolute local systems directory. The same root supplies hardware and collected FPM data. Recommendation preserves it in exported prediction configs. Moving the plan to another machine requires updating absolute paths or regenerating it there. The single-root `engine.systems_path` input is an alias added by onboarding; saved configurations use the canonical `engine.systems_paths` list. Both spellings resolve relative local paths against the working directory when the configuration is loaded, so saving and reloading from another directory preserves the selected roots. The canonical list's `default` entry continues to select packaged data. A worker’s `timing.systems_paths` overrides engine-level roots; performance metadata records those effective roots. Empty or whitespace-only roots are rejected by the shared root type in both prediction and recommendation configurations. + +These configurations target AISimulate's standalone `predict` and `recommend` commands. Dynamo's replay Planner AIC session adapter does not currently forward custom systems roots; using that downstream path requires a separate adapter update. + +Ordinary `predict` and `recommend` commands retain their existing behavior and defaults. Their generated configs can be loaded and edited through the public configuration schema. In this revision, missing model registration or data may still prevent execution; a successful simulation is not an accuracy result. Compare its output with an independent run of the same model, runtime, topology, and workload to assess accuracy. + +## Choose the model execution route + +Hand off the saved request, pinned model configuration, and plan. Both routes need a canonical checkpoint identity, effective precision and topology, correct weight and KV-cache accounting, and matching whole-forward FPM measurements. Collected timings alone do not establish memory fit. + +- **Registered-model/SOL route, available in this revision:** reuse a compatible analytical class or follow [How to Add a New Model](../python/aisimulate/docs/add_a_new_model.md) when choosing to add one. Verify its operation graph, memory and cache accounting, and native FPM SOL execution. A dedicated class is an option for this route, not the intended prerequisite for every FPM onboarding. +- **Class-independent route, planned in the separate decoupling change:** resolve model and resource metadata without constructing an operation graph. Use direct measured-time interpolation, including wider two-sided KV brackets at the same batch size when both neighboring prompt curves cover the query. Missing brackets or unsupported metadata must produce explicit errors. This route is not implemented by the guided foundation; 2D interpolation remains experimental. + +Per-operation silicon profiling described in the model guide is not required by either FPM route. The intended self-service workflow collects whole-forward timings, then verifies prediction and recommendation for the exact target deployment. Collection bootstrap can already resolve some unregistered model configurations; that does not establish that this revision's ordinary FPM prediction path can construct them. diff --git a/python/aisimulate/collector/fpm_forward/cli.py b/python/aisimulate/collector/fpm_forward/cli.py index fd674867f..7096f5933 100644 --- a/python/aisimulate/collector/fpm_forward/cli.py +++ b/python/aisimulate/collector/fpm_forward/cli.py @@ -16,7 +16,7 @@ from .config import add_fpm_arguments, add_fpm_generator_arguments from .entry import resolve_inputs, resolve_run_inputs, run_resolved -_INPUT_ERRORS = (FileNotFoundError, RuntimeError, TypeError, ValueError) +_INPUT_ERRORS = (OSError, RuntimeError, TypeError, ValueError) def _parser() -> argparse.ArgumentParser: diff --git a/python/aisimulate/src/aisimulate/capacity.py b/python/aisimulate/src/aisimulate/capacity.py index 21b85cdf7..c7d48bd13 100644 --- a/python/aisimulate/src/aisimulate/capacity.py +++ b/python/aisimulate/src/aisimulate/capacity.py @@ -252,7 +252,7 @@ def estimate_num_gpu_blocks( attention_backend: str | None = None, enable_eplb: bool = False, wideep_num_slots: int | None = None, - systems_path: str | None = None, + systems_path: str | list[str] | None = None, cuda_graph_reserved_bytes: int = 0, diagnostics: dict[str, Any] | None = None, ) -> int: diff --git a/python/aisimulate/src/aisimulate/cli_args.py b/python/aisimulate/src/aisimulate/cli_args.py index 4f8b7f348..4947d8094 100644 --- a/python/aisimulate/src/aisimulate/cli_args.py +++ b/python/aisimulate/src/aisimulate/cli_args.py @@ -11,6 +11,7 @@ import yaml from .detail import parse_detail_sections +from .support.cli import add_support_parser class _CliConfigError(ValueError): @@ -75,6 +76,7 @@ def build_parser() -> argparse.ArgumentParser: action="store_true", help="pace prediction against the real wall clock instead of virtual time", ) + add_support_parser(subparsers) return parser diff --git a/python/aisimulate/src/aisimulate/compiler.py b/python/aisimulate/src/aisimulate/compiler.py index fbdbcdaaa..7089d0e09 100644 --- a/python/aisimulate/src/aisimulate/compiler.py +++ b/python/aisimulate/src/aisimulate/compiler.py @@ -420,6 +420,9 @@ def _worker_performance_model_metadata( if value is not None: config[field] = value config["database_mode"] = worker.timing.database_mode or engine.database_mode + systems_paths = worker.timing.systems_paths or engine.systems_paths + if systems_paths is not None: + config["systems_paths"] = systems_paths if engine.speculation is not None: config["speculation"] = engine.speculation.cost_config() if worker.timing.fpm_parquet_path is not None: @@ -488,6 +491,10 @@ def _worker_engine_args( payload["speculation"] = engine.speculation.model_dump(mode="json") if engine.backend_version is not None: payload["aic_backend_version"] = engine.backend_version + if engine.systems_paths is not None and worker.timing.type != "default": + from .sweeper.forward_pass_estimator import resolve_systems_paths + + payload["systems_path"] = list(resolve_systems_paths(engine.systems_paths)) if engine.decoder_replay: payload["aic_decoder_replay"] = True for field in ("database_mode", "enable_shared_layer", "strict_provenance"): @@ -543,6 +550,7 @@ def _worker_engine_args( if capacity.type == "default" and cache.state_cache is None: payload = materialize_aic_num_gpu_blocks(payload) for name in ( + "systems_path", "aic_backend_version", "aic_system", "aic_model_path", diff --git a/python/aisimulate/src/aisimulate/config/common.py b/python/aisimulate/src/aisimulate/config/common.py index 3a5d59da2..1d9fa8df1 100644 --- a/python/aisimulate/src/aisimulate/config/common.py +++ b/python/aisimulate/src/aisimulate/config/common.py @@ -11,13 +11,33 @@ from typing import Annotated, Any, Generic, Literal, TypeVar import yaml -from pydantic import BaseModel, ConfigDict, Field, field_validator, model_validator +from pydantic import AfterValidator, BaseModel, ConfigDict, Field, field_validator, model_validator class StrictModel(BaseModel): model_config = ConfigDict(extra="forbid") +def _absolute_systems_path(value: str) -> str: + if not value.strip() or "," in value: + raise ValueError("systems_path must name one nonempty local directory, not a comma-separated list") + return str(Path(value).expanduser().resolve()) + + +SystemsPath = Annotated[str, Field(strict=True, min_length=1), AfterValidator(_absolute_systems_path)] + + +def _normalize_systems_root(value: str) -> str: + if not value.strip(): + raise ValueError("systems_paths entries must be nonempty") + if value.lower() == "default": + return "default" + return str(Path(value).expanduser().resolve()) + + +SystemsRoot = Annotated[str, Field(strict=True, min_length=1), AfterValidator(_normalize_systems_root)] + + T = TypeVar("T") PositiveFiniteFloat = Annotated[float, Field(strict=True, gt=0, allow_inf_nan=False)] PositiveStrictInt = Annotated[int, Field(strict=True, gt=0)] diff --git a/python/aisimulate/src/aisimulate/config/engine.py b/python/aisimulate/src/aisimulate/config/engine.py index 0ad481544..df3658305 100644 --- a/python/aisimulate/src/aisimulate/config/engine.py +++ b/python/aisimulate/src/aisimulate/config/engine.py @@ -16,6 +16,8 @@ IntegerRange, NumericRange, StrictModel, + SystemsPath, + SystemsRoot, is_active_engine_model_control, ) @@ -264,7 +266,7 @@ class TimingConfig(StrictModel): estimation_mode: Literal["auto", "op_level", "fpm_interpolation", "fpm_regression"] | None = None fallback_policy: Literal["deny", "allow"] | None = None estimator_config: dict[str, Any] | None = None - systems_paths: list[str] | None = Field(default=None, min_length=1) + systems_paths: list[SystemsRoot] | None = Field(default=None, min_length=1) database_mode: Literal["SILICON", "HYBRID", "EMPIRICAL", "SOL"] | None = None transfer_policy: str | list[str] | None = None @@ -412,25 +414,25 @@ def _validate_engine_controls(self): database_mode: Literal["SILICON", "HYBRID", "EMPIRICAL", "SOL"] = "SILICON" transfer_policy: str | list[str] | None = None - systems_paths: list[str] | None = None + systems_paths: list[SystemsRoot] | None = Field(default=None, min_length=1) + systems_path: SystemsPath | None = Field(default=None, exclude=True) estimation_mode: Literal["auto", "op_level", "fpm_interpolation", "fpm_regression"] = "auto" fallback_policy: Literal["deny", "allow"] = "deny" estimator_config: dict[str, Any] = Field(default_factory=dict) + @model_validator(mode="after") + def _migrate_systems_path(self): + if self.systems_path is not None: + if self.systems_paths is not None and self.systems_paths != [self.systems_path]: + raise ValueError("systems_path conflicts with systems_paths") + self.systems_paths = [self.systems_path] + return self + @field_validator("database_mode", mode="before") @classmethod def _normalize_database_mode(cls, value): return value.upper() if isinstance(value, str) else value - @field_validator("systems_paths") - @classmethod - def _nonempty_system_roots(cls, value): - if value is None: - return value - if not value or any(not path.strip() for path in value): - raise ValueError("systems_paths must contain at least one nonempty root") - return value - @model_validator(mode="after") def _supported_estimator_policies(self): mode = getattr(self, "mode", "aggregated") @@ -439,7 +441,6 @@ def _supported_estimator_policies(self): custom_policy = ( self.database_mode != "SILICON" or self.transfer_policy is not None - or self.systems_paths not in (None, ["default"]) or self.estimation_mode != "auto" or self.fallback_policy != "deny" or bool(self.estimator_config) @@ -449,6 +450,9 @@ def _supported_estimator_policies(self): ) roles = [getattr(workers, role, None) for role in ("aggregated", "prefill", "decode")] unsupported_provider = "afd" in modes or getattr(workers, "encoder", None) is not None + if self.systems_paths not in (None, ["default"]) and unsupported_provider: + unsupported = "AFD" if "afd" in modes else "analytical encoder pools" + raise ValueError(f"engine.systems_paths does not support {unsupported}") if unsupported_provider and any( worker is not None and ( diff --git a/python/aisimulate/src/aisimulate/main.py b/python/aisimulate/src/aisimulate/main.py index 99804c640..817679d15 100644 --- a/python/aisimulate/src/aisimulate/main.py +++ b/python/aisimulate/src/aisimulate/main.py @@ -48,6 +48,7 @@ workload_bounds, ) from .stack import StackResolutionError, resolve_runner_factory +from .support.cli import run_support_command from .sweeper.provider import AdapterReplaySpec from .sweeper.replay import ReplayOutputRequirements @@ -299,6 +300,13 @@ def _write_resource_plan(args, plan: dict[str, Any]) -> None: def main(argv: Sequence[str] | None = None) -> int: parser = build_parser() args = parser.parse_args(list(sys.argv[1:] if argv is None else argv)) + if args.command == "onboard": + try: + return run_support_command(args) + except (ValidationError, ValueError, OSError) as exc: + parser.error(str(exc)) + except KeyboardInterrupt: + return 130 # Stack resolution deliberately precedes opening the configuration file. try: factory = resolve_runner_factory(args.stack) diff --git a/python/aisimulate/src/aisimulate/recommend.py b/python/aisimulate/src/aisimulate/recommend.py index cccb8e454..ff902df5c 100644 --- a/python/aisimulate/src/aisimulate/recommend.py +++ b/python/aisimulate/src/aisimulate/recommend.py @@ -796,6 +796,8 @@ def _candidate_prediction( "context_length": sample.get("context_length") or "max", "workers": {}, } + if sample.get("systems_paths") is not None: + engine["systems_paths"] = sample["systems_paths"] for name in (*ENGINE_MODEL_CONTROL_FIELDS, "enable_chunked_prefill", "nextn_accepted"): if sample.get(name) is not None: engine[name] = sample[name] diff --git a/python/aisimulate/src/aisimulate/supervision.py b/python/aisimulate/src/aisimulate/supervision.py index 1c8438ea3..682f86153 100644 --- a/python/aisimulate/src/aisimulate/supervision.py +++ b/python/aisimulate/src/aisimulate/supervision.py @@ -342,6 +342,15 @@ def main(argv: Sequence[str] | None = None) -> int: arguments = list(sys.argv[1:] if argv is None else argv) parser = build_parser() args = parser.parse_args(arguments) + if args.command == "onboard": + from .support.cli import run_support_command + + try: + return run_support_command(args) + except (ValueError, OSError) as exc: + parser.error(str(exc)) + except KeyboardInterrupt: + return 130 raw = None try: raw = _load_mapping(args.config) diff --git a/python/aisimulate/src/aisimulate/support/__init__.py b/python/aisimulate/src/aisimulate/support/__init__.py new file mode 100644 index 000000000..e5bce66e6 --- /dev/null +++ b/python/aisimulate/src/aisimulate/support/__init__.py @@ -0,0 +1,8 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Request and planning helpers for local FPM onboarding.""" + +from .schema import SupportRequest + +__all__ = ["SupportRequest"] diff --git a/python/aisimulate/src/aisimulate/support/cli.py b/python/aisimulate/src/aisimulate/support/cli.py new file mode 100644 index 000000000..e5a06ce1d --- /dev/null +++ b/python/aisimulate/src/aisimulate/support/cli.py @@ -0,0 +1,332 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Terminal setup and local planning for FPM onboarding.""" + +from __future__ import annotations + +import argparse +import json +import os +import shlex +import sys +import tempfile +from pathlib import Path +from typing import Any, get_args + +import yaml +from pydantic import ValidationError + +from .fpm import run_fpm +from .plan import create_plan, request_id +from .schema import FPMDeployment, SearchProfile, SloSpec, SupportIdentity, SupportRequest, WorkloadSpec + + +def add_support_parser(subparsers: Any) -> None: + onboard = subparsers.add_parser( + "onboard", + help="Onboard a model for FPM simulation on a target hardware platform.", + description=( + "Onboard a model for FPM simulation on a target hardware platform. " + "Declare the model and hardware, generate a plan, and preview or execute FPM collection." + ), + ) + actions = onboard.add_subparsers(dest="support_action", required=True) + init = actions.add_parser( + "init", + help="Declare a model, target hardware, and FPM pilot workload.", + description=( + "Use --interactive for terminal prompts. Scripted setup requires --model, --model-revision, " + "--model-kind, --framework-version, --gpu, --gpu-count, and --interconnect. " + "One node and a small TP1 synthetic pilot are defaults; review the saved scope before collection." + ), + ) + init.add_argument("--interactive", action="store_true", help="Guide setup with terminal prompts.") + init.add_argument("--profile", choices=("onboarding",), help="Optional alias; all setup uses onboarding.") + init.add_argument("--model") + init.add_argument("--model-revision", help="Pinned model revision; moving labels such as main are unsupported.") + init.add_argument("--model-kind", choices=("dense", "moe")) + init.add_argument("--framework", choices=("vllm",), help="Collection runtime (default: vllm).") + init.add_argument("--framework-version", help="Pinned vLLM version in the collection environment.") + init.add_argument("--gpu", help="Target GPU system name, for example h200_sxm.") + init.add_argument("--gpu-count", type=int, help="Total GPUs available.") + init.add_argument("--node-count", type=int, help="Number of nodes (default: 1).") + init.add_argument("--gpus-per-node", type=int, help="GPUs available per node; defaults to gpu-count on one node.") + init.add_argument("--interconnect", help="Interconnect, for example nvswitch, pcie, or none.") + init.add_argument("--sm", type=int, help="GPU SM version override.") + init.add_argument("--tokenizer-revision") + init.add_argument("--chat-template-revision") + init.add_argument("--aisimulate-revision", help="Optional pinned AISimulate source revision.") + init.add_argument("--tensor-parallel", type=int, help="GPUs per pure-TP worker (default: 1).") + init.add_argument("--input-tokens", type=int, help="Input tokens per request (default: 1024).") + init.add_argument("--output-tokens", type=int, help="Output tokens per request (default: 128).") + init.add_argument("--concurrency", type=int, help="Concurrent requests (default: 1).") + init.add_argument("--request-count", type=int, help="Synthetic request count (default: 4).") + init.add_argument("--context-length", type=int, help="Pilot context limit in tokens (default: 16384).") + init.add_argument("--ttft-ms", type=float, help="Target time to first token in ms (default: 1000).") + init.add_argument("--tpot-ms", type=float, help="Target time per output token in ms (default: 100).") + init.add_argument("--max-candidates", type=int, help="One worker, or also the largest replica count (1 or 2).") + init.add_argument( + "--objective", + choices=get_args(SearchProfile.model_fields["objective"].annotation), + help="Recommendation objective (default: throughput).", + ) + init.add_argument("--seed", type=int, help="Recommendation search seed (default: 42).") + init.add_argument("--output", default="support-request.yaml") + init.add_argument("--overwrite", action="store_true") + + plan = actions.add_parser( + "plan", help="Plan FPM collection and write predict/recommend configs for the target hardware." + ) + plan.add_argument("-c", "--config", required=True) + plan.add_argument("--output-dir", default="./aisimulate-support") + plan.add_argument("--overwrite", action="store_true") + plan.add_argument("--format", choices=("table", "json"), default="table") + + collect = actions.add_parser( + "collect-fpm", help="Preview FPM collection on the target hardware, or run with --execute." + ) + collect.add_argument("-c", "--config", required=True) + collect.add_argument("--output-dir", default="./aisimulate-support") + collect.add_argument( + "--execute", action="store_true", help="Execute collection using a matching saved onboarding plan." + ) + collect.add_argument( + "--smoke", action="store_true", help="Diagnostic smoke collection; does not publish formal data." + ) + collect.add_argument( + "--limit", type=int, help="Limit diagnostic cases; requires --smoke, no formal data publication." + ) + collect.add_argument("--resume", action="store_true", help="Resume the existing collector checkpoint.") + collect.add_argument("--checkpoint-dir", help="Checkpoint directory inside the plan's fpm-checkpoint directory.") + deployment = collect.add_argument_group("Collector deployment") + deployment.add_argument("--dynamo-version", help="Target Dynamo release used to resolve collector templates.") + deployment.add_argument("--image", help="Collector container image; prefer an immutable digest.") + deployment.add_argument("--namespace", help="Kubernetes namespace for collector resources.") + deployment.add_argument("--model-cache", metavar="NAME[:MOUNT[:SUBPATH]]", help="Model-cache PVC and mount.") + deployment.add_argument("--transport", choices=("nvlink", "ib", "efa")) + deployment.add_argument("--image-pull-secret", help="Kubernetes secret for pulling the collector image.") + + +def _values(args: argparse.Namespace, model: Any) -> dict[str, Any]: + return {name: getattr(args, name) for name in model.model_fields if getattr(args, name, None) is not None} + + +def _request_from_args(args: argparse.Namespace) -> SupportRequest: + identity = _values(args, SupportIdentity) + if args.gpus_per_node is None and identity.get("node_count", 1) == 1: + identity["gpus_per_node"] = args.gpu_count + return SupportRequest.model_validate( + { + "identity": identity, + "workload": {**_values(args, WorkloadSpec), "slo": _values(args, SloSpec)}, + "search": _values(args, SearchProfile), + } + ) + + +def _request_target(path: str | Path, *, overwrite: bool) -> Path: + target = Path(path).expanduser().absolute() + if target.exists() and not target.is_file(): + raise ValueError(f"onboarding request {target} must be a file path") + if (target.exists() or target.is_symlink()) and not overwrite: + raise ValueError(f"onboarding request {target} exists; pass --overwrite to replace it") + parent = target.parent + while not parent.exists() and not parent.is_symlink(): + parent = parent.parent + if not parent.is_dir() or not os.access(parent, os.W_OK): + raise ValueError(f"onboarding request parent {parent} must be a writable directory") + return target + + +def _write_request(request: SupportRequest, path: str | Path, *, overwrite: bool) -> Path: + target = _request_target(path, overwrite=overwrite) + payload = yaml.safe_dump(request.model_dump(mode="json", exclude_none=True), sort_keys=False) + target.parent.mkdir(parents=True, exist_ok=True) + temporary = None + try: + with tempfile.NamedTemporaryFile(mode="w", encoding="utf-8", dir=target.parent, delete=False) as stream: + temporary = Path(stream.name) + stream.write(payload) + if overwrite: + os.replace(temporary, target) + else: + os.link(temporary, target) + finally: + if temporary is not None: + temporary.unlink(missing_ok=True) + return target + + +def _print(value: dict[str, Any], output_format: str = "table") -> None: + if output_format == "json": + print(json.dumps(value, sort_keys=True)) + return + for key, item in value.items(): + rendered = json.dumps(item, sort_keys=True) if isinstance(item, (dict, list)) else str(item) + print(f"{key}: {rendered}") + + +_PROMPTS = { + "model": ("Model name or path", str), + "model_revision": ("Pinned model revision (not main/latest)", str), + "model_kind": ("Model kind (dense/moe)", str), + "framework_version": ("Pinned vLLM version", str), + "gpu": ("Target GPU system name (for example h200_sxm)", str), + "gpu_count": ("Total GPUs available", int), + "interconnect": ("GPU interconnect (for example nvswitch, pcie, or none)", str), + "tensor_parallel": ("GPUs per pure-TP worker", int), + "input_tokens": ("Input tokens per request", int), + "output_tokens": ("Output tokens per request", int), + "concurrency": ("Concurrent requests", int), + "context_length": ("Pilot context limit in tokens", int), + "ttft_ms": ("Target time to first token (ms)", float), + "tpot_ms": ("Target time per output token (ms)", float), +} +_CORRECTION_PROMPTS = { + **_PROMPTS, + "framework": ("Runtime (vllm)", str), + "tokenizer_revision": ("Pinned tokenizer revision", str), + "chat_template_revision": ("Pinned chat-template revision", str), + "aisimulate_revision": ("Pinned AISimulate revision", str), + "node_count": ("Number of nodes", int), + "gpus_per_node": ("GPUs per node", int), + "sm": ("GPU SM version", int), + "request_count": ("Number of synthetic requests", int), + "max_candidates": ("Maximum candidates (1 or 2)", int), + "objective": ("Recommendation objective", str), + "seed": ("Recommendation search seed", int), +} + + +def _prompt(args: argparse.Namespace, name: str) -> None: + label, convert = _CORRECTION_PROMPTS[name] + default = getattr(args, name) + if default is None: + for model in (SupportIdentity, WorkloadSpec, SloSpec, SearchProfile): + field = model.model_fields.get(name) + if field is not None and not field.is_required(): + default = field.default + if name == "gpu_count": + default = 1 + while True: + answer = input(f"{label}{f' [{default}]' if default is not None else ''}: ").strip() + if not answer and default is None: + print("A value is required.") + continue + try: + value = convert(answer) if answer else default + except ValueError: + print(f"Enter a valid {'integer' if convert is int else 'number'}.") + continue + setattr(args, name, value) + return + + +def _guided_request(args: argparse.Namespace) -> SupportRequest: + if not sys.stdin.isatty(): + raise ValueError( + "--interactive requires a terminal; for automation supply options (aisimulate onboard init --help)" + ) + print("Onboard a model for FPM simulation on a target hardware platform.") + print("Start with a pure-TP worker and a small synthetic workload on vLLM.") + print("Enter accepts a displayed default. Ctrl-C cancels without saving. Supplied options skip their prompts.") + for name in _PROMPTS: + if getattr(args, name) is None: + _prompt(args, name) + while True: + try: + return _request_from_args(args) + except ValidationError as exc: + error = exc.errors(include_url=False)[0] + name = error["loc"][-1] if error["loc"] else None + print(f"{name or 'Request'}: {error['msg']}") + while name not in _CORRECTION_PROMPTS: + name = ( + input("Option to correct (for example --tensor-parallel): ") + .strip() + .removeprefix("--") + .replace("-", "_") + ) + if name not in _CORRECTION_PROMPTS: + print("Choose one of: " + ", ".join(key.replace("_", "-") for key in _CORRECTION_PROMPTS)) + _prompt(args, name) + + +def _init(args: argparse.Namespace) -> int: + _request_target(args.output, overwrite=args.overwrite) + try: + request = _guided_request(args) if args.interactive else _request_from_args(args) + path = _write_request(request, args.output, overwrite=args.overwrite) + except (EOFError, KeyboardInterrupt): + print("Setup cancelled.", file=sys.stderr) + return 130 + workload = request.workload + print( + f"Scope: FPM simulation for {request.identity.model} ({request.identity.model_kind}), " + f"vLLM {request.identity.framework_version}, " + f"{request.identity.gpu_count} {request.identity.gpu} GPU(s) on {request.identity.node_count} node(s); " + f"TP{request.search.tensor_parallel}, {workload.input_tokens}/{workload.output_tokens} tokens, " + f"concurrency {workload.concurrency}, {workload.request_count} requests, " + f"up to {request.search.max_candidates} candidate(s)." + ) + print( + "Request saved. Model integration, runtime compatibility and FPM data are unchecked; " + "accuracy is not assessed. Setup has not launched GPU work." + ) + _print( + { + "request_id": request_id(request), + "request": str(path), + "next": shlex.join(["aisimulate", "onboard", "plan", "--config", str(path)]), + } + ) + return 0 + + +def _plan(args: argparse.Namespace) -> int: + request = SupportRequest.from_yaml(args.config) + root = Path(args.output_dir).expanduser().resolve() + plan = create_plan(request, root, overwrite=args.overwrite) + _print( + { + "request_id": plan["request_id"], + "candidate_count": plan["search"]["candidate_count"], + "plan": str(root / "support-plan.json"), + "commands": plan["outputs"]["commands"], + "prerequisites": "unchecked", + "accuracy": "not assessed", + "next": shlex.join( + [ + "aisimulate", + "onboard", + "collect-fpm", + "--config", + plan["outputs"]["request"], + "--output-dir", + str(root), + ] + ), + }, + args.format, + ) + return 0 + + +def run_support_command(args: argparse.Namespace) -> int: + if args.support_action == "init": + return _init(args) + if args.support_action == "plan": + return _plan(args) + if args.support_action == "collect-fpm": + return run_fpm( + SupportRequest.from_yaml(args.config), + output_dir=args.output_dir, + execute=args.execute, + smoke=args.smoke, + limit=args.limit, + resume=args.resume, + checkpoint_dir=args.checkpoint_dir, + deployment=FPMDeployment(**_values(args, FPMDeployment)), + ) + raise AssertionError(f"unhandled onboarding action {args.support_action!r}") diff --git a/python/aisimulate/src/aisimulate/support/fpm.py b/python/aisimulate/src/aisimulate/support/fpm.py new file mode 100644 index 000000000..af53c8fda --- /dev/null +++ b/python/aisimulate/src/aisimulate/support/fpm.py @@ -0,0 +1,181 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Preview or explicitly invoke the existing packaged FPM collector.""" + +from __future__ import annotations + +import json +import shlex +import sys +from pathlib import Path + +from .schema import FPMDeployment, SupportRequest + + +def fpm_cli_args( + request: SupportRequest, + *, + output_dir: str | Path, + plan_only: bool, + smoke: bool = False, + limit: int | None = None, + resume: bool = False, + checkpoint_dir: str | Path | None = None, + deployment: FPMDeployment | None = None, +) -> list[str]: + if limit is not None and (not smoke or type(limit) is not int or limit < 1): + raise ValueError("limit must be a positive cell count and requires smoke=True") + root = Path(output_dir).expanduser().resolve() + checkpoint = Path(checkpoint_dir).expanduser().resolve() if checkpoint_dir else root / "fpm-checkpoint" + if checkpoint != root / "fpm-checkpoint" and root / "fpm-checkpoint" not in checkpoint.parents: + raise ValueError("checkpoint_dir must stay within the plan's fpm-checkpoint directory") + command = [ + "python3", + "-m", + "collector.fpm_forward", + "--backend", + request.identity.framework, + "--model-path", + request.identity.model, + "--gpu", + request.identity.gpu, + "--fpm-max-gpus", + str(request.search.tensor_parallel), + "--fpm-gpu-counts", + str(request.search.tensor_parallel), + "--fpm-parallel-presets", + "pure_tp" if request.identity.model_kind == "moe" else "tp", + "--fpm-max-prefill-isl", + str(max(2, request.workload.input_tokens * request.workload.concurrency)), + "--fpm-max-prefill-batch-size", + str(request.workload.concurrency), + "--checkpoint-dir", + str(checkpoint), + "--fpm-artifact-root", + str(root / "fpm-artifacts"), + "--fpm-database-root", + str(root / "systems/data"), + ] + if request.identity.sm is not None: + command.extend(("--sm", str(request.identity.sm))) + if deployment is not None: + for name, value in deployment.model_dump(exclude_none=True).items(): + if name == "image": + command.extend(("--generator-set", f"K8sConfig.k8s_image={json.dumps(value)}")) + else: + command.extend(("--" + name.replace("_", "-"), value)) + if plan_only: + command.append("--plan-only") + if smoke: + command.append("--smoke") + if limit is not None: + command.extend(("--limit", str(limit))) + if resume: + command.append("--resume") + return command + + +def _check_campaign_outputs(root: Path, *, smoke: bool, resume: bool, checkpoint_dir: str | Path | None) -> None: + checkpoint_name = "fpm_forward_smoke.json" if smoke else "fpm_forward.json" + checkpoint_root = root / "fpm-checkpoint" + artifact_paths = (root / "fpm-artifacts").glob("*/smoke/*" if smoke else "*/*") + occupied = ( + any(smoke or path.name != "smoke" for path in artifact_paths) + or any(checkpoint_root.rglob(checkpoint_name)) + or (not smoke and any((root / "systems/data").glob("*"))) + ) + if not occupied: + return + if not resume: + raise ValueError( + "existing campaign outputs require --resume with a matching checkpoint; " + "choose a new output directory otherwise" + ) + selected = Path(checkpoint_dir).expanduser().resolve() if checkpoint_dir else checkpoint_root + checkpoint = selected / checkpoint_name + try: + if checkpoint.is_symlink(): + raise ValueError("symlinked checkpoint") + payload = json.loads(checkpoint.read_text(encoding="utf-8")) + if ( + not isinstance(payload, dict) + or not isinstance(payload.get("schema"), str) + or not payload["schema"] + or not isinstance(payload.get("plan_sha256"), str) + or len(payload["plan_sha256"]) != 64 + or not isinstance(payload.get("cells"), dict) + ): + raise ValueError("invalid checkpoint document") + except (OSError, ValueError) as exc: + raise ValueError( + f"existing campaign outputs require a readable matching checkpoint at {checkpoint}; " + "choose a new output directory otherwise" + ) from exc + # The collector remains responsible for the checkpoint schema and frozen-plan identity. + + +def run_fpm( + request: SupportRequest, + *, + output_dir: str | Path, + execute: bool = False, + smoke: bool = False, + limit: int | None = None, + resume: bool = False, + checkpoint_dir: str | Path | None = None, + deployment: FPMDeployment | None = None, +) -> int: + """Preview without side effects; execution requires the matching saved plan.""" + + from .plan import check_plan, plan_lock + + root = Path(output_dir).expanduser().resolve() + command = fpm_cli_args( + request, + output_dir=root, + plan_only=not execute, + smoke=smoke, + limit=limit, + resume=resume, + checkpoint_dir=checkpoint_dir, + deployment=deployment, + ) + if not execute: + if resume or (root.exists() and any(root.iterdir())): + check_plan(request, root) + print(shlex.join(command)) + return 0 + check_plan(request, root) + with plan_lock(root): + check_plan(request, root) + _check_campaign_outputs(root, smoke=smoke, resume=resume, checkpoint_dir=checkpoint_dir) + # The collector reports input/plan failures through argparse before + # entering run_resolved; imports and execution failures return exit 1. + try: + from collector.fpm_forward.cli import main as fpm_main + + result = fpm_main(command[3:]) + if result == 0 and not smoke: + selected = Path(checkpoint_dir).expanduser().resolve() if checkpoint_dir else root / "fpm-checkpoint" + checkpoint = json.loads((selected / "fpm_forward.json").read_text(encoding="utf-8")) + publication = checkpoint.get("database") + if not isinstance(publication, dict) or publication.get("status") != "passed": + raise RuntimeError("collector did not record a completed FPM database publication") + # The collector validates this commit record, including on resume + # after raw-artifact reclamation. Check its actual runtime version, + # not an unrelated directory left by an earlier collection. + metadata = json.loads(Path(publication["metadata"]).read_text(encoding="utf-8")) + observed = metadata.get("backend_version") + expected = request.identity.framework_version + if observed != expected: + raise RuntimeError( + f"pod-reported {request.identity.framework} version {observed!r} does not match " + f"framework_version {expected!r}; generated configs cannot use this publication. " + "Collected artifacts are preserved. Create a new plan in a new output directory " + "using a matching runtime or the observed framework_version." + ) + return result + except Exception as exc: + print(f"aisimulate onboard collect-fpm failed: {exc}", file=sys.stderr) + return 1 diff --git a/python/aisimulate/src/aisimulate/support/plan.py b/python/aisimulate/src/aisimulate/support/plan.py new file mode 100644 index 000000000..46b08c406 --- /dev/null +++ b/python/aisimulate/src/aisimulate/support/plan.py @@ -0,0 +1,340 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Create ordinary prediction/recommendation inputs without resolving a model.""" + +from __future__ import annotations + +import fcntl +import hashlib +import json +import os +import stat +from collections.abc import Iterator +from contextlib import contextmanager +from importlib.resources import files +from importlib.util import find_spec, module_from_spec +from pathlib import Path +from typing import Any + +import yaml + +from aisimulate.config import CorePredictionConfig, CoreRecommendationConfig + +from .fpm import fpm_cli_args +from .schema import SupportRequest + +_LOCK_NAME = ".support.lock" + + +def request_id(request: SupportRequest) -> str: + """Bind the entire declared request, including workload and search settings.""" + + payload = json.dumps(request.model_dump(mode="json", exclude_none=True), sort_keys=True, separators=(",", ":")) + return "onboarding-" + hashlib.sha256(payload.encode()).hexdigest() + + +def _parallelism(request: SupportRequest, replicas: int) -> dict[str, int]: + return { + "replicas": replicas, + "tensor": request.search.tensor_parallel, + "pipeline": 1, + "attention_data": 1, + "moe_tensor": request.search.tensor_parallel if request.identity.model_kind == "moe" else 1, + "moe_expert": 1, + } + + +def _configs( + request: SupportRequest, root: Path +) -> tuple[dict[str, Any], dict[str, dict[str, Any]], list[dict[str, int]]]: + max_replicas = request.identity.node_count * (request.identity.gpus_per_node // request.search.tensor_parallel) + replicas = list(dict.fromkeys((1, max_replicas)))[: request.search.max_candidates] + presets = [_parallelism(request, count) for count in replicas] + workload = request.workload + common = { + "traffic": { + "source": { + "type": "synthetic", + "input_tokens": workload.input_tokens, + "output_tokens": workload.output_tokens, + }, + "load": {"type": "concurrency", "concurrency": workload.concurrency}, + "stop": {"requests": workload.request_count}, + }, + "evaluation": {"sla": {"ttft_ms": workload.slo.ttft_ms, "itl_ms": workload.slo.tpot_ms}}, + } + engine = { + "mode": "aggregated", + "model": request.identity.model, + "hardware": request.identity.gpu, + "backend": request.identity.framework, + "backend_version": request.identity.framework_version, + "systems_paths": [str(root / "systems")], + "context_length": request.search.context_length, + } + # Recommendation otherwise expands scheduler defaults into extra domains. + worker = { + "scheduler": {"max_batched_tokens": 8192, "max_sequences": 256}, + "timing": {"type": "default", "estimation_mode": "fpm_interpolation", "fallback_policy": "deny"}, + } + prediction = { + **common, + "engine": {**engine, "workers": {"aggregated": {**worker, "parallelism": presets[0]}}}, + } + CorePredictionConfig.model_validate(prediction) + recommendations = {} + for preset in presets: + recommendation = { + **common, + "engine": {**engine, "workers": {"aggregated": {**worker, "parallelism": {"preset": [preset]}}}}, + "optimization": { + "target": request.search.objective, + "constraints": {"max_candidate_gpus": request.identity.gpu_count}, + }, + "optimizer": {"algorithm": "random", "max_trials": 1, "parallelism": 1, "seed": request.search.seed}, + } + CoreRecommendationConfig.model_validate(recommendation) + name = "pilot" if preset["replicas"] == 1 else f"replicas-{preset['replicas']}" + recommendations[name] = recommendation + return prediction, recommendations, presets + + +def _commands(request: SupportRequest, root: Path, recommendation_names: list[str]) -> dict[str, Any]: + return { + "fpm_plan_local": fpm_cli_args(request, output_dir=root, plan_only=True), + "fpm_run_local": [ + "aisimulate", + "onboard", + "collect-fpm", + "--config", + str(root / "request.yaml"), + "--output-dir", + str(root), + "--execute", + ], + "predict": [ + [ + "aisimulate", + "predict", + "--config", + str(root / "predict/pilot.yaml"), + "--output-dir", + str(root / "predict-results/pilot"), + "--format", + "json", + ] + ], + "recommend": [ + [ + "aisimulate", + "recommend", + "--config", + str(root / f"recommend/{name}.yaml"), + "--output-dir", + str(root / f"recommend-results/{name}"), + "--format", + "json", + ] + for name in recommendation_names + ], + } + + +def _plan_documents(request: SupportRequest, root: Path) -> tuple[dict[str, Any], dict[Path, bytes]]: + package_spec = find_spec("aisimulate_core") + if package_spec is None: + raise ValueError("aisimulate_core package resources are required for packaged system specifications") + # Consult the resource loader without executing the package's native imports. + source_spec = files(module_from_spec(package_spec)) / "systems" / f"{request.identity.gpu}.yaml" + if not source_spec.is_file(): + raise ValueError( + f"GPU {request.identity.gpu!r} has no packaged system specification; " + "new GPU architectures require integration" + ) + prediction, recommendations, presets = _configs(request, root) + commands = _commands(request, root, list(recommendations)) + plan = { + "schema_version": "aisimulate-support-plan/v1", + "request_id": request_id(request), + "search": { + "candidate_count": len(presets), + "max_candidates": request.search.max_candidates, + "candidates": [ + {"parallelism": preset, "total_gpus": preset["replicas"] * preset["tensor"]} for preset in presets + ], + "detail": ( + "A single selected worker is planned; recommendation does not compare alternative configurations." + if len(presets) == 1 + else "Run one independent one-trial recommendation for the selected worker and one for the maximum " + "same-worker replicas that fit within each node. Compare their separate results; no combined ranking " + "or additional topology search is generated." + ), + "baseline_rule": "Selected single worker; model parallelism legality and memory fit remain unchecked.", + }, + "fpm": { + "status": "planned", + "worker_gpus": request.search.tensor_parallel, + "plan_command": commands["fpm_plan_local"], + "sampling": ( + "Prefill is bounded by max(2, input_tokens * concurrency) and concurrency. Decode sampling and " + "runtime context auto-fitting use the collector's existing profile. The synthetic request count " + "does not bound timing samples. Smoke and limited runs do not publish formal FPM data." + ), + }, + "prerequisites": [ + { + "id": "model_integration", + "status": "not_checked", + "detail": ( + "Verify pinned model metadata, memory/cache accounting, and chosen parallelism. This revision's " + "FPM predict/recommend uses registered analytical classes and SOL transfer; follow " + "python/aisimulate/docs/add_a_new_model.md for that route. Class-independent FPM with direct " + "timing interpolation is planned separately. Per-operation silicon data is not required." + ), + }, + { + "id": "runtime", + "status": "not_checked", + "detail": ( + "Prepare and verify the declared model/tokenizer revisions, vLLM version, visible GPUs, " + "model access, and the packaged Dynamo/Kubernetes/Generator collector runtime. Local invocation " + "uses that existing runtime. The collector does not apply revision/version declarations; " + "use a pinned local model snapshot where needed." + ), + }, + { + "id": "fpm_data", + "status": "not_checked", + "detail": "Collect matching prefill/decode timings into the local systems tree before prediction or " + "recommendation. Creating a plan establishes neither timing coverage nor memory fit.", + }, + ], + "accuracy": { + "status": "not_assessed", + "detail": "No measured accuracy or E2E validation is claimed by this plan.", + }, + "outputs": { + "request": str(root / "request.yaml"), + "prediction_configs": [str(root / "predict/pilot.yaml")], + "recommendation_configs": [str(root / f"recommend/{name}.yaml") for name in recommendations], + "recommendation_results": [str(root / f"recommend-results/{name}") for name in recommendations], + "commands": str(root / "commands.json"), + "systems_root": str(root / "systems"), + }, + } + yaml_documents = { + Path("request.yaml"): request.model_dump(mode="json", exclude_none=True), + Path("predict/pilot.yaml"): prediction, + **{Path(f"recommend/{name}.yaml"): config for name, config in recommendations.items()}, + } + documents = {path: yaml.safe_dump(data, sort_keys=False).encode() for path, data in yaml_documents.items()} + documents[Path("systems") / source_spec.name] = source_spec.read_bytes() + documents[Path("commands.json")] = (json.dumps(commands, indent=2, sort_keys=True) + "\n").encode() + documents[Path("support-plan.json")] = (json.dumps(plan, indent=2, sort_keys=True) + "\n").encode() + return plan, documents + + +@contextmanager +def plan_lock(root: Path) -> Iterator[None]: + """Keep one stable lock inode; the OS releases ownership on process exit.""" + + lock = root / _LOCK_NAME + descriptor = os.open(lock, os.O_CREAT | os.O_RDWR | os.O_NOFOLLOW, 0o600) + try: + if not stat.S_ISREG(os.fstat(descriptor).st_mode): + raise ValueError(f"onboarding lock must be a regular file: {lock}") + try: + fcntl.flock(descriptor, fcntl.LOCK_EX | fcntl.LOCK_NB) + except BlockingIOError as exc: + raise ValueError(f"another onboarding operation holds {lock}") from exc + yield + finally: + os.close(descriptor) + + +def _check_paths(root: Path) -> None: + for relative in ( + "request.yaml", + "support-plan.json", + "commands.json", + "predict", + "predict/pilot.yaml", + "recommend", + "recommend/pilot.yaml", + "systems", + "systems/data", + "fpm-checkpoint", + "fpm-artifacts", + ): + if (root / relative).is_symlink(): + raise ValueError(f"refusing symlinked plan output {root / relative}") + pending = [root / relative for relative in ("systems/data", "fpm-checkpoint", "fpm-artifacts")] + while pending: + path = pending.pop() + if path.is_symlink(): + raise ValueError(f"refusing symlinked plan output {path}") + if path.is_dir(): + pending.extend(path.iterdir()) + + +def check_plan(request: SupportRequest, root: Path) -> None: + """Verify saved identity before collecting or reusing any local output.""" + + _check_paths(root) + try: + prior = json.loads((root / "support-plan.json").read_text(encoding="utf-8")) + except (OSError, ValueError) as exc: + raise ValueError( + f"a readable onboarding plan is required in {root}; run aisimulate onboard plan first" + ) from exc + if not isinstance(prior, dict) or prior.get("request_id") != request_id(request): + raise ValueError(f"refusing to mix a different request identity in {root}; choose a new output directory") + _check_saved_request(request, root) + + +def _check_saved_request(request: SupportRequest, root: Path) -> None: + try: + saved = SupportRequest.from_yaml(root / "request.yaml") + except ValueError as exc: + raise ValueError(f"saved request in {root} was modified or cannot be read") from exc + if saved != request: + raise ValueError(f"saved request identity in {root} differs from the requested plan") + + +def create_plan(request: SupportRequest, output_dir: str | Path, *, overwrite: bool = False) -> dict[str, Any]: + """Write a new plan, or repair missing files for an identical saved request.""" + + root = Path(output_dir).expanduser().resolve() + plan, documents = _plan_documents(request, root) + root.mkdir(parents=True, exist_ok=True) + with plan_lock(root): + _check_paths(root) + nonempty = any(path.name != _LOCK_NAME for path in root.iterdir()) + if nonempty: + if not overwrite: + raise ValueError(f"output directory {root} is nonempty; use overwrite only for the same request") + if (root / "support-plan.json").exists(): + check_plan(request, root) + else: + # request.yaml is written first, so it can establish identity + # after an interruption before the completed plan was saved. + _check_saved_request(request, root) + # Inspect every generated file before writing any file. Existing data, + # checkpoints, results, and edited inputs are never replaced or deleted. + pending: dict[Path, bytes] = {} + for relative, content in documents.items(): + destination = root / relative + if destination.is_symlink(): + raise ValueError(f"refusing symlinked plan output {destination}") + if not destination.exists(): + pending[relative] = content + elif destination.read_bytes() != content: + raise ValueError(f"generated plan input {destination} was modified; choose a new output directory") + for relative, content in pending.items(): + destination = root / relative + destination.parent.mkdir(parents=True, exist_ok=True) + with destination.open("xb") as handle: + handle.write(content) + (root / "systems/data").mkdir(parents=True, exist_ok=True) + return plan diff --git a/python/aisimulate/src/aisimulate/support/schema.py b/python/aisimulate/src/aisimulate/support/schema.py new file mode 100644 index 000000000..603f6356c --- /dev/null +++ b/python/aisimulate/src/aisimulate/support/schema.py @@ -0,0 +1,166 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Declared onboarding inputs; validity does not establish simulation readiness.""" + +from __future__ import annotations + +import re +from pathlib import Path, PurePosixPath +from typing import Literal + +from packaging.version import Version +from pydantic import Field, field_validator, model_validator + +from aisimulate.config.common import PositiveFiniteFloat, PositiveStrictInt, StrictModel, load_yaml + + +class SupportIdentity(StrictModel): + model: str + model_revision: str + model_kind: Literal["dense", "moe"] + framework: Literal["vllm"] = "vllm" + framework_version: str + gpu: str + gpu_count: PositiveStrictInt + node_count: PositiveStrictInt = 1 + gpus_per_node: PositiveStrictInt + interconnect: str + sm: PositiveStrictInt | None = None + tokenizer_revision: str | None = None + chat_template_revision: str | None = None + aisimulate_revision: str | None = None + + @field_validator( + "model", + "model_revision", + "framework_version", + "gpu", + "interconnect", + "tokenizer_revision", + "chat_template_revision", + "aisimulate_revision", + ) + @classmethod + def _nonempty(cls, value: str | None) -> str | None: + if value is None: + return None + value = value.strip() + if not value or "\x00" in value: + raise ValueError("identity values must be nonempty and contain no NUL characters") + return value + + @field_validator( + "model_revision", "framework_version", "tokenizer_revision", "chat_template_revision", "aisimulate_revision" + ) + @classmethod + def _pinned_revision(cls, value: str | None) -> str | None: + if value is not None and value.lower() in {"main", "master", "latest", "head"}: + raise ValueError("declare a pinned revision or version rather than a mutable default") + return value + + @field_validator("gpu") + @classmethod + def _system_name(cls, value: str) -> str: + if not re.fullmatch(r"[A-Za-z0-9][A-Za-z0-9_-]*", value): + raise ValueError("gpu must be a packaged system name, not a path") + return value + + @model_validator(mode="after") + def _allocation(self) -> SupportIdentity: + if self.node_count * self.gpus_per_node != self.gpu_count: + raise ValueError("node_count * gpus_per_node must equal gpu_count") + return self + + +class SloSpec(StrictModel): + ttft_ms: PositiveFiniteFloat = 1000.0 + tpot_ms: PositiveFiniteFloat = 100.0 + + +class WorkloadSpec(StrictModel): + """One fixed synthetic shape; request_count is not a timing sample count.""" + + input_tokens: PositiveStrictInt = 1024 + output_tokens: PositiveStrictInt = 128 + concurrency: PositiveStrictInt = 1 + request_count: PositiveStrictInt = 4 + slo: SloSpec = Field(default_factory=SloSpec) + + @model_validator(mode="after") + def _request_count(self) -> WorkloadSpec: + if self.request_count < self.concurrency: + raise ValueError("request_count must be at least concurrency") + return self + + +class SearchProfile(StrictModel): + tensor_parallel: PositiveStrictInt = 1 + context_length: PositiveStrictInt = 16384 + max_candidates: int = Field(default=1, strict=True, ge=1, le=2) + objective: Literal[ + "throughput", + "throughput_per_gpu", + "throughput_per_user", + "goodput", + "goodput_per_gpu", + "ttft", + "e2e_latency", + "pareto", + ] = "throughput" + seed: int = Field(default=42, strict=True, ge=0) + + +_DNS_LABEL = r"[a-z0-9](?:[-a-z0-9]{0,61}[a-z0-9])?" +_DNS_SUBDOMAIN = rf"{_DNS_LABEL}(?:\.{_DNS_LABEL})*" + + +class FPMDeployment(StrictModel): + """Deployment-only inputs included in the collector's frozen-plan identity.""" + + dynamo_version: str | None = None + image: str | None = Field(default=None, pattern=r"^[^\s\x00]+$") + namespace: str | None = Field(default=None, pattern=rf"^{_DNS_LABEL}$") + model_cache: str | None = None + transport: Literal["nvlink", "ib", "efa"] | None = None + image_pull_secret: str | None = Field(default=None, pattern=rf"^{_DNS_SUBDOMAIN}$", max_length=253) + + @field_validator("dynamo_version") + @classmethod + def _release_version(cls, value: str | None) -> str | None: + if value is not None: + Version(value) + return value + + @field_validator("model_cache") + @classmethod + def _model_cache(cls, value: str | None) -> str | None: + if value is None: + return None + parts = value.split(":") + if len(parts) > 3 or len(parts[0]) > 253 or not re.fullmatch(_DNS_SUBDOMAIN, parts[0]): + raise ValueError("model_cache must be NAME[:MOUNT[:SUBPATH]] with a valid PVC name") + if len(parts) > 1 and parts[1] and not PurePosixPath(parts[1]).is_absolute(): + raise ValueError("model_cache MOUNT must be an absolute container path") + if "\x00" in value or "\n" in value or "\r" in value: + raise ValueError("model_cache must contain no NUL characters or newlines") + return value + + +class SupportRequest(StrictModel): + schema_version: Literal["aisimulate-support-request/v1"] = "aisimulate-support-request/v1" + identity: SupportIdentity + workload: WorkloadSpec = Field(default_factory=WorkloadSpec) + search: SearchProfile = Field(default_factory=SearchProfile) + + @model_validator(mode="after") + def _shape(self) -> SupportRequest: + if self.search.tensor_parallel > self.identity.gpus_per_node: + raise ValueError("search.tensor_parallel must fit within gpus_per_node") + if self.workload.input_tokens + self.workload.output_tokens > self.search.context_length: + raise ValueError("search.context_length must cover the input and output tokens") + return self + + @classmethod + def from_yaml(cls, path: str | Path) -> SupportRequest: + return cls.model_validate(load_yaml(path)) diff --git a/python/aisimulate/src/aisimulate/sweeper/config.py b/python/aisimulate/src/aisimulate/sweeper/config.py index ce3e82227..30865d343 100644 --- a/python/aisimulate/src/aisimulate/sweeper/config.py +++ b/python/aisimulate/src/aisimulate/sweeper/config.py @@ -29,7 +29,7 @@ from aisimulate.config.traffic import AgenticSnapshotOptions -from ..config.common import ENGINE_MODEL_CONTROL_FIELDS, is_active_engine_model_control +from ..config.common import ENGINE_MODEL_CONTROL_FIELDS, SystemsPath, SystemsRoot, is_active_engine_model_control from ..config.engine import NgramSpeculationConfig @@ -525,7 +525,8 @@ class SearchSpace(BaseModel): hardware_sku: str # e.g. "h200_sxm" database_mode: Literal["SILICON", "HYBRID", "EMPIRICAL", "SOL"] = "SILICON" transfer_policy: str | list[str] | None = None - systems_paths: list[str] | None = Field(default=None, min_length=1) + systems_paths: list[SystemsRoot] | None = Field(default=None, min_length=1) + systems_path: SystemsPath | None = Field(default=None, exclude=True) estimation_mode: Literal["auto", "op_level", "fpm_interpolation", "fpm_regression"] = "auto" fallback_policy: Literal["deny", "allow"] = "deny" estimator_config: dict[str, Any] = Field(default_factory=dict) @@ -650,6 +651,10 @@ def _migrate_legacy_estimator_selection(cls, value): @model_validator(mode="after") def _validate_search_choices(self) -> SearchSpace: """Every backend dimension is a non-empty subset of its allowed choices.""" + if self.systems_path is not None: + if self.systems_paths is not None and self.systems_paths != [self.systems_path]: + raise ValueError("systems_path conflicts with systems_paths") + self.systems_paths = [self.systems_path] if self.speculation is not None: if self.aic_nextn is not None: raise ValueError("speculation cannot be combined with aic_nextn") @@ -831,15 +836,6 @@ def _validate_moe_kernel_source(cls, value): return normalize_kernel_source(value, "moe_kernel_source") - @field_validator("systems_paths") - @classmethod - def _validate_estimator_roots(cls, value): - if value is None: - return value - if any(not path.strip() for path in value): - raise ValueError("systems_paths entries must be nonempty") - return value - @model_validator(mode="after") def _validate_estimator_controls(self): from aisimulate_core.sdk.common import resolve_transfer_policy @@ -892,7 +888,6 @@ def _validate_estimator_controls(self): nondefault = ( self.database_mode != "SILICON" or self.transfer_policy is not None - or self.systems_paths not in (None, ["default"]) or self.estimation_mode != "auto" or self.fallback_policy != "deny" or bool(self.estimator_config) @@ -900,6 +895,9 @@ def _validate_estimator_controls(self): roles = ({"agg"} if "agg" in self.deployment_mode else set()) | ( {"prefill", "decode"} if "disagg" in self.deployment_mode else set() ) + if self.systems_paths not in (None, ["default"]) and self._uses_legacy_estimator_provider(): + unsupported = "AFD" if set(self.deployment_mode) & {"afd", "afd+pd"} else "analytical encoder pools" + raise ValueError(f"systems_paths does not support {unsupported}") if nondefault and ( set(self.deployment_mode) & {"afd", "afd+pd"} or self.encoder is not None diff --git a/python/aisimulate/src/aisimulate/sweeper/deploy.py b/python/aisimulate/src/aisimulate/sweeper/deploy.py index d92f28a15..129376e9e 100644 --- a/python/aisimulate/src/aisimulate/sweeper/deploy.py +++ b/python/aisimulate/src/aisimulate/sweeper/deploy.py @@ -52,6 +52,8 @@ def _performance_model_metadata(sample: dict[str, Any], role: str, *, backend_ve config["fpm_parquet_path"] = sample[f"{role}_fpm_parquet_path"] if sample.get("speculation") is not None: config["speculation"] = NgramSpeculationConfig.model_validate(sample["speculation"]).cost_config() + if sample.get("systems_paths") is not None: + config["systems_paths"] = sample["systems_paths"] return {"provider": "aic", "config": config} @@ -99,6 +101,10 @@ def _engine_args_payload( memory_fraction_field: float(memory_fraction), "enable_prefix_caching": bool(sample[f"{role}_enable_prefix_caching"]), } + if sample.get("systems_paths") is not None and sample.get(f"{role}_timing_model") is not None: + from .forward_pass_estimator import resolve_systems_paths + + payload["systems_path"] = list(resolve_systems_paths(sample["systems_paths"])) if sample.get("context_length") is not None: payload["max_model_len"] = int(sample["context_length"]) if moe_tp * moe_ep > 1: @@ -126,6 +132,7 @@ def _engine_args_payload( if sample.get(f"{role}_num_gpu_blocks") is None: payload = materialize_aic_num_gpu_blocks(payload) for name in ( + "systems_path", "aic_backend_version", "aic_system", "aic_model_path", diff --git a/python/aisimulate/src/aisimulate/sweeper/kv_load.py b/python/aisimulate/src/aisimulate/sweeper/kv_load.py index 40eae1004..f6a9cb259 100644 --- a/python/aisimulate/src/aisimulate/sweeper/kv_load.py +++ b/python/aisimulate/src/aisimulate/sweeper/kv_load.py @@ -92,6 +92,8 @@ def _role_capacity_tokens( roots = resolved.get("systems_paths") if not roots and resolved.get("systems_path"): roots = [resolved["systems_path"]] + if not roots: + roots = sample.get("systems_paths") per_rank_tokens = _per_rank_capacity_tokens( config.shape, model_name=str(resolved.get("model", resolved.get("model_path", sample["model_name"]))), diff --git a/python/aisimulate/src/aisimulate/sweeper/sample.py b/python/aisimulate/src/aisimulate/sweeper/sample.py index 4d1d48f59..a4dd74aa4 100644 --- a/python/aisimulate/src/aisimulate/sweeper/sample.py +++ b/python/aisimulate/src/aisimulate/sweeper/sample.py @@ -136,6 +136,8 @@ def unroll_sample( for key in _DEPLOYMENT_PINNED: sample[key] = getattr(search_space, key) + if search_space.systems_paths is not None: + sample["systems_paths"] = search_space.systems_paths if mode == "disagg": sample["prefill_hardware_sku"] = search_space.hardware_sku_for("prefill") diff --git a/python/aisimulate/src/aisimulate/sweeper/search.py b/python/aisimulate/src/aisimulate/sweeper/search.py index 364b74028..e726bd306 100644 --- a/python/aisimulate/src/aisimulate/sweeper/search.py +++ b/python/aisimulate/src/aisimulate/sweeper/search.py @@ -616,10 +616,18 @@ def _materialize_one( if sample["deployment_mode"] == "disagg": role_hardware = {role: sample[f"{role}_hardware_sku"] for role in ("prefill", "decode")} if len(set(role_hardware.values())) == 1: - backend_version = resolve_backend_version(role_hardware["prefill"], selection["backend"]) + backend_version = resolve_backend_version( + role_hardware["prefill"], + selection["backend"], + systems_paths=config.search_space.systems_paths_for("prefill"), + ) else: role_versions = { - role: resolve_backend_version(hardware, selection["backend"]) + role: resolve_backend_version( + hardware, + selection["backend"], + systems_paths=config.search_space.systems_paths_for(role), + ) for role, hardware in role_hardware.items() } if len(set(role_versions.values())) != 1: @@ -630,7 +638,11 @@ def _materialize_one( ) backend_version = role_versions["prefill"] else: - backend_version = resolve_backend_version(config.search_space.hardware_sku, selection["backend"]) + backend_version = resolve_backend_version( + config.search_space.hardware_sku, + selection["backend"], + systems_paths=config.search_space.systems_paths_for("agg"), + ) # The resolved perf-model version is part of the evaluated contract. Keep it # on the candidate so downstream artifact generation cannot independently # select a different backend version. diff --git a/python/aisimulate/src/aisimulate_core/sdk/engine.py b/python/aisimulate/src/aisimulate_core/sdk/engine.py index a898b3824..8e95cae61 100644 --- a/python/aisimulate/src/aisimulate_core/sdk/engine.py +++ b/python/aisimulate/src/aisimulate_core/sdk/engine.py @@ -276,8 +276,8 @@ def _literal_backend_version( (an omitted version means the ``current`` slot). Slot-policy errors (unlisted versions, unpopulated aliases) PROPAGATE — the spec builder is a user-level surface and must not smuggle ungated coordinates onto the - wire. Trees without a slots file (synthetic/external) keep the ungated - passthrough. + wire. Trees without a slots file (synthetic/external) keep explicit + versions unchanged and resolve an omitted version to the latest database. """ resolved = getattr(database, "version", None) if database is not None else None if resolved: @@ -286,7 +286,11 @@ def _literal_backend_version( slots = perf_database.get_version_slots(system, backend, systems_paths=systems_path) if slots is None: - return backend_version + return ( + backend_version + if backend_version is not None + else perf_database.get_latest_database_version(system, backend, systems_paths=systems_path) + ) requested = "current" if backend_version is None else backend_version return perf_database.resolve_query_version(system, backend, requested, systems_paths=systems_path) diff --git a/python/aisimulate/tests/unit/support_matrix/test_script_imports.py b/python/aisimulate/tests/unit/support_matrix/test_script_imports.py new file mode 100644 index 000000000..a2a4d8abb --- /dev/null +++ b/python/aisimulate/tests/unit/support_matrix/test_script_imports.py @@ -0,0 +1,38 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Importing support-matrix scripts must preserve the caller's package search path.""" + +import subprocess +import sys +from pathlib import Path + +import pytest + +pytestmark = pytest.mark.unit + + +@pytest.mark.parametrize("script", ["compare_support_matrix", "generate_support_matrix"]) +def test_script_import_preserves_package_search_path(script): + result = subprocess.run( + [ + sys.executable, + "-c", + """ +import importlib +import sys + +original_path = list(sys.path) +importlib.import_module(sys.argv[1]) +assert sys.path == original_path, (original_path, sys.path) +""", + f"tools.support_matrix.{script}", + ], + cwd=Path(__file__).resolve().parents[3], + capture_output=True, + text=True, + timeout=30, + check=False, + ) + + assert result.returncode == 0, result.stdout + result.stderr diff --git a/python/aisimulate/tests/unit/test_support_cli.py b/python/aisimulate/tests/unit/test_support_cli.py new file mode 100644 index 000000000..4333cb712 --- /dev/null +++ b/python/aisimulate/tests/unit/test_support_cli.py @@ -0,0 +1,955 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Public onboarding CLI behavior without a model download or GPU collection.""" + +from __future__ import annotations + +import builtins +import json +import os +import shlex +import shutil +import subprocess +import sys +from pathlib import Path +from types import SimpleNamespace + +import pytest + +import aisimulate.main as cli +from aisimulate.config import CorePredictionConfig, CoreRecommendationConfig +from aisimulate.support.schema import SupportRequest + +pytestmark = pytest.mark.unit + +_REQUIRED = { + "model": "example/unintegrated-model", + "model_revision": "revision-123", + "model_kind": "dense", + "framework_version": "0.24.0", + "gpu": "h200_sxm", + "gpu_count": 4, + "interconnect": "nvswitch", +} +_PILOT = { + "tensor_parallel": 2, + "input_tokens": 1024, + "output_tokens": 128, + "concurrency": 1, + "context_length": 16384, + "ttft_ms": 1000, + "tpot_ms": 100, +} + + +def _init_args(output: Path, *, full: bool = False, **changes) -> list[str]: + options = {**_REQUIRED, **(_PILOT if full else {}), **changes} + return ["onboard", "init", "--output", str(output)] + [ + part + for name, value in options.items() + if value is not None + for part in ("--" + name.replace("_", "-"), str(value)) + ] + + +def _terminal(monkeypatch, answers=()) -> list[str]: + prompts = [] + remaining = iter(answers) + monkeypatch.setattr(sys.stdin, "isatty", lambda: True) + + def respond(prompt: str) -> str: + prompts.append(prompt) + try: + answer = next(remaining) + except StopIteration: + pytest.fail(f"unexpected prompt: {prompt}") + if isinstance(answer, BaseException): + raise answer + return answer + + monkeypatch.setattr("builtins.input", respond) + return prompts + + +def test_supervised_entrypoint_runs_onboarding_before_prediction_config_handling(tmp_path, monkeypatch): + from aisimulate import supervision + + def unexpected_supervision(*args, **kwargs): + pytest.fail("interactive onboarding must retain the calling terminal") + + monkeypatch.setattr(supervision, "run_process", unexpected_supervision) + output = tmp_path / "request.yaml" + assert supervision.main(_init_args(output)) == 0 + assert SupportRequest.from_yaml(output).identity.model == _REQUIRED["model"] + + +@pytest.mark.parametrize("command", [[], ["onboard"]]) +def test_help_explains_model_fpm_and_target_hardware_scope(capsys, command) -> None: + with pytest.raises(SystemExit) as result: + cli.main([*command, "--help"]) + + assert result.value.code == 0 + output = " ".join(capsys.readouterr().out.split()) + assert "Onboard a model for FPM simulation on a target hardware platform." in output + assert "aisimulate" in output + + +@pytest.mark.parametrize("action", ["init", "plan", "collect-fpm"]) +def test_retired_support_command_is_rejected_before_dispatch(tmp_path, monkeypatch, capsys, action) -> None: + def unexpected_dispatch(*args, **kwargs): + pytest.fail("the retired command must be rejected before setup or collector dispatch") + + monkeypatch.setattr(cli, "run_support_command", unexpected_dispatch) + monkeypatch.setattr(cli, "resolve_runner_factory", unexpected_dispatch) + request = tmp_path / "request.yaml" + if action == "init": + args = _init_args(request)[1:] + else: + args = [action, "--config", str(request), "--output-dir", str(tmp_path / "plan")] + if action == "collect-fpm": + args.append("--execute") + + with pytest.raises(SystemExit) as result: + cli.main(["support", *args]) + + assert result.value.code == 2 + assert "invalid choice: 'support'" in capsys.readouterr().err + assert not list(tmp_path.iterdir()) + + +def test_guided_and_scripted_setup_produce_the_same_request(tmp_path, monkeypatch, capsys) -> None: + guided = tmp_path / "guided request.yaml" + scripted = tmp_path / "scripted.yaml" + prompts = _terminal( + monkeypatch, + ["example/unintegrated-model", "revision-123", "dense", "0.24.0", "h200_sxm", "4", "nvswitch", "2"] + [""] * 6, + ) + + assert cli.main(["onboard", "init", "--interactive", "--output", str(guided)]) == 0 + assert cli.main(_init_args(scripted, tensor_parallel=2)) == 0 + + assert SupportRequest.from_yaml(guided) == SupportRequest.from_yaml(scripted) + assert any("Input tokens per request [1024]" in prompt for prompt in prompts) + assert any("Target time to first token (ms) [1000.0]" in prompt for prompt in prompts) + request = SupportRequest.from_yaml(guided) + assert request.workload.request_count == 4 + assert request.identity.gpus_per_node == 4 + assert request.identity.aisimulate_revision is None + assert request.identity.tokenizer_revision is None + output = capsys.readouterr().out + assert "accuracy is not assessed" in output + assert "are unchecked" in output + + +def test_supplied_options_skip_prompts_and_onboarding_alias_is_optional(tmp_path, monkeypatch) -> None: + prompts = _terminal(monkeypatch) + guided = tmp_path / "guided.yaml" + scripted = tmp_path / "scripted.yaml" + options = {"request_count": 8, "seed": 123, "max_candidates": 2, "objective": "goodput"} + + assert cli.main(_init_args(guided, full=True, **options) + ["--interactive", "--profile", "onboarding"]) == 0 + assert cli.main(_init_args(scripted, full=True, **options)) == 0 + + assert prompts == [] + assert SupportRequest.from_yaml(guided) == SupportRequest.from_yaml(scripted) + + +def test_guided_setup_recovers_numeric_input_and_shared_field_validation(tmp_path, monkeypatch, capsys) -> None: + output = tmp_path / "request.yaml" + prompts = _terminal(monkeypatch, ["many", "4", "revision-123"]) + + assert cli.main(_init_args(output, full=True, model_revision="main", gpu_count=None) + ["--interactive"]) == 0 + + request = SupportRequest.from_yaml(output) + assert request.identity.gpu_count == 4 + assert request.identity.model_revision == "revision-123" + assert len(prompts) == 3 + transcript = capsys.readouterr().out + assert "Enter a valid integer" in transcript + assert "declare a pinned revision" in transcript + + +@pytest.mark.parametrize( + ("changes", "answers", "expected"), + [ + ({"tensor_parallel": 8}, ["unknown-option", "--tensor-parallel", "2"], "tensor_parallel"), + ({"context_length": 100}, ["context-length", "4096"], "context_length"), + ({"concurrency": 8}, ["request-count", "8"], "request_count"), + ({"node_count": 2, "gpus_per_node": 4}, ["gpu-count", "8"], "gpu_count"), + ], +) +def test_guided_setup_recovers_cross_field_validation( + tmp_path, monkeypatch, capsys, changes, answers, expected +) -> None: + output = tmp_path / "request.yaml" + prompts = _terminal(monkeypatch, answers) + + assert cli.main(_init_args(output, full=True, **changes) + ["--interactive"]) == 0 + + SupportRequest.from_yaml(output) + assert any("Option to correct" in prompt for prompt in prompts) + assert expected in capsys.readouterr().out + + +@pytest.mark.parametrize( + "changes", + [ + {"model": None}, + {"model_revision": "latest"}, + {"gpu_count": 0}, + {"ttft_ms": "nan"}, + {"request_count": 0}, + {"context_length": 100}, + {"tensor_parallel": 8}, + {"max_candidates": 3}, + {"objective": "invented"}, + {"node_count": 2}, + ], +) +def test_scripted_setup_uses_shared_validation_and_never_writes_invalid_requests(tmp_path, capsys, changes) -> None: + output = tmp_path / "new" / "request.yaml" + + with pytest.raises(SystemExit) as error: + cli.main(_init_args(output, **changes)) + + assert error.value.code == 2 + assert not output.parent.exists() + assert "error:" in capsys.readouterr().err + + +@pytest.mark.parametrize("interruption", [EOFError, KeyboardInterrupt]) +@pytest.mark.parametrize("existing", [False, True]) +def test_cancelled_setup_preserves_files_and_creates_no_partial_request( + tmp_path, monkeypatch, capsys, interruption, existing +) -> None: + output = tmp_path / "new" / "request.yaml" + if existing: + output.parent.mkdir() + output.write_text("existing contents\n") + _terminal(monkeypatch, ["example/unintegrated-model", interruption()]) + + assert cli.main(["onboard", "init", "--interactive", "--output", str(output), "--overwrite"]) == 130 + + if existing: + assert output.read_text() == "existing contents\n" + assert list(output.parent.iterdir()) == [output] + else: + assert not output.parent.exists() + assert "Setup cancelled" in capsys.readouterr().err + + +def test_interactive_requires_a_terminal_without_reading_input(tmp_path, monkeypatch, capsys) -> None: + _terminal(monkeypatch) + monkeypatch.setattr(sys.stdin, "isatty", lambda: False) + output = tmp_path / "request.yaml" + + with pytest.raises(SystemExit) as error: + cli.main(["onboard", "init", "--interactive", "--output", str(output)]) + + assert error.value.code == 2 + assert "requires a terminal" in capsys.readouterr().err + assert not output.exists() + + +@pytest.mark.parametrize("target_kind", ["existing_file", "directory", "parent_file"]) +def test_invalid_output_is_rejected_before_prompting(tmp_path, monkeypatch, target_kind) -> None: + _terminal(monkeypatch) + target = tmp_path / "request.yaml" + if target_kind == "existing_file": + target.write_text("keep me") + elif target_kind == "directory": + target.mkdir() + else: + parent = tmp_path / "file" + parent.write_text("keep me") + target = parent / "request.yaml" + + with pytest.raises(SystemExit) as error: + cli.main(["onboard", "init", "--interactive", "--output", str(target)]) + + assert error.value.code == 2 + + +def test_output_install_does_not_overwrite_a_file_created_during_setup(tmp_path, monkeypatch) -> None: + output = tmp_path / "request.yaml" + real_link = os.link + + def competing_writer(source, destination) -> None: + Path(destination).write_text("concurrent writer\n") + real_link(source, destination) + + monkeypatch.setattr(os, "link", competing_writer) + + with pytest.raises(SystemExit) as error: + cli.main(_init_args(output)) + + assert error.value.code == 2 + assert output.read_text() == "concurrent writer\n" + assert list(tmp_path.iterdir()) == [output] + + +def test_explicit_overwrite_replaces_a_symlink_without_modifying_its_target(tmp_path) -> None: + existing = tmp_path / "existing.yaml" + existing.write_text("keep me\n") + output = tmp_path / "request.yaml" + output.symlink_to(existing) + + assert cli.main(_init_args(output) + ["--overwrite"]) == 0 + + assert not output.is_symlink() + SupportRequest.from_yaml(output) + assert existing.read_text() == "keep me\n" + + +def test_init_next_command_quotes_the_request_path(tmp_path, capsys) -> None: + output = tmp_path / "request 'with spaces'; $(unused).yaml" + + assert cli.main(_init_args(output)) == 0 + + next_command = next( + line.removeprefix("next: ") for line in capsys.readouterr().out.splitlines() if line.startswith("next: ") + ) + assert shlex.split(next_command) == ["aisimulate", "onboard", "plan", "--config", str(output)] + + +def test_next_command_handles_a_relative_filename_starting_with_a_dash(tmp_path, monkeypatch, capsys) -> None: + monkeypatch.chdir(tmp_path) + args = _init_args(Path("unused.yaml")) + args[2:4] = ["--output=-request.yaml"] + + assert cli.main(args) == 0 + + next_command = next( + line.removeprefix("next: ") for line in capsys.readouterr().out.splitlines() if line.startswith("next: ") + ) + parsed = cli.build_parser().parse_args(shlex.split(next_command)[1:]) + assert Path(parsed.config) == tmp_path / "-request.yaml" + + +@pytest.mark.parametrize("output_dir", ["-plan", "plan 'quoted'; $(unused)"]) +def test_printed_plan_next_command_runs_in_a_shell(tmp_path, monkeypatch, capsys, output_dir) -> None: + monkeypatch.chdir(tmp_path) + request = tmp_path / "request.yaml" + assert cli.main(_init_args(request)) == 0 + capsys.readouterr() + assert cli.main(["onboard", "plan", "-c", str(request), f"--output-dir={output_dir}", "--format", "json"]) == 0 + summary = json.loads(capsys.readouterr().out) + + result = subprocess.run( + ["/bin/sh", "-c", summary["next"]], + cwd=tmp_path, + env={**os.environ, "PATH": str(Path(sys.executable).parent) + os.pathsep + os.environ.get("PATH", "")}, + stdin=subprocess.DEVNULL, + capture_output=True, + text=True, + timeout=30, + ) + + assert result.returncode == 0, result.stderr + assert Path(summary["plan"]) == tmp_path / output_dir / "support-plan.json" + assert "collector.fpm_forward" in result.stdout + assert "--plan-only" in result.stdout + + +@pytest.mark.parametrize("model_kind", ["dense", "moe"]) +def test_init_plan_and_preview_use_real_public_configs_without_launching_collection( + tmp_path, monkeypatch, capsys, model_kind +) -> None: + def unexpected_launch(*args, **kwargs): + pytest.fail("setup, plan, and preview must not launch a collector") + + monkeypatch.setattr(subprocess, "run", unexpected_launch) + monkeypatch.setitem(sys.modules, "collector.fpm_forward.cli", SimpleNamespace(main=unexpected_launch)) + monkeypatch.setattr(cli, "resolve_runner_factory", unexpected_launch) + request_path = tmp_path / "request 'quoted'.yaml" + output = tmp_path / "plan 'quoted'" + assert cli.main(_init_args(request_path, model_kind=model_kind, tensor_parallel=2)) == 0 + capsys.readouterr() + + assert ( + cli.main(["onboard", "plan", "--config", str(request_path), "--output-dir", str(output), "--format", "json"]) + == 0 + ) + summary = json.loads(capsys.readouterr().out) + plan = json.loads((output / "support-plan.json").read_text()) + assert summary["request_id"] == plan["request_id"] + assert summary["candidate_count"] == 1 + prediction = CorePredictionConfig.from_yaml(plan["outputs"]["prediction_configs"][0]) + recommendation = CoreRecommendationConfig.from_yaml(plan["outputs"]["recommendation_configs"][0]) + assert prediction.engine.workers.aggregated.timing.estimation_mode == "fpm_interpolation" + assert recommendation.engine.workers.aggregated.timing.estimation_mode == "fpm_interpolation" + assert prediction.engine.systems_paths == [plan["outputs"]["systems_root"]] + assert recommendation.engine.systems_paths == prediction.engine.systems_paths + assert prediction.traffic.source.input_tokens == 1024 + + next_command = shlex.split(summary["next"]) + assert next_command[:3] == ["aisimulate", "onboard", "collect-fpm"] + assert cli.main(next_command[1:]) == 0 + preview = capsys.readouterr().out + assert "collector.fpm_forward" in preview + assert "--fpm-max-gpus 2 --fpm-gpu-counts 2" in preview + assert "--fpm-parallel-presets " + ("pure_tp" if model_kind == "moe" else "tp") in preview + + +def test_explicit_execute_forwards_diagnostic_options_and_exit_status(tmp_path, monkeypatch) -> None: + calls = [] + + def collector_main(argv): + calls.append(argv) + return 7 + + monkeypatch.setitem(sys.modules, "collector.fpm_forward.cli", SimpleNamespace(main=collector_main)) + request_path = tmp_path / "request.yaml" + output = tmp_path / "plan" + assert cli.main(_init_args(request_path, tensor_parallel=2)) == 0 + assert cli.main(["onboard", "plan", "-c", str(request_path), "--output-dir", str(output)]) == 0 + + assert ( + cli.main( + [ + "onboard", + "collect-fpm", + "-c", + str(request_path), + "--output-dir", + str(output), + "--execute", + "--smoke", + "--limit", + "1", + "--resume", + ] + ) + == 7 + ) + + assert len(calls) == 1 + assert "--plan-only" not in calls[0] + assert "--smoke" in calls[0] + assert "--resume" in calls[0] + assert calls[0][calls[0].index("--limit") + 1] == "1" + + +def _local_collection_command(tmp_path, **changes): + model = tmp_path / "model with spaces" + model.mkdir() + (model / "config.json").write_text( + json.dumps( + { + "architectures": ["LlamaForCausalLM"], + "model_type": "llama", + "num_hidden_layers": 2, + "hidden_size": 128, + "intermediate_size": 256, + "num_attention_heads": 4, + "num_key_value_heads": 2, + "vocab_size": 512, + "max_position_embeddings": 16384, + "torch_dtype": "bfloat16", + } + ) + ) + request_path = tmp_path / "request.yaml" + root = tmp_path / "plan" + assert cli.main(_init_args(request_path, model=model, **changes)) == 0 + assert cli.main(["onboard", "plan", "-c", str(request_path), "--output-dir", str(root)]) == 0 + return ["onboard", "collect-fpm", "-c", str(request_path), "--output-dir", str(root), "--execute"] + + +def _synthetic_collector(monkeypatch, version): + """Exercise real collection, provenance validation, and publication without GPUs.""" + from collector.fpm_forward import planner, runner + + executions = [] + + def render_cell(_plan, _cell, cell_dir, _overrides, **_kwargs): + (cell_dir / "k8s_deploy.yaml").write_text("apiVersion: v1\nkind: Pod\nmetadata:\n name: test-cell\n") + for name in ("run.sh", "fpm_env.sh", "collector-runtime-env.sh"): + (cell_dir / name).write_text("#!/bin/sh\n") + + class LocalResource: + def __init__(self, _manifest, cell_dir): + self.cell = json.loads((cell_dir / "cell.json").read_text()) + self.raw = cell_dir / "raw/pod-0" + + def apply(self): + self.raw.mkdir(parents=True) + + def wait_ready(self, _expected_nodes): + return ["pod-0"] + + def stage(self, _pods, _files): + pass + + def prepare_attempt(self, _pods, **identity): + (self.raw / "collector-provenance.json").write_text( + json.dumps( + { + "schema_name": "aic_fpm_collector_provenance", + "schema_version": 1, + **identity, + "runtime": {"backend": "vllm", "backend_version": version}, + } + ) + ) + + def execute(self, _pods): + executions.append(self.cell["cell_id"]) + phase = self.cell["workload_kind"] + prefill = int(phase == "prefill") + point = { + "point_type": phase, + "benchmark_id": 1, + "batch_size": 1, + "total_prefill_tokens": prefill, + "total_kv_read_tokens": 1, + } + fpm = { + "counter_id": 1, + "dp_rank": 0, + "wall_time": 0.01, + "scheduled_requests": { + "num_prefill_requests": prefill, + "sum_prefill_tokens": prefill, + "sum_prefill_kv_tokens": prefill, + "num_decode_requests": 1 - prefill, + "sum_decode_kv_tokens": 1 - prefill, + }, + } + (self.raw / "benchmark.json").write_text( + json.dumps( + { + "schema_version": 2, + "artifact_type": "rank", + "status": "complete", + "valid": True, + "usable": True, + "timing_valid": True, + "run_id": "synthetic-run", + "grid_digest": "synthetic-grid", + "config": {"mode": phase}, + "coverage": {"expected_points": 1, "completed_points": 1, "skipped_points": 0}, + "dp": {"rank": 0, "size": 1}, + "results": [{"point": point, "fpms": [fpm]}], + "iteration_groups": [ + { + "benchmark_id": 1, + "point": point, + "expected_dp_ranks": [0], + "complete": True, + "wall_time": 0.01, + "rank_results": [{"dp_rank": 0, "fpms": [fpm]}], + } + ], + "skipped_points": [], + "missing_phases": [], + "timing": {"benchmark_elapsed_seconds": 1.0, "measured_iteration_seconds": 0.01}, + "kvwarm": {"enabled": True, "warm_eligible": True, "skip_reason": None}, + } + ) + ) + + def collect(self, _pods, *, require_benchmark=True): + pass + + def cleanup(self): + pass + + monkeypatch.setattr(planner, "_git_revision", lambda: "test-source-revision") + monkeypatch.setattr(runner, "_render_cell", render_cell) + monkeypatch.setattr(runner, "KubernetesCellRunner", LocalResource) + return executions + + +@pytest.mark.parametrize("smoke", [False, True]) +@pytest.mark.parametrize( + "declared,observed", + [ + ("0.25.1", "0.25.1"), + ("0.25.1", "0.25.1+cu128"), + ("0.25.1", "0.26.0"), + ("0.25.1+cu128", "0.25.1+cu128"), + ], +) +def test_formal_collection_checks_published_version_and_preserves_evidence_on_resume( + tmp_path, monkeypatch, capsys, smoke, declared, observed +): + executions = _synthetic_collector(monkeypatch, observed) + command = _local_collection_command(tmp_path, framework_version=declared) + root = tmp_path / "plan" + if declared != observed: + command += ["--checkpoint-dir", str(root / "fpm-checkpoint/custom")] + if smoke: + command += ["--smoke", "--limit", "1"] + capsys.readouterr() + expected = int(not smoke and declared != observed) + + assert cli.main(command) == expected + + stderr = capsys.readouterr().err + if expected: + assert f"pod-reported vllm version {observed!r}" in stderr + assert f"framework_version {declared!r}" in stderr + assert "preserved" in stderr + assert "new plan" in stderr + assert "Traceback" not in stderr + provenance = list((root / "fpm-artifacts").rglob("collector-provenance.json")) + assert provenance + assert {json.loads(path.read_text())["runtime"]["backend_version"] for path in provenance} == {observed} + metadata = root / "systems/data/h200_sxm/vllm" / observed / "fpm_forward_perf.metadata.json" + if smoke: + assert not list((root / "systems/data").iterdir()) + else: + assert json.loads(metadata.read_text())["backend_version"] == observed + if declared != observed: + # A stale matching directory must not hide this campaign's different runtime. + (root / "systems/data/h200_sxm/vllm" / declared).mkdir() + preserved = { + path: path.read_bytes() for path in root.rglob("*") if path.is_file() and path.name not in {"run-manifest.json"} + } + initial_executions = executions.copy() + + assert cli.main([*command, "--resume"]) == expected + assert executions == initial_executions + assert all(path.read_bytes() == contents for path, contents in preserved.items()) + if not smoke: + # A completed formal checkpoint remains resumable after raw-artifact reclamation. + for path in provenance: + shutil.rmtree(path.parent.parent) + assert cli.main([*command, "--resume"]) == expected + assert executions == initial_executions + assert metadata.read_bytes() == preserved[metadata] + + +def test_formal_collection_requires_publication_evidence(tmp_path, monkeypatch, capsys): + from collector.fpm_forward import runner + + command = _local_collection_command(tmp_path) + monkeypatch.setattr(runner, "run_collection", lambda *args, **kwargs: []) + capsys.readouterr() + + assert cli.main(command) == 1 + stderr = capsys.readouterr().err + assert "aisimulate onboard collect-fpm failed:" in stderr + assert "Traceback" not in stderr + + +def test_deployment_options_reach_frozen_collector_plan(tmp_path, monkeypatch): + from collector.fpm_forward import runner + + calls = [] + monkeypatch.setattr(runner, "run_collection", lambda plan, **kwargs: calls.append((plan, kwargs)) or []) + command = _local_collection_command(tmp_path) + command += [ + "--smoke", + "--dynamo-version", + "1.2.0", + "--image", + "registry.example/fpm@sha256:" + "a" * 64, + "--namespace", + "pilot-collection", + "--model-cache", + "model-cache:/models:checkpoint", + "--transport", + "nvlink", + "--image-pull-secret", + "registry-secret", + ] + + assert cli.main(command) == 0 + + [(plan, kwargs)] = calls + assert kwargs["generator_overrides"] == { + "generator_dynamo_version": "1.2.0", + "K8sConfig": { + "k8s_image": "registry.example/fpm@sha256:" + "a" * 64, + "k8s_namespace": "pilot-collection", + "k8s_pvc_name": "model-cache", + "k8s_pvc_mount_path": "/models", + "k8s_model_path_in_pvc": "checkpoint", + "transport": "nvlink", + "k8s_image_pull_secret": "registry-secret", + }, + } + assert plan.cells + assert kwargs["artifact_root"] == str(tmp_path / "plan/fpm-artifacts") + assert kwargs["database_root"] == str(tmp_path / "plan/systems/data") + + +@pytest.mark.parametrize( + "option,value", + [ + ("--dynamo-version", "latest"), + ("--image", "bad image"), + ("--image", ""), + ("--namespace", "invalid namespace"), + ("--model-cache", ":/models"), + ("--model-cache", "cache:a:b:c"), + ("--model-cache", "cache:relative"), + ("--transport", "pcie"), + ("--image-pull-secret", "invalid secret"), + ("--generator-set", "params.agg.tp=8"), + ("--generator-config", "unrestricted.yaml"), + ], +) +def test_deployment_options_reject_invalid_or_unrestricted_inputs(tmp_path, capsys, option, value): + command = _local_collection_command(tmp_path) + command.remove("--execute") + + with pytest.raises(SystemExit) as error: + cli.main([*command, option, value]) + + assert error.value.code == 2 + assert "Traceback" not in capsys.readouterr().err + assert not (tmp_path / "plan/fpm-checkpoint").exists() + assert not (tmp_path / "plan/fpm-artifacts").exists() + + +@pytest.mark.parametrize( + "option,first,changed", + [ + ("--dynamo-version", "1.2.0", "1.3.0"), + ("--image", "registry.example/fpm:first", "registry.example/fpm:second"), + ("--namespace", "first", "second"), + ("--model-cache", "cache:/models:first", "cache:/models:second"), + ("--transport", "nvlink", "ib"), + ("--image-pull-secret", "first", "second"), + ], +) +def test_deployment_resume_requires_the_same_frozen_identity(tmp_path, monkeypatch, capsys, option, first, changed): + from collector.fpm_forward import planner, runner + + command = [*_local_collection_command(tmp_path), "--smoke", option, first] + checkpoint = tmp_path / "plan/fpm-checkpoint/fpm_forward_smoke.json" + plans = [] + actual_run = runner.run_collection + + def checkpoint_only(plan, **kwargs): + # No silicon data is produced: exercise planning and checkpoint guards + # with a synthetic empty campaign checkpoint in place of GPU execution. + plans.append(plan) + checkpoint.parent.mkdir(exist_ok=True) + checkpoint.write_text(json.dumps({"schema": runner.CHECKPOINT_SCHEMA, "plan_sha256": plan.sha256, "cells": {}})) + return [] + + monkeypatch.setattr(planner, "_git_revision", lambda: "test-source-revision") + monkeypatch.setattr(runner, "run_collection", checkpoint_only) + assert cli.main(command) == 0 + saved = checkpoint.read_bytes() + with pytest.raises(SystemExit) as error: + cli.main(command) + assert error.value.code == 2 + assert len(plans) == 1 + assert cli.main([*command, "--resume"]) == 0 + assert plans[0].sha256 == plans[1].sha256 + assert checkpoint.read_bytes() == saved + + # The real collector must reject a changed identity before GPU execution. + monkeypatch.setattr(runner, "run_collection", actual_run) + monkeypatch.setattr(runner, "KubernetesCellRunner", lambda *args, **kwargs: pytest.fail("must not launch GPU work")) + capsys.readouterr() + assert cli.main([*command, "--resume", option, changed]) == 1 + assert "checkpoint does not match the current frozen plan" in capsys.readouterr().err + assert checkpoint.read_bytes() == saved + + +@pytest.mark.parametrize("stage", ["input", "execution"]) +@pytest.mark.parametrize("error_type", [RuntimeError, ValueError, OSError, KeyboardInterrupt]) +def test_collector_failures_keep_public_cli_exit_codes(tmp_path, monkeypatch, capsys, stage, error_type): + from collector.fpm_forward import cli as collector_cli + + request_path = tmp_path / "request.yaml" + root = tmp_path / "plan" + assert cli.main(_init_args(request_path)) == 0 + assert cli.main(["onboard", "plan", "-c", str(request_path), "--output-dir", str(root)]) == 0 + capsys.readouterr() + + def fail(*args, **kwargs): + raise error_type("test collection failure") + + monkeypatch.setattr(collector_cli, "build_collection_case_plan", lambda **kwargs: SimpleNamespace(model_path=None)) + monkeypatch.setattr(collector_cli, "resolve_run_inputs", fail if stage == "input" else lambda *args: None) + monkeypatch.setattr(collector_cli, "run_resolved", fail) + command = ["onboard", "collect-fpm", "-c", str(request_path), "--output-dir", str(root), "--execute"] + + if error_type is KeyboardInterrupt: + assert cli.main(command) == 130 + elif stage == "input": + with pytest.raises(SystemExit) as error: + cli.main(command) + assert error.value.code == 2 + else: + assert cli.main(command) == 1 + stderr = capsys.readouterr().err + assert "Traceback" not in stderr + if stage == "execution" and error_type is not KeyboardInterrupt: + assert "aisimulate onboard collect-fpm failed: test collection failure" in stderr + + +@pytest.mark.parametrize("error_type", [ImportError, KeyboardInterrupt]) +def test_collector_import_failures_keep_public_cli_exit_codes(tmp_path, monkeypatch, capsys, error_type): + command = _local_collection_command(tmp_path) + real_import = builtins.__import__ + + def fail_collector_import(name, *args, **kwargs): + if name == "collector.fpm_forward.cli": + raise error_type("collector dependency unavailable") + return real_import(name, *args, **kwargs) + + monkeypatch.setattr(builtins, "__import__", fail_collector_import) + capsys.readouterr() + + assert cli.main(command) == (130 if error_type is KeyboardInterrupt else 1) + stderr = capsys.readouterr().err + assert "Traceback" not in stderr + if error_type is ImportError: + assert "aisimulate onboard collect-fpm failed: collector dependency unavailable" in stderr + + +def test_malformed_resumed_checkpoint_fails_cleanly_without_launching_collection(tmp_path, capsys): + from collector.fpm_forward.runner import CHECKPOINT_SCHEMA + + command = _local_collection_command(tmp_path) + capsys.readouterr() + assert cli.main([argument for argument in command if argument != "--execute"]) == 0 + preview = shlex.split(capsys.readouterr().out) + launch_marker = tmp_path / "cluster-command-launched" + bin_dir = tmp_path / "bin" + bin_dir.mkdir() + kubectl = bin_dir / "kubectl" + kubectl.write_text('#!/bin/sh\n: > "$FPM_TEST_LAUNCH_MARKER"\nexit 99\n') + kubectl.chmod(0o755) + environment = { + **os.environ, + "PATH": str(bin_dir) + os.pathsep + os.environ.get("PATH", ""), + "FPM_KUBECTL": "kubectl", + "FPM_TEST_LAUNCH_MARKER": str(launch_marker), + } + planned = subprocess.run( + [sys.executable, *preview[1:]], + cwd=tmp_path, + env=environment, + stdin=subprocess.DEVNULL, + capture_output=True, + text=True, + timeout=30, + ) + assert planned.returncode == 0, planned.stderr + frozen = json.loads(planned.stdout) + root = tmp_path / "plan" + checkpoint = root / "fpm-checkpoint/fpm_forward.json" + checkpoint.parent.mkdir() + checkpoint.write_text( + json.dumps( + { + "schema": CHECKPOINT_SCHEMA, + "plan_sha256": frozen["sha256"], + "cells": {frozen["cells"][0]["cell_id"]: []}, + } + ) + ) + original = checkpoint.read_bytes() + + resumed = subprocess.run( + [sys.executable, "-m", "aisimulate", *command, "--resume"], + cwd=tmp_path, + env=environment, + stdin=subprocess.DEVNULL, + capture_output=True, + text=True, + timeout=30, + ) + + assert resumed.returncode == 1, resumed.stderr + assert "Traceback" not in resumed.stderr + assert "aisimulate onboard collect-fpm failed:" in resumed.stderr + assert checkpoint.read_bytes() == original + assert not launch_marker.exists() + assert not list((root / "fpm-artifacts").rglob("*.yaml")) + + +def test_real_collector_artifact_symlink_rejected_before_first_write(tmp_path, monkeypatch, capsys): + from collector.fpm_forward import runner + + command = _local_collection_command(tmp_path) + capsys.readouterr() + assert cli.main([argument for argument in command if argument != "--execute"]) == 0 + preview = shlex.split(capsys.readouterr().out) + planned = subprocess.run( + [sys.executable, *preview[1:]], + cwd=tmp_path, + stdin=subprocess.DEVNULL, + capture_output=True, + text=True, + timeout=30, + ) + assert planned.returncode == 0, planned.stderr + frozen = json.loads(planned.stdout) + root = tmp_path / "plan" + external = tmp_path / "external" + external.mkdir() + (root / "fpm-artifacts").mkdir() + (root / "fpm-artifacts" / frozen["sha256"][:16]).symlink_to(external, target_is_directory=True) + reached_checkpoint = [] + + def stop_before_gpu(*args, **kwargs): + reached_checkpoint.append(True) + raise RuntimeError("collector reached checkpoint loading") + + # The real collector writes collection-plan.json before loading its + # checkpoint. Keep that path intact and stop before any GPU work. + monkeypatch.setattr(runner, "_load_checkpoint", stop_before_gpu) + + with pytest.raises(SystemExit) as error: + cli.main(command) + + assert error.value.code == 2 + assert "refusing symlinked plan output" in capsys.readouterr().err + assert not reached_checkpoint + assert not list(external.iterdir()) + assert not (root / "fpm-checkpoint").exists() + + +@pytest.mark.parametrize("contents", ["[one, two]\n", "identity: [\n", "identity: {}\n"]) +def test_bad_request_yaml_is_reported_without_a_traceback(tmp_path, capsys, contents) -> None: + request = tmp_path / "bad.yaml" + request.write_text(contents) + + with pytest.raises(SystemExit) as error: + cli.main(["onboard", "plan", "--config", str(request), "--output-dir", str(tmp_path / "plan")]) + + assert error.value.code == 2 + assert "Traceback" not in capsys.readouterr().err + assert not (tmp_path / "plan").exists() + + +@pytest.mark.parametrize("command", ["predict", "recommend"]) +def test_ordinary_cli_options_and_numerical_defaults_are_preserved(command) -> None: + args = cli.build_parser().parse_args([command, "--config", "ordinary.yaml"]) + assert args.command == command + assert args.stack == "engine" + assert args.output_dir == "./aisimulate-output" + assert args.overrides == [] + assert args.overwrite is False + assert args.format == "table" + config = CorePredictionConfig.model_validate( + {"engine": {"model": "example/model", "hardware": "h200_sxm", "workers": {"aggregated": {}}}} + ) + assert config.engine.workers.aggregated.timing.forward_model == "op_level" + assert config.traffic.load.concurrency == 10 + assert config.traffic.stop.requests == 100 + + +def test_installed_module_scripted_init_works_without_a_terminal(tmp_path) -> None: + output = tmp_path / "request.yaml" + result = subprocess.run( + [sys.executable, "-m", "aisimulate.main", *_init_args(output)], + stdin=subprocess.DEVNULL, + capture_output=True, + text=True, + timeout=30, + ) + assert result.returncode == 0, result.stderr + assert SupportRequest.from_yaml(output).identity.model == _REQUIRED["model"] + assert "next:" in result.stdout diff --git a/python/aisimulate/tests/unit/test_support_plan.py b/python/aisimulate/tests/unit/test_support_plan.py new file mode 100644 index 000000000..c07453354 --- /dev/null +++ b/python/aisimulate/tests/unit/test_support_plan.py @@ -0,0 +1,1061 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""The onboarding plan stays usable before model integration or GPU collection.""" + +from __future__ import annotations + +import builtins +import fcntl +import inspect +import json +import os +import select +import shlex +import subprocess +import sys +import zipfile +from contextlib import contextmanager +from importlib.machinery import PathFinder +from importlib.resources import files +from pathlib import Path +from typing import ClassVar + +import pytest +from pydantic import ValidationError + +from aisimulate.config import CorePredictionConfig, CoreRecommendationConfig +from aisimulate.support.fpm import run_fpm +from aisimulate.support.plan import create_plan, plan_lock, request_id +from aisimulate.support.schema import SupportRequest + +pytestmark = pytest.mark.unit + + +class _OfflineRunner: + def run(self, spec): + from aisimulate.sweeper.replay import ReplayReport + + return ReplayReport(metrics={"output_throughput_tok_s": 10.0, "gpu_hours": 1.0}) + + def close(self): + pass + + +class _OfflineRunnerFactory: + def capabilities(self): + from aisimulate.sweeper.replay import RunnerCapabilities + + return RunnerCapabilities(supported_backend_topologies=(("vllm", "agg"),)) + + def create(self, worker_id): + return _OfflineRunner() + + +class _OfflinePerfModel: + requests: ClassVar[list[dict]] = [] + + @classmethod + def best_available(cls, request): + import aisimulate_core + + payload = request.to_dict() + cls.requests.append(payload) + config = json.loads(aisimulate_core.RustForwardPassPerfModel.normalize_config(json.dumps(payload))) + return cls(config) + + def __init__(self, config): + self.config = config + + def diagnostics(self): + return { + "readiness": "ready", + "provenance": { + "config": self.config, + "selected_systems_root": self.config["systems_paths"][0], + "requested_estimation_mode": "fpm_interpolation", + "selected_estimation_mode": "fpm_interpolation", + "selection_failures": [], + }, + } + + def close(self): + pass + + +def _assert_estimator_requests(requests, systems_root, moe_tensor): + expected = { + "model": "unregistered/Example Model", + "system": "h200_sxm", + "backend": "vllm", + "backend_version": "0.25.1", + "worker_type": "aggregated", + "tp": 4, + "pp": 1, + "attention_dp": 1, + "moe_tp_size": moe_tensor if moe_tensor > 1 else None, + "moe_ep_size": 1 if moe_tensor > 1 else None, + "kv_block_size": 64, + "systems_paths": [str(systems_root)], + "estimation_mode": "fpm_interpolation", + "fallback_policy": "deny", + "database_mode": "SILICON", + "estimator_config": {}, + } + for request in requests: + assert {key: request[key] for key in expected} == expected + + +def _request(**updates) -> SupportRequest: + payload = { + "identity": { + "model": "unregistered/Example Model", + "model_revision": "checkpoint-2026-09-14", + "model_kind": "moe", + "framework_version": "0.25.1", + "gpu": "h200_sxm", + "gpu_count": 12, + "node_count": 2, + "gpus_per_node": 6, + "interconnect": "NVLink", + }, + "search": {"tensor_parallel": 4}, + } + for section, values in updates.items(): + payload.setdefault(section, {}).update(values) + return SupportRequest.model_validate(payload) + + +@pytest.mark.parametrize( + "updates", + [ + {"identity": {"model_kind": "auto"}}, + {"identity": {"model_revision": "main"}}, + {"identity": {"framework_version": " "}}, + {"identity": {"gpu": "../h200_sxm"}}, + {"identity": {"gpu_count": 13}}, + {"search": {"tensor_parallel": 8}}, + {"search": {"tensor_parallel": True}}, + {"search": {"max_candidates": 3}}, + {"search": {"context_length": 1151}}, + {"workload": {"input_tokens": 0}}, + {"workload": {"request_count": 0}}, + {"workload": {"concurrency": 5}}, + {"workload": {"slo": {"ttft_ms": float("inf"), "tpot_ms": 1.0}}}, + {"workload": {"slo": {"ttft_ms": 1.0, "tpot_ms": True}}}, + ], +) +def test_request_rejects_invalid_allocation_workload_and_identity(updates): + with pytest.raises(ValidationError): + _request(**updates) + + +def test_plan_does_not_resolve_unknown_model_and_reloads_public_configs(tmp_path, monkeypatch): + real_import = builtins.__import__ + + def no_model_or_collector_import(name, *args, **kwargs): + assert not name.startswith(("collector", "huggingface_hub", "transformers")), name + assert ".sdk.models" not in name, name + return real_import(name, *args, **kwargs) + + monkeypatch.setattr(builtins, "__import__", no_model_or_collector_import) + request = _request() + root = tmp_path / "plan with 'quotes'" + plan = create_plan(request, root) + prediction = CorePredictionConfig.from_yaml(root / "predict/pilot.yaml") + recommendation = CoreRecommendationConfig.from_yaml(root / "recommend/pilot.yaml") + assert SupportRequest.from_yaml(root / "request.yaml") == request + assert json.loads((root / "support-plan.json").read_text()) == plan + assert prediction.engine.systems_paths == [str(root / "systems")] + assert recommendation.engine.systems_paths == prediction.engine.systems_paths + assert prediction.engine.model == "unregistered/Example Model" + worker = prediction.engine.workers.aggregated + assert worker.timing.type == "default" + assert worker.timing.estimation_mode == "fpm_interpolation" + assert worker.timing.fallback_policy == "deny" + assert worker.parallelism.model_dump() == { + "replicas": 1, + "tensor": 4, + "pipeline": 1, + "attention_data": 1, + "moe_tensor": 4, + "moe_expert": 1, + "decode_context": None, + } + assert recommendation.engine.workers.aggregated.timing.estimation_mode == "fpm_interpolation" + assert recommendation.engine.workers.aggregated.timing.fallback_policy == "deny" + assert recommendation.optimizer.max_trials == 1 + assert recommendation.optimizer.parallelism == 1 + assert plan["search"]["candidate_count"] == 1 + assert "single" in plan["search"]["detail"] + assert {item["status"] for item in plan["prerequisites"]} == {"not_checked"} + assert plan["accuracy"]["status"] == "not_assessed" + assert (root / "systems/h200_sxm.yaml").is_file() + assert not list((root / "systems/data").iterdir()) + assert not (root / "evidence.yaml").exists() + + +def test_onboarding_plan_and_preview_do_not_import_estimator_runtime(tmp_path): + request = tmp_path / "request.json" + request.write_text(_request().model_dump_json()) + root = tmp_path / "plan" + result = subprocess.run( + [ + sys.executable, + "-c", + """ +import importlib.abc +import sys +from pathlib import Path + +blocked = ("aisimulate._runtime", "aisimulate_core", "aiconfigurator", "aiconfigurator_core", + "numpy", "pandas", "torch", "transformers", "huggingface_hub", "collector") + +class NoEstimatorRuntime(importlib.abc.MetaPathFinder): + def find_spec(self, fullname, path=None, target=None): + if fullname == "aisimulate_core": + return None # Resource lookup may inspect the package spec without importing it. + assert not any(fullname == name or fullname.startswith(name + ".") for name in blocked), fullname + +sys.meta_path.insert(0, NoEstimatorRuntime()) +from aisimulate.supervision import main + +request, output = sys.argv[1:] +assert main(["onboard", "plan", "--config", request, "--output-dir", output, "--format", "json"]) == 0 +assert main(["onboard", "collect-fpm", "--config", request, "--output-dir", output]) == 0 +assert (Path(output) / "systems/h200_sxm.yaml").is_file() +assert not any(name == prefix or name.startswith(prefix + ".") for name in sys.modules for prefix in blocked) +""", + str(request), + str(root), + ], + cwd=tmp_path, + capture_output=True, + text=True, + timeout=30, + check=False, + ) + + assert result.returncode == 0, result.stdout + result.stderr + assert not (root / "fpm-checkpoint").exists() + assert not (root / "fpm-artifacts").exists() + + +@pytest.mark.parametrize("model_kind,preset,moe_tensor", [("dense", "tp", 1), ("moe", "pure_tp", 4)]) +def test_plan_collects_one_chosen_worker_and_bounds_recommendation(tmp_path, model_kind, preset, moe_tensor): + from aisimulate.recommend import recommendation_to_sweeper + + request = _request(identity={"model_kind": model_kind}, search={"max_candidates": 2}) + plan = create_plan(request, tmp_path) + configs = [CoreRecommendationConfig.from_yaml(path) for path in plan["outputs"]["recommendation_configs"]] + candidates = [candidate for config in configs for candidate in config.engine.workers.aggregated.parallelism.preset] + assert [(c.replicas, c.tensor, c.moe_tensor) for c in candidates] == [(1, 4, moe_tensor), (2, 4, moe_tensor)] + assert all(config.optimizer.max_trials == 1 for config in configs) + assert plan["search"]["candidate_count"] == 2 + for config in configs: + search = recommendation_to_sweeper(config).search_space + assert search.agg_max_num_batched_tokens == [8192] + assert search.agg_max_num_seqs == [256] + commands = json.loads((tmp_path / "commands.json").read_text()) + command = commands["fpm_plan_local"] + assert command[command.index("--fpm-gpu-counts") + 1] == "4" + assert command[command.index("--fpm-max-gpus") + 1] == "4" + assert command[command.index("--fpm-parallel-presets") + 1] == preset + assert command[command.index("--fpm-max-prefill-isl") + 1] == "1024" + assert command[command.index("--fpm-max-prefill-batch-size") + 1] == "1" + assert "--plan-only" in command + assert "--execute" in commands["fpm_run_local"] + + +def test_one_worker_allocation_keeps_one_pilot_even_with_two_candidate_limit(tmp_path): + request = _request(identity={"node_count": 1, "gpus_per_node": 4, "gpu_count": 4}, search={"max_candidates": 2}) + plan = create_plan(request, tmp_path) + commands = json.loads((tmp_path / "commands.json").read_text()) + + assert plan["search"]["candidate_count"] == 1 + assert plan["outputs"]["recommendation_configs"] == [str(tmp_path / "recommend/pilot.yaml")] + assert len(commands["recommend"]) == 1 + assert "single" in plan["search"]["detail"] + + +@pytest.mark.parametrize("seed", [0, 7, 42]) +@pytest.mark.parametrize("model_kind,moe_tensor", [("dense", 1), ("moe", 4)]) +def test_emitted_recommendations_evaluate_both_replica_choices(tmp_path, monkeypatch, seed, model_kind, moe_tensor): + from aisimulate.main import build_parser + from aisimulate.recommend import run_recommendation + from aisimulate.sweeper.parallel_enum import ParallelShape, ReplicaParallelConfig + from aisimulate_core.sdk import RustForwardPassPerfModel + + # Keep process-local doubles in this orchestration test. The real lowering, + # estimator resolver, sampler, resource admission and scoring still run. + monkeypatch.setattr("aisimulate.supervision.in_supervised_process", lambda: True) + monkeypatch.setattr(_OfflinePerfModel, "requests", []) + monkeypatch.setattr(RustForwardPassPerfModel, "best_available", _OfflinePerfModel.best_available) + legal = [ + ReplicaParallelConfig(shape=ParallelShape(tp=4, pp=1, dp=1, moe_tp=moe_tensor, moe_ep=1), replicas=count) + for count in (1, 2) + ] + monkeypatch.setattr("aisimulate.sweeper.search_space.parallel_configs_for", lambda *args, **kwargs: legal) + request = _request(identity={"model_kind": model_kind}, search={"max_candidates": 2, "seed": seed}) + plan = create_plan(request, tmp_path) + commands = json.loads((tmp_path / "commands.json").read_text())["recommend"] + evaluated = [] + results = [] + configs = [] + outputs = [] + for command in commands: + args = build_parser().parse_args(command[1:]) + config = CoreRecommendationConfig.from_yaml(args.config) + result = run_recommendation(config, stack="engine", runner_factory=_OfflineRunnerFactory(), show_progress=False) + evaluated.extend(candidate.config["replicas"] for candidate in result.candidates) + results.append(result) + configs.append(args.config) + outputs.append(args.output_dir) + assert evaluated == [1, 2] + assert configs == plan["outputs"]["recommendation_configs"] + assert outputs == plan["outputs"]["recommendation_results"] + assert len(set(outputs)) == 2 + assert all(result.counts.evaluated == 1 and result.counts.cache_hits == 0 for result in results) + assert len(_OfflinePerfModel.requests) == 2 + _assert_estimator_requests(_OfflinePerfModel.requests, tmp_path / "systems", moe_tensor) + + +@pytest.mark.parametrize( + "max_candidates,missing_pilot_samples,statuses", + [(1, False, [0]), (2, False, [0, 0]), (2, True, [1, 0])], +) +@pytest.mark.parametrize("tampered_commands", [False, True]) +def test_documented_recommendation_loop_attempts_every_candidate( + tmp_path, max_candidates, missing_pilot_samples, statuses, tampered_commands +): + guide = Path(__file__).resolve().parents[4] / "docs/fpm-self-service.md" + snippet = guide.read_text().split("python3 - <<'PY'\n", 1)[1].split("\nPY\n", 1)[0] + workdir = tmp_path / "work with 'quotes'; $(literal)" + workdir.mkdir() + request = _request( + identity={"model_kind": "dense", "gpu_count": 8, "node_count": 1, "gpus_per_node": 8}, + search={"max_candidates": max_candidates, "objective": "ttft"}, + ) + plan = create_plan(request, workdir / "aisimulate-support") + commands = json.loads(Path(plan["outputs"]["commands"]).read_text())["recommend"] + injected_marker = tmp_path / "injected-command-ran" + if tampered_commands: + command_file = Path(plan["outputs"]["commands"]) + payload = json.loads(command_file.read_text()) + payload["recommend"].insert( + 0, + [ + sys.executable, + "-c", + "import sys; from pathlib import Path; Path(sys.argv[1]).write_text('unexpected execution')", + str(injected_marker), + ], + ) + command_file.write_text(json.dumps(payload)) + evaluated = tmp_path / "evaluated.txt" + estimator_requests = tmp_path / "estimator-requests.jsonl" + wrapper = tmp_path / "bin/aisimulate" + wrapper.parent.mkdir() + # Install the same Core double in each supervised child. The snippet executes + # real subprocesses, CLI dispatch, estimator resolution, sampling and export. + wrapper.write_text( + f"#!{sys.executable}\n" + f""" +import os +import sys +import json +from pathlib import Path +from typing import ClassVar + +from aisimulate_core.sdk import RustForwardPassPerfModel +import aisimulate.main as cli +import aisimulate.sweeper.search_space as search_space +from aisimulate.sweeper.parallel_enum import ParallelShape, ReplicaParallelConfig +from aisimulate.sweeper.replay import ReplayReport, RunnerCapabilities + +{inspect.getsource(_OfflinePerfModel)} + +def best_available(request): + model = _OfflinePerfModel.best_available(request) + with Path(os.environ["SUPPORT_TEST_ESTIMATOR_REQUESTS"]).open("a") as stream: + stream.write(json.dumps(request.to_dict()) + "\\n") + return model + +RustForwardPassPerfModel.best_available = best_available + +legal = [ + ReplicaParallelConfig(shape=ParallelShape(tp=4, pp=1, dp=1, moe_tp=1, moe_ep=1), replicas=count) + for count in (1, 2) +] +search_space.parallel_configs_for = lambda *args, **kwargs: legal + +class Runner: + def run(self, spec): + replicas = spec.backend_deployment.num_workers + with Path(os.environ["SUPPORT_TEST_EVALUATED"]).open("a") as stream: + stream.write(str(replicas) + "\\n") + samples = 0.0 if {missing_pilot_samples!r} and replicas == 1 else 4.0 + return ReplayReport(metrics={{"mean_ttft_ms": 10.0, "num_ttft_samples": samples}}) + + def close(self): + pass + +class Factory: + def capabilities(self): + return RunnerCapabilities(supported_backend_topologies=(("vllm", "agg"),)) + + def create(self, worker_id): + return Runner() + +cli.resolve_runner_factory = lambda stack: Factory() +if __name__ == "__main__": + raise SystemExit(cli.main(sys.argv[1:])) +""" + ) + wrapper.chmod(0o755) + result = subprocess.run( + [sys.executable, "-c", snippet], + cwd=workdir, + env={ + **os.environ, + "PATH": str(wrapper.parent) + os.pathsep + os.environ.get("PATH", ""), + "SUPPORT_TEST_EVALUATED": str(evaluated), + "SUPPORT_TEST_ESTIMATOR_REQUESTS": str(estimator_requests), + }, + capture_output=True, + text=True, + timeout=30, + check=False, + ) + + assert not injected_marker.exists() + assert result.returncode == (1 if any(statuses) else 0), result.stderr + assert evaluated.read_text().splitlines() == [str(index) for index in range(1, max_candidates + 1)] + resolved_requests = [json.loads(line) for line in estimator_requests.read_text().splitlines()] + assert len(resolved_requests) == max_candidates + _assert_estimator_requests(resolved_requests, workdir / "aisimulate-support/systems", 1) + for index, (output, status) in enumerate(zip(plan["outputs"]["recommendation_results"], statuses, strict=True), 1): + root = Path(output) + payload = json.loads((root / "recommendation.json").read_text()) + assert payload["counts"]["evaluated"] == 1 + assert payload["counts"]["feasible"] == (0 if status else 1) + assert payload["counts"]["infeasible"] == (1 if status else 0) + assert payload["candidates"][0]["config"]["replicas"] == index + exported = root / "recommendations/0001.yaml" + if status: + assert "no qualifying samples" in payload["candidates"][0]["reason"] + assert not exported.exists() + else: + prediction = CorePredictionConfig.from_yaml(exported) + assert prediction.engine.workers.aggregated.parallelism.replicas == index + assert prediction.engine.systems_paths == [str(workdir / "aisimulate-support/systems")] + timing = prediction.engine.workers.aggregated.timing + assert timing.estimation_mode == "fpm_interpolation" + assert timing.fallback_policy == "deny" + assert timing.systems_paths == prediction.engine.systems_paths + assert [line for line in result.stdout.splitlines() if line.startswith("exit ")] == [ + f"exit {status}: {shlex.join(command)}" for command, status in zip(commands, statuses, strict=True) + ] + + +def test_prefill_bounds_follow_concurrent_workload_and_collector_minimum(tmp_path): + request = _request(workload={"input_tokens": 10, "concurrency": 3, "request_count": 4}) + command = create_plan(request, tmp_path)["fpm"]["plan_command"] + assert command[command.index("--fpm-max-prefill-isl") + 1] == "30" + assert command[command.index("--fpm-max-prefill-batch-size") + 1] == "3" + small = _request(workload={"input_tokens": 1}) + command = create_plan(small, tmp_path / "small")["fpm"]["plan_command"] + assert command[command.index("--fpm-max-prefill-isl") + 1] == "2" + + +@pytest.mark.parametrize( + "updates", + [ + {"search": {"tensor_parallel": 2}}, + {"search": {"max_candidates": 2}}, + {"search": {"seed": 7}}, + {"workload": {"input_tokens": 64}}, + {"identity": {"model_revision": "other-checkpoint"}}, + {"identity": {"framework_version": "0.26.0"}}, + {"identity": {"tokenizer_revision": "tokenizer-v2"}}, + ], +) +@pytest.mark.parametrize("complete", [False, True]) +def test_overwrite_rejects_different_requests_without_touching_data(tmp_path, updates, complete): + request = _request() + create_plan(request, tmp_path) + if not complete: + (tmp_path / "support-plan.json").unlink() + data = tmp_path / "systems/data/collected.parquet" + data.write_bytes(b"timings") + before = {str(p.relative_to(tmp_path)): p.read_bytes() for p in tmp_path.rglob("*") if p.is_file()} + changed = _request(**updates) + assert request_id(changed) != request_id(request) + with pytest.raises(ValueError, match="identity|different request|mix"): + create_plan(changed, tmp_path, overwrite=True) + after = {str(p.relative_to(tmp_path)): p.read_bytes() for p in tmp_path.rglob("*") if p.is_file()} + assert before == after + + +def test_same_request_overwrite_preserves_timings_and_rejects_modified_inputs(tmp_path): + request = _request() + plan = create_plan(request, tmp_path) + data = tmp_path / "systems/data/collected.parquet" + data.write_bytes(b"timings") + with pytest.raises(ValueError, match="nonempty|already exists|overwrite"): + create_plan(request, tmp_path) + assert create_plan(request, tmp_path, overwrite=True) == plan + assert data.read_bytes() == b"timings" + (tmp_path / "request.yaml").write_text("modified: true\n") + with pytest.raises(ValueError, match="modified|request|identity"): + create_plan(request, tmp_path, overwrite=True) + assert data.read_bytes() == b"timings" + + +@pytest.mark.parametrize("modification", ["retired_command", "missing_execute", "formatting"]) +@pytest.mark.parametrize("complete", [False, True]) +def test_overwrite_rejects_modified_commands_before_repairing_missing_files(tmp_path, modification, complete): + request = _request() + create_plan(request, tmp_path) + if not complete: + (tmp_path / "support-plan.json").unlink() + command_file = tmp_path / "commands.json" + commands = json.loads(command_file.read_text()) + if modification == "retired_command": + commands["fpm_run_local"][1] = "support" + elif modification == "missing_execute": + commands["fpm_run_local"].remove("--execute") + command_file.write_text( + json.dumps(commands, indent=4 if modification == "formatting" else 2, sort_keys=True) + "\n" + ) + _seed_campaign(tmp_path) + prediction = tmp_path / "predict/pilot.yaml" + prediction.unlink() + before_repair = _file_contents(tmp_path) + + with pytest.raises(ValueError, match="generated plan input .*commands.json was modified"): + create_plan(request, tmp_path, overwrite=True) + assert _file_contents(tmp_path) == before_repair + + +@pytest.mark.parametrize("layout", ["directory", "zip"]) +def test_plan_copies_system_specification_from_core_package_resources(tmp_path, monkeypatch, layout): + from aisimulate.support import plan + + system_spec = files("aisimulate_core").joinpath("systems", "h200_sxm.yaml").read_bytes() + system_spec += b"\n# Packaged separately from the onboarding module.\n" + package_init = "raise AssertionError('reading system resources must not execute the core package')\n" + if layout == "directory": + distribution = tmp_path / "separate-core-package" + package = distribution / "aisimulate_core" + (package / "systems").mkdir(parents=True) + (package / "__init__.py").write_text(package_init) + (package / "systems/h200_sxm.yaml").write_bytes(system_spec) + else: + distribution = tmp_path / "core.zip" + with zipfile.ZipFile(distribution, "w") as handle: + handle.writestr("aisimulate_core/__init__.py", package_init) + handle.writestr("aisimulate_core/systems/h200_sxm.yaml", system_spec) + package_spec = PathFinder.find_spec("aisimulate_core", [str(distribution)]) + assert package_spec is not None + + def core_spec(name): + assert name == "aisimulate_core" + return package_spec + + monkeypatch.setattr(plan, "find_spec", core_spec) + root = tmp_path / "plan" + + create_plan(_request(), root) + + assert (root / "systems/h200_sxm.yaml").read_bytes() == system_spec + + +def test_missing_core_package_is_rejected_before_creating_outputs(tmp_path, monkeypatch): + from aisimulate.support import plan + + monkeypatch.setattr(plan, "find_spec", lambda name: None) + root = tmp_path / "plan" + + with pytest.raises(ValueError, match="aisimulate_core package resources"): + create_plan(_request(), root) + + assert not root.exists() + + +def test_new_gpu_is_rejected_before_creating_outputs(tmp_path): + root = tmp_path / "new-plan" + with pytest.raises(ValueError, match="packaged.*system|system specification"): + create_plan(_request(identity={"gpu": "unreleased_gpu"}), root) + assert not root.exists() + + +def test_preview_is_shell_safe_and_does_not_create_outputs_or_import_collector(tmp_path, monkeypatch, capsys): + request = _request() + root = tmp_path / "plan with 'quotes'" + real_import = builtins.__import__ + + def no_collector_import(name, *args, **kwargs): + assert not name.startswith("collector"), name + return real_import(name, *args, **kwargs) + + monkeypatch.setattr(builtins, "__import__", no_collector_import) + assert run_fpm(request, output_dir=root) == 0 + command = shlex.split(capsys.readouterr().out) + assert command[1:3] == ["-m", "collector.fpm_forward"] + assert command[command.index("--model-path") + 1] == request.identity.model + assert command[command.index("--fpm-database-root") + 1] == str(root / "systems/data") + assert "--plan-only" in command + assert not root.exists() + + +def test_execute_delegates_only_for_matching_plan(tmp_path, monkeypatch): + from collector.fpm_forward import cli + + calls = [] + monkeypatch.setattr(cli, "main", lambda argv: calls.append(argv) or 7) + request = _request() + with pytest.raises(ValueError, match="plan"): + run_fpm(request, execute=True, output_dir=tmp_path) + assert not calls + create_plan(request, tmp_path) + assert run_fpm(request, execute=True, output_dir=tmp_path, smoke=True, limit=1) == 7 + assert calls[0][calls[0].index("--limit") + 1] == "1" + assert "--smoke" in calls[0] + assert "--plan-only" not in calls[0] + with pytest.raises(ValueError, match="identity|different request|mix"): + run_fpm(_request(search={"tensor_parallel": 2}), execute=True, resume=True, output_dir=tmp_path) + assert len(calls) == 1 + + +def _seed_campaign(root, *, smoke=False, checkpoint_dir=None): + sha = "a" * 64 + artifact = root / "fpm-artifacts" / sha[:16] + if smoke: + artifact /= "smoke" + timing = artifact / "cells/cell-prefill/raw/existing-timing.json" + timing.parent.mkdir(parents=True) + timing.write_bytes(b"previous measured data") + checkpoint = (checkpoint_dir or root / "fpm-checkpoint") / ( + "fpm_forward_smoke.json" if smoke else "fpm_forward.json" + ) + checkpoint.parent.mkdir(parents=True, exist_ok=True) + checkpoint.write_text(json.dumps({"schema": "aic-fpm-collector-checkpoint-v3", "plan_sha256": sha, "cells": {}})) + if not smoke: + (root / "systems/data/collected.parquet").write_bytes(b"formal timings") + return checkpoint, timing + + +def _file_contents(root): + return {str(path.relative_to(root)): path.read_bytes() for path in root.rglob("*") if path.is_file()} + + +def _reject_collector_import(monkeypatch): + real_import = builtins.__import__ + + def no_collector_import(name, *args, **kwargs): + assert not name.startswith("collector"), "occupied campaigns must be rejected before collector import" + return real_import(name, *args, **kwargs) + + monkeypatch.setattr(builtins, "__import__", no_collector_import) + + +def test_published_execution_command_rejects_an_occupied_campaign(tmp_path, monkeypatch, capsys): + from aisimulate.main import main + + create_plan(_request(), tmp_path) + _seed_campaign(tmp_path) + before = _file_contents(tmp_path) + command = json.loads((tmp_path / "commands.json").read_text())["fpm_run_local"] + assert command[:3] == ["aisimulate", "onboard", "collect-fpm"] + _reject_collector_import(monkeypatch) + + with pytest.raises(SystemExit) as error: + main(command[1:]) + + assert error.value.code == 2 + assert "require --resume" in capsys.readouterr().err + assert _file_contents(tmp_path) == before + + +@pytest.mark.parametrize("smoke", [False, True]) +def test_repeated_execute_requires_resume_without_touching_campaign(tmp_path, monkeypatch, smoke): + create_plan(_request(), tmp_path) + _seed_campaign(tmp_path, smoke=smoke) + before = _file_contents(tmp_path) + _reject_collector_import(monkeypatch) + + with pytest.raises(ValueError, match="resume.*new output|new output.*resume"): + run_fpm(SupportRequest.from_yaml(tmp_path / "request.yaml"), output_dir=tmp_path, execute=True, smoke=smoke) + + assert _file_contents(tmp_path) == before + with plan_lock(tmp_path): + pass + + +@pytest.mark.parametrize( + "smoke,relative", + [ + (False, "systems/data/collected.parquet"), + (False, "fpm-checkpoint/custom/fpm_forward.json"), + (True, "fpm-checkpoint/custom/fpm_forward_smoke.json"), + (False, "fpm-artifacts/aaaaaaaaaaaaaaaa/cells/raw.json"), + (True, "fpm-artifacts/aaaaaaaaaaaaaaaa/smoke/cells/raw.json"), + ], +) +def test_partial_campaign_output_alone_prevents_restart(tmp_path, monkeypatch, smoke, relative): + request = _request() + create_plan(request, tmp_path) + sentinel = tmp_path / relative + sentinel.parent.mkdir(parents=True, exist_ok=True) + sentinel.write_bytes(b"interrupted campaign") + before = _file_contents(tmp_path) + _reject_collector_import(monkeypatch) + + with pytest.raises(ValueError, match="resume"): + run_fpm(request, output_dir=tmp_path, execute=True, smoke=smoke) + + assert _file_contents(tmp_path) == before + + +@pytest.mark.parametrize("smoke", [False, True]) +@pytest.mark.parametrize("checkpoint_state", ["missing", "invalid_json", "empty_object", "new_directory"]) +def test_resume_occupied_campaign_requires_usable_selected_checkpoint(tmp_path, monkeypatch, smoke, checkpoint_state): + request = _request() + create_plan(request, tmp_path) + checkpoint, _ = _seed_campaign(tmp_path, smoke=smoke) + selected_dir = checkpoint.parent + if checkpoint_state == "missing": + checkpoint.unlink() + elif checkpoint_state == "new_directory": + selected_dir /= "new-campaign" + else: + checkpoint.write_text("invalid JSON" if checkpoint_state == "invalid_json" else "{}") + before = _file_contents(tmp_path) + _reject_collector_import(monkeypatch) + + with pytest.raises(ValueError, match="checkpoint"): + run_fpm(request, output_dir=tmp_path, execute=True, smoke=smoke, resume=True, checkpoint_dir=selected_dir) + + assert _file_contents(tmp_path) == before + with plan_lock(tmp_path): + pass + + +@pytest.mark.parametrize("smoke", [False, True]) +def test_matching_explicit_resume_delegates_and_preserves_campaign(tmp_path, monkeypatch, smoke): + from collector.fpm_forward import cli + + request = _request() + create_plan(request, tmp_path) + checkpoint, _ = _seed_campaign(tmp_path, smoke=smoke, checkpoint_dir=tmp_path / "fpm-checkpoint/custom") + before = _file_contents(tmp_path) + calls = [] + monkeypatch.setattr(cli, "main", lambda argv: calls.append(argv) or 7) + + assert ( + run_fpm(request, output_dir=tmp_path, execute=True, smoke=smoke, resume=True, checkpoint_dir=checkpoint.parent) + == 7 + ) + + assert "--resume" in calls[0] + assert calls[0][calls[0].index("--checkpoint-dir") + 1] == str(checkpoint.parent) + assert _file_contents(tmp_path) == before + + +@pytest.mark.parametrize("smoke", [False, True]) +def test_fresh_smoke_and_formal_campaigns_are_independent(tmp_path, monkeypatch, smoke): + from collector.fpm_forward import cli + + request = _request() + create_plan(request, tmp_path) + _seed_campaign(tmp_path, smoke=not smoke) + before = _file_contents(tmp_path) + calls = [] + monkeypatch.setattr(cli, "main", lambda argv: calls.append(argv) or 7) + + assert run_fpm(request, output_dir=tmp_path, execute=True, smoke=smoke) == 7 + + assert len(calls) == 1 + assert _file_contents(tmp_path) == before + + +def test_execute_checks_for_campaign_outputs_while_holding_lock(tmp_path, monkeypatch): + from aisimulate.support import plan + + request = _request() + create_plan(request, tmp_path) + real_lock = plan.plan_lock + + @contextmanager + def raced_lock(root): + with real_lock(root): + _seed_campaign(root) + yield + + monkeypatch.setattr(plan, "plan_lock", raced_lock) + _reject_collector_import(monkeypatch) + + with pytest.raises(ValueError, match="resume"): + run_fpm(request, output_dir=tmp_path, execute=True) + + assert (tmp_path / "systems/data/collected.parquet").read_bytes() == b"formal timings" + with plan_lock(tmp_path): + pass + + +@pytest.mark.parametrize("smoke,limit", [(False, 1), (True, 0), (True, -1), (True, True)]) +def test_invalid_limited_collection_fails_before_preview_or_execution(tmp_path, smoke, limit): + with pytest.raises(ValueError, match="limit"): + run_fpm(_request(), output_dir=tmp_path, smoke=smoke, limit=limit) + + +@pytest.mark.parametrize("complete", [False, True]) +def test_plan_rejects_symlinked_output_without_touching_target(tmp_path, complete): + external = tmp_path / "external" + external.mkdir() + data = external / "collected.parquet" + data.write_bytes(b"timings") + root = tmp_path / "plan" + create_plan(_request(), root) + if not complete: + (root / "support-plan.json").unlink() + (root / "systems/data").rmdir() + (root / "systems/data").symlink_to(external, target_is_directory=True) + with pytest.raises(ValueError, match="symlink"): + create_plan(_request(), root, overwrite=True) + assert data.read_bytes() == b"timings" + + +@pytest.mark.parametrize( + "relative", + [ + "fpm-artifacts/aaaaaaaaaaaaaaaa/smoke/cells/output", + "fpm-checkpoint/custom/output", + "systems/data/h200_sxm/vllm/0.25.1/output", + ], +) +@pytest.mark.parametrize("target_kind", ["directory", "file", "dangling_directory", "dangling_file"]) +@pytest.mark.parametrize("operation", ["repair", "execute"]) +def test_nested_symlinks_rejected_before_repair_or_collector_import( + tmp_path, monkeypatch, relative, target_kind, operation +): + request = _request() + root = tmp_path / "plan" + create_plan(request, root) + _seed_campaign(root) + (root / "predict/pilot.yaml").unlink() + external = tmp_path / "external" + if target_kind == "directory": + external.mkdir() + elif target_kind == "file": + external.write_bytes(b"outside data") + link = root / relative + link.parent.mkdir(parents=True, exist_ok=True) + link.symlink_to(external, target_is_directory=target_kind.endswith("directory")) + before = _file_contents(tmp_path) + _reject_collector_import(monkeypatch) + + with pytest.raises(ValueError, match="refusing symlinked plan output"): + if operation == "repair": + create_plan(request, root, overwrite=True) + else: + run_fpm(request, output_dir=root, execute=True, resume=True) + + assert _file_contents(tmp_path) == before + assert link.is_symlink() + assert link.readlink() == external + if target_kind == "directory": + assert not list(external.iterdir()) + elif target_kind.startswith("dangling"): + assert not external.exists() + + +@pytest.mark.parametrize("relative", ["../shared", "fpm-checkpoint/../../shared"]) +def test_checkpoint_override_cannot_escape_request_directory(tmp_path, relative): + with pytest.raises(ValueError, match="checkpoint_dir"): + run_fpm(_request(), output_dir=tmp_path, checkpoint_dir=tmp_path / relative) + + +@pytest.mark.parametrize("kill_owner", [False, True]) +def test_plan_lock_recovers_after_owner_exit_and_preserves_existing_files(tmp_path, kill_owner): + request = _request() + create_plan(request, tmp_path) + before = (tmp_path / "support-plan.json").read_bytes() + process = subprocess.Popen( + [ + sys.executable, + "-c", + "import sys; from pathlib import Path; from aisimulate.support.plan import plan_lock; " + "lock = plan_lock(Path(sys.argv[1])); lock.__enter__(); " + "print('locked', flush=True); sys.stdin.read(); lock.__exit__(None, None, None)", + str(tmp_path), + ], + stdin=subprocess.PIPE, + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + text=True, + ) + try: + assert select.select([process.stdout], [], [], 10)[0], "lock owner did not become ready" + assert process.stdout.readline().strip() == "locked" + with pytest.raises(ValueError, match="another onboarding operation"): + create_plan(request, tmp_path, overwrite=True) + if kill_owner: + process.kill() + process.communicate(timeout=10) + finally: + if process.poll() is None: + process.kill() + process.communicate(timeout=10) + + (tmp_path / "predict/pilot.yaml").unlink() + create_plan(request, tmp_path, overwrite=True) + assert (tmp_path / "predict/pilot.yaml").is_file() + assert (tmp_path / "support-plan.json").read_bytes() == before + + +def test_plan_lock_does_not_unlink_the_inode_used_by_a_waiting_process(tmp_path): + create_plan(_request(), tmp_path) + with (tmp_path / ".support.lock").open("r") as waiting: + with plan_lock(tmp_path): + pass + fcntl.flock(waiting, fcntl.LOCK_EX | fcntl.LOCK_NB) + with pytest.raises(ValueError, match="another onboarding operation"): + create_plan(_request(), tmp_path, overwrite=True) + + +@pytest.mark.parametrize("existing_target", [False, True]) +def test_plan_lock_does_not_follow_a_symlink(tmp_path, existing_target): + root = tmp_path / "plan" + root.mkdir() + external = tmp_path / "external" + if existing_target: + external.write_text("preserve this") + (root / ".support.lock").symlink_to(external) + + with pytest.raises(OSError): + create_plan(_request(), root) + + assert not (root / "request.yaml").exists() + assert external.read_text() == "preserve this" if existing_target else not external.exists() + + +def test_racing_writer_cannot_be_overwritten_or_reported_as_success(tmp_path, monkeypatch): + original_open = Path.open + destination = tmp_path / "predict/pilot.yaml" + + def racing_open(path, mode="r", *args, **kwargs): + if path == destination and mode == "xb": + with original_open(path, "wb") as handle: + handle.write(b"another writer's config") + return original_open(path, mode, *args, **kwargs) + + monkeypatch.setattr(Path, "open", racing_open) + with pytest.raises(FileExistsError): + create_plan(_request(), tmp_path) + assert destination.read_bytes() == b"another writer's config" + assert not (tmp_path / "support-plan.json").exists() + with plan_lock(tmp_path): + pass + + before_repair = _file_contents(tmp_path) + with pytest.raises(ValueError, match="generated plan input .*pilot.yaml was modified"): + create_plan(_request(), tmp_path, overwrite=True) + assert _file_contents(tmp_path) == before_repair + + monkeypatch.setattr(Path, "open", original_open) + destination.unlink() + create_plan(_request(), tmp_path, overwrite=True) + assert CorePredictionConfig.from_yaml(destination).engine.model == _request().identity.model + + +@pytest.mark.parametrize( + "interrupted_file", + ["predict/pilot.yaml", "recommend/pilot.yaml", "systems/h200_sxm.yaml", "commands.json", "support-plan.json"], +) +def test_overwrite_repairs_interrupted_plan_and_preserves_existing_outputs(tmp_path, monkeypatch, interrupted_file): + request = _request() + original_open = Path.open + + def interrupted_open(path, mode="r", *args, **kwargs): + if path == tmp_path / interrupted_file and mode == "xb": + raise InterruptedError("interrupted plan generation") + return original_open(path, mode, *args, **kwargs) + + monkeypatch.setattr(Path, "open", interrupted_open) + with pytest.raises(InterruptedError): + create_plan(request, tmp_path) + monkeypatch.setattr(Path, "open", original_open) + assert SupportRequest.from_yaml(tmp_path / "request.yaml") == request + assert not (tmp_path / "support-plan.json").exists() + (tmp_path / "systems/data").mkdir(parents=True, exist_ok=True) + _seed_campaign(tmp_path) + before_repair = _file_contents(tmp_path) + + with pytest.raises(ValueError, match="nonempty"): + create_plan(request, tmp_path) + _reject_collector_import(monkeypatch) + with pytest.raises(ValueError, match="readable onboarding plan"): + run_fpm(request, execute=True, resume=True, output_dir=tmp_path) + assert _file_contents(tmp_path) == before_repair + + repaired = create_plan(request, tmp_path, overwrite=True) + + assert json.loads((tmp_path / "support-plan.json").read_text()) == repaired + assert repaired["request_id"] == request_id(request) + assert CorePredictionConfig.from_yaml(tmp_path / "predict/pilot.yaml").engine.model == request.identity.model + assert CoreRecommendationConfig.from_yaml(tmp_path / "recommend/pilot.yaml").engine.model == request.identity.model + assert json.loads((tmp_path / "commands.json").read_text())["fpm_run_local"][1:3] == ["onboard", "collect-fpm"] + assert (tmp_path / "systems/h200_sxm.yaml").is_file() + assert {name: (tmp_path / name).read_bytes() for name in before_repair} == before_repair + assert create_plan(request, tmp_path, overwrite=True) == repaired + + +def test_plan_can_restart_after_interruption_before_saved_request(tmp_path, monkeypatch): + original_open = Path.open + + def interrupted_open(path, mode="r", *args, **kwargs): + if path == tmp_path / "request.yaml" and mode == "xb": + raise InterruptedError("interrupted plan generation") + return original_open(path, mode, *args, **kwargs) + + monkeypatch.setattr(Path, "open", interrupted_open) + with pytest.raises(InterruptedError): + create_plan(_request(), tmp_path) + monkeypatch.setattr(Path, "open", original_open) + + assert [path.name for path in tmp_path.iterdir()] == [".support.lock"] + assert create_plan(_request(), tmp_path)["request_id"] == request_id(_request()) + + +@pytest.mark.parametrize("metadata", ["invalid JSON", "{}", '{"request_id": "another-request"}']) +def test_overwrite_rejects_existing_invalid_metadata_before_repair(tmp_path, metadata): + request = _request() + create_plan(request, tmp_path) + (tmp_path / "support-plan.json").write_text(metadata) + (tmp_path / "predict/pilot.yaml").unlink() + before_repair = _file_contents(tmp_path) + + with pytest.raises(ValueError, match="readable onboarding plan|request identity"): + create_plan(request, tmp_path, overwrite=True) + + assert _file_contents(tmp_path) == before_repair + + +@pytest.mark.parametrize("saved_request", [None, "identity: broken\n"]) +def test_overwrite_requires_frozen_request_for_interrupted_plan(tmp_path, saved_request): + request = _request() + create_plan(request, tmp_path) + (tmp_path / "support-plan.json").unlink() + request_path = tmp_path / "request.yaml" + if saved_request is None: + request_path.unlink() + else: + request_path.write_text(saved_request) + before_repair = _file_contents(tmp_path) + + with pytest.raises(ValueError, match="saved request"): + create_plan(request, tmp_path, overwrite=True) + + assert _file_contents(tmp_path) == before_repair diff --git a/python/aisimulate/tests/unit/test_systems_path.py b/python/aisimulate/tests/unit/test_systems_path.py new file mode 100644 index 000000000..5e6870ab0 --- /dev/null +++ b/python/aisimulate/tests/unit/test_systems_path.py @@ -0,0 +1,563 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Public prediction and recommendation use an isolated local systems root.""" + +from __future__ import annotations + +import hashlib +import json +import subprocess +import sys +from copy import deepcopy +from importlib.resources import files +from pathlib import Path + +import pyarrow as pa +import pyarrow.parquet as pq +import pytest +import yaml +from pydantic import ValidationError + +from aisimulate.compiler import prediction_to_replay_spec +from aisimulate.config import CorePredictionConfig, CoreRecommendationConfig +from aisimulate.output import write_recommendations +from aisimulate.recommend import recommendation_to_sweeper, run_recommendation +from aisimulate.runner import EngineReplayRunnerFactory +from aisimulate.sweeper.config import SearchSpace +from aisimulate_core.sdk import ForwardPassPerfModelConfig, RustForwardPassPerfModel, models, perf_database +from aisimulate_core.sdk.config import ModelConfig +from aisimulate_core.sdk.errors import PerfDataNotAvailableError +from aisimulate_core.sdk.operations.fpm_forward import _CELL_MATCH_COLUMNS + +pytestmark = pytest.mark.unit + +_SYSTEM = "local_fpm_test" +_VERSION = "0.0.1" + + +def _request() -> dict: + return { + "engine": { + "mode": "aggregated", + "model": "Qwen/Qwen3-0.6B", + "hardware": "h200_sxm", + "backend": "vllm", + "context_length": 128, + "workers": {"aggregated": {}}, + }, + } + + +@pytest.mark.parametrize("config_type", [CorePredictionConfig, CoreRecommendationConfig]) +@pytest.mark.parametrize("root_key", ["systems_path", "systems_paths"]) +def test_systems_root_is_absolute_and_reloadable_before_collection(config_type, root_key, tmp_path, monkeypatch): + monkeypatch.chdir(tmp_path) + raw = _request() + raw["engine"][root_key] = "local profiles" if root_key == "systems_path" else ["local profiles"] + if config_type is CoreRecommendationConfig: + raw["optimization"] = {"target": "throughput"} + + config = config_type.model_validate(raw) + saved = config.model_dump(mode="json", exclude_none=True) + monkeypatch.chdir(tmp_path.parent) + + assert "systems_path" not in saved["engine"] + assert saved["engine"]["systems_paths"] == [str(tmp_path / "local profiles")] + assert config_type.model_validate(saved).engine.systems_paths == config.engine.systems_paths + + +@pytest.mark.parametrize("config_type", [CorePredictionConfig, CoreRecommendationConfig]) +@pytest.mark.parametrize("scope", ["engine", "worker"]) +def test_canonical_systems_roots_preserve_order_and_default_on_reload(config_type, scope, tmp_path, monkeypatch): + monkeypatch.chdir(tmp_path) + raw = _request() + controls = raw["engine"] if scope == "engine" else raw["engine"]["workers"]["aggregated"].setdefault("timing", {}) + controls["systems_paths"] = ["second root", "DEFAULT", "~/first root", "with,comma", "second root"] + if config_type is CoreRecommendationConfig: + raw["optimization"] = {"target": "throughput"} + expected = [ + str(tmp_path / "second root"), + "default", + str(Path.home() / "first root"), + str(tmp_path / "with,comma"), + str(tmp_path / "second root"), + ] + config = config_type.model_validate(raw) + saved = config.model_dump(mode="json", exclude_none=True) + monkeypatch.chdir(tmp_path.parent) + + reloaded = config_type.model_validate(saved) + controls = reloaded.engine if scope == "engine" else reloaded.engine.workers.aggregated.timing + assert controls.systems_paths == expected + if config_type is CoreRecommendationConfig: + assert recommendation_to_sweeper(reloaded).search_space.systems_paths_for("agg") == expected + + +@pytest.mark.parametrize("config_type", [CorePredictionConfig, CoreRecommendationConfig]) +def test_systems_root_alias_accepts_equivalent_canonical_roots(config_type, tmp_path, monkeypatch): + monkeypatch.chdir(tmp_path) + raw = _request() + raw["engine"].update(systems_path="local profiles", systems_paths=["./local profiles"]) + if config_type is CoreRecommendationConfig: + raw["optimization"] = {"target": "throughput"} + + assert config_type.model_validate(raw).engine.systems_paths == [str(tmp_path / "local profiles")] + + +@pytest.mark.parametrize("root_key", ["systems_path", "systems_paths"]) +def test_search_space_roots_are_normalized_before_reload(root_key, tmp_path, monkeypatch): + monkeypatch.chdir(tmp_path) + roots = {root_key: "shared" if root_key == "systems_path" else ["shared", "default", "other"]} + space = SearchSpace( + model_name="model", + hardware_sku="system", + role_estimator_controls={"prefill": {"systems_paths": ["worker", "default"]}}, + **roots, + ) + saved = space.model_dump(mode="json", exclude_none=True) + monkeypatch.chdir(tmp_path.parent) + + reloaded = SearchSpace.model_validate(saved) + expected = [str(tmp_path / "shared")] + if root_key == "systems_paths": + expected.extend(["default", str(tmp_path / "other")]) + assert "systems_path" not in saved + assert reloaded.systems_paths_for("decode") == expected + assert reloaded.systems_paths_for("prefill") == [str(tmp_path / "worker"), "default"] + + +@pytest.mark.parametrize("config_type", [CorePredictionConfig, CoreRecommendationConfig]) +@pytest.mark.parametrize("value", ["", " ", 7, ["one"], {"choices": ["one"]}, "one,two"]) +def test_systems_root_rejects_empty_or_multiple_roots(config_type, value): + raw = _request() + raw["engine"]["systems_path"] = value + if config_type is CoreRecommendationConfig: + raw["optimization"] = {"target": "throughput"} + + with pytest.raises(ValidationError, match="systems_path"): + config_type.model_validate(raw) + + +@pytest.mark.parametrize("config_type", [CorePredictionConfig, CoreRecommendationConfig]) +def test_legacy_systems_root_rejects_conflicting_canonical_roots(config_type, tmp_path): + raw = _request() + raw["engine"].update(systems_path=str(tmp_path / "legacy"), systems_paths=[str(tmp_path / "canonical")]) + if config_type is CoreRecommendationConfig: + raw["optimization"] = {"target": "throughput"} + + with pytest.raises(ValidationError, match="systems_path conflicts with systems_paths"): + config_type.model_validate(raw) + + +@pytest.mark.parametrize("config_type", [CorePredictionConfig, CoreRecommendationConfig]) +@pytest.mark.parametrize("root_key", ["systems_path", "systems_paths"]) +@pytest.mark.parametrize("timing", [{"type": "fixed", "prefill_ms": 1, "decode_ms": 2}, {"type": "polynomial"}]) +def test_shared_roots_with_custom_timing_preserve_estimator_control_guards(config_type, root_key, timing, tmp_path): + raw = _request() + raw["engine"][root_key] = str(tmp_path) if root_key == "systems_path" else [str(tmp_path)] + raw["engine"]["workers"]["aggregated"]["timing"] = timing + if config_type is CoreRecommendationConfig: + raw["optimization"] = {"target": "throughput"} + + config = config_type.model_validate(raw) + saved = config.model_dump(mode="json", exclude_none=True) + assert saved["engine"]["systems_paths"] == [str(tmp_path)] + assert config_type.model_validate(saved).engine.workers.aggregated.timing.type == timing["type"] + + raw["engine"]["estimation_mode"] = "fpm_interpolation" + with pytest.raises(ValidationError, match="estimator policies require.*default timing"): + config_type.model_validate(raw) + + +@pytest.mark.parametrize("config_type", [CorePredictionConfig, CoreRecommendationConfig]) +@pytest.mark.parametrize("unsupported", ["AFD", "analytical encoder pools"]) +def test_systems_root_rejects_modes_that_cannot_consume_it(config_type, unsupported, tmp_path): + raw = _request() + if unsupported == "AFD": + raw["engine"]["mode"] = "afd" + raw["engine"]["workers"] = {} + raw["engine"]["afd"] = {"phase": "both", "combined_with_pd": False, "a_batch_size": 1} + if config_type is CorePredictionConfig: + raw["engine"]["afd"].update(n_a_nodes=1, n_f_nodes=1, tp_a=1) + else: + raw["engine"]["workers"]["encoder"] = {} + raw["traffic"] = { + "source": {"type": "synthetic", "images": {"height": 32, "width": 32}}, + "load": {"type": "concurrency", "concurrency": 1}, + "stop": {"requests": 1}, + } + if config_type is CoreRecommendationConfig: + raw["optimization"] = {"target": "throughput"} + + config_type.model_validate(raw) + raw["engine"]["systems_path"] = str(tmp_path) + with pytest.raises(ValidationError, match=f"systems_path.*{unsupported}"): + config_type.model_validate(raw) + + +@pytest.fixture +def local_profiles(tmp_path): + """Two roots share an identity but have distinct timings and GPU memory.""" + model = models.get_model("Qwen/Qwen3-0.6B", ModelConfig(forward_model="fpm"), "vllm") + identity = dict(zip(_CELL_MATCH_COLUMNS, model.context_ops[0]._match_identity, strict=True)) + for name in ("tp", "pp", "dp", "moe_tp", "moe_ep", "cp"): + identity[name] = int(identity[name]) + for name in ("enable_wideep", "enable_eplb"): + identity[name] = identity[name] == "True" + roots = [] + for label, latency, memory in (("fast", 20.0, 16 << 30), ("slow", 200.0, 32 << 30)): + root = tmp_path / label + data = root / "data" / _SYSTEM / "vllm" / _VERSION + data.mkdir(parents=True) + system = yaml.safe_load((files("aisimulate_core") / "systems/h200_sxm.yaml").read_text()) + system["data_dir"] = f"data/{_SYSTEM}" + system["gpu"]["mem_capacity"] = memory + (root / f"{_SYSTEM}.yaml").write_text(yaml.safe_dump(system)) + rows = [] + for phase, tokens, kv in (("prefill", 16, 0), ("decode", 0, 0), ("decode", 0, 128)): + rows.append( + identity + | { + "cell_id": phase, + "model_path": model.model_path, + "system": _SYSTEM, + "backend": "vllm", + "backend_version": _VERSION, + "weight_quantization": identity["gemm_quant_mode"], + "workload_kind": phase, + "batch_size": 1, + "total_prefill_tokens": tokens, + "total_kv_read_tokens": kv, + "partition_policy": "balanced_v1", + "latency_ms": latency, + } + ) + parquet = data / "fpm_forward_perf.parquet" + pq.write_table(pa.Table.from_pylist(rows), parquet) + (data / "fpm_forward_perf.metadata.json").write_text( + json.dumps( + { + "schema_name": "aic_fpm_forward_perf", + "schema_version": 6, + "coordinate_system": "iteration_totals_balanced_v1", + "measurement_policy": "dynamo_native_single_sample_v1", + "row_count": len(rows), + "parquet_sha256": hashlib.sha256(parquet.read_bytes()).hexdigest(), + "system": _SYSTEM, + "backend": "vllm", + "backend_version": _VERSION, + } + ) + ) + roots.append(root) + return roots + + +def _local_request(root, mode="aggregated") -> dict: + raw = _request() + raw["engine"].update(mode=mode, hardware=_SYSTEM, systems_path=str(root), backend_version=_VERSION) + worker = { + "scheduler": {"max_batched_tokens": 16, "max_sequences": 1}, + "kv_cache": {"prefix_caching": False}, + "timing": {"forward_model": "fpm"}, + } + roles = ("aggregated",) if mode == "aggregated" else ("prefill", "decode") + raw["engine"]["workers"] = {role: deepcopy(worker) for role in roles} + raw["traffic"] = { + "source": {"type": "synthetic", "input_tokens": 16, "output_tokens": 2}, + "load": {"type": "concurrency", "concurrency": 1}, + "stop": {"requests": 1}, + } + return raw + + +def _predict(raw): + spec = prediction_to_replay_spec(CorePredictionConfig.model_validate(raw)) + runner = EngineReplayRunnerFactory().create(0) + try: + return runner.run(spec) + finally: + runner.close() + + +@pytest.mark.parametrize("mode", ["aggregated", "disaggregated"]) +@pytest.mark.parametrize("engine_roots", [False, True]) +def test_worker_timing_roots_match_runtime_and_metadata(local_profiles, mode, engine_roots): + fast, slow = local_profiles + raw = _local_request(fast, mode) + if not engine_roots: + raw["engine"].pop("systems_path") + for worker in raw["engine"]["workers"].values(): + worker["timing"]["systems_paths"] = [str(slow)] + + config = CorePredictionConfig.model_validate(raw) + saved = config.model_dump(mode="json", exclude_none=True) + deployment = prediction_to_replay_spec(CorePredictionConfig.model_validate(saved)).backend_deployment + for metadata in deployment.performance_model_metadata.values(): + assert metadata["config"]["systems_paths"] == [str(slow)] + report = _predict(saved) + assert report.metrics["completed_requests"] == 1 + assert report.metrics["mean_ttft_ms"] == pytest.approx(200.0 if mode == "aggregated" else 400.0) + assert report.metrics["mean_itl_ms"] == pytest.approx(200.0) + + +@pytest.mark.parametrize("root_key", ["systems_path", "systems_paths"]) +def test_reloaded_prediction_consumes_relative_systems_roots(local_profiles, root_key, tmp_path, monkeypatch): + fast, slow = local_profiles + monkeypatch.chdir(fast.parent) + raw = _local_request(fast) + raw["engine"].pop("systems_path") + raw["engine"][root_key] = fast.name if root_key == "systems_path" else [fast.name, slow.name] + saved = CorePredictionConfig.model_validate(raw).model_dump(mode="json", exclude_none=True) + other_directory = tmp_path / "other" + other_directory.mkdir() + monkeypatch.chdir(other_directory) + + report = _predict(saved) + + assert report.metrics["completed_requests"] == 1 + # The first token is emitted by prefill; only later tokens need decode. + assert report.metrics["mean_ttft_ms"] == pytest.approx(20.0) + assert report.metrics["mean_itl_ms"] == pytest.approx(20.0) + deployment = prediction_to_replay_spec(CorePredictionConfig.model_validate(saved)).backend_deployment + expected = [str(fast)] if root_key == "systems_path" else [str(fast), str(slow)] + assert deployment.performance_model_metadata["aggregated"]["config"]["systems_paths"] == expected + + +@pytest.mark.parametrize("mode", ["aggregated", "disaggregated"]) +def test_prediction_consumes_local_fpm_timings_without_changing_default_roots(local_profiles, mode): + default_roots = perf_database.get_systems_paths() + fast, slow = local_profiles + + fast_report = _predict(_local_request(fast, mode)) + slow_report = _predict(_local_request(slow, mode)) + repeat_report = _predict(_local_request(fast, mode)) + + assert fast_report.metrics["completed_requests"] == slow_report.metrics["completed_requests"] == 1 + assert slow_report.metrics["mean_ttft_ms"] > 5 * fast_report.metrics["mean_ttft_ms"] + assert repeat_report.metrics["duration_ms"] == fast_report.metrics["duration_ms"] + assert perf_database.get_systems_paths() == default_roots + deployment = prediction_to_replay_spec( + CorePredictionConfig.model_validate(_local_request(fast, mode)) + ).backend_deployment + assert deployment.performance_model_metadata + for metadata in deployment.performance_model_metadata.values(): + assert metadata["config"]["systems_paths"] == [str(fast)] + + +@pytest.mark.parametrize("root_key", ["systems_path", "systems_paths"]) +@pytest.mark.parametrize("timing", [{"type": "fixed", "prefill_ms": 1, "decode_ms": 2}, {"type": "polynomial"}]) +@pytest.mark.parametrize("fixed_capacity", [False, True]) +def test_custom_timing_can_use_local_profiles_for_kv_capacity(local_profiles, timing, root_key, fixed_capacity): + capacities = [] + for root in local_profiles: + raw = _local_request(root) + worker = raw["engine"]["workers"]["aggregated"] + worker["timing"] = timing + if fixed_capacity: + worker["kv_cache"]["capacity"] = {"type": "fixed", "blocks": 128} + if root_key == "systems_paths": + raw["engine"]["systems_paths"] = [raw["engine"].pop("systems_path")] + + deployment = prediction_to_replay_spec(CorePredictionConfig.model_validate(raw)).backend_deployment + assert "systems_path" not in deployment.agg_engine_args + assert deployment.performance_model_metadata["aggregated"]["config"]["systems_paths"] == [str(root)] + capacities.append(deployment.agg_engine_args["num_gpu_blocks"]) + assert _predict(raw).metrics["completed_requests"] == 1 + + if fixed_capacity: + assert capacities == [128, 128] + else: + assert capacities[1] > capacities[0] > 0 + + +@pytest.mark.parametrize("backend_version", [None, _VERSION]) +def test_canonical_estimator_resolves_no_slot_local_version(local_profiles, backend_version): + root = str(local_profiles[0]) + assert perf_database.get_version_slots(_SYSTEM, "vllm", systems_paths=root) is None + request = ForwardPassPerfModelConfig( + model="Qwen/Qwen3-0.6B", + system=_SYSTEM, + backend="vllm", + backend_version=backend_version, + worker_type="aggregated", + kv_block_size=64, + estimation_mode="fpm_interpolation", + systems_paths=(root,), + ) + model = RustForwardPassPerfModel.best_available(request) + try: + diagnostics = model.diagnostics() + assert diagnostics["readiness"] == "ready" + provenance = diagnostics["provenance"] + assert provenance["selected_systems_root"] == root + assert provenance["config"]["backend_version"] == _VERSION + assert provenance["config"]["systems_paths"] == [root] + assert provenance["config"]["estimation_mode"] == "fpm_interpolation" + assert provenance["config"]["fallback_policy"] == "deny" + assert model.estimate_forward_pass_time_ms( + {"scheduled_requests": {"num_decode_requests": 1, "sum_decode_kv_tokens": 0}} + ) == pytest.approx(20.0) + finally: + model.close() + + +@pytest.mark.parametrize("backend_version", [None, _VERSION]) +def test_canonical_estimator_rejects_no_slot_root_without_data(local_profiles, tmp_path, backend_version): + root = tmp_path / "uncollected" + root.mkdir() + (root / f"{_SYSTEM}.yaml").write_text((local_profiles[0] / f"{_SYSTEM}.yaml").read_text()) + assert perf_database.get_latest_database_version(_SYSTEM, "vllm", systems_paths=str(root)) is None + request = ForwardPassPerfModelConfig( + model="Qwen/Qwen3-0.6B", + system=_SYSTEM, + backend="vllm", + backend_version=backend_version, + worker_type="aggregated", + estimation_mode="fpm_interpolation", + systems_paths=(str(root),), + ) + error_type = ValueError if backend_version is None else PerfDataNotAvailableError + message = "no backend version is available" if backend_version is None else "perf data directory not found" + with pytest.raises(error_type, match=message): + RustForwardPassPerfModel.best_available(request) + + +@pytest.mark.parametrize("timing", [{"type": "fixed", "prefill_ms": 1, "decode_ms": 2}, {"type": "polynomial"}]) +@pytest.mark.parametrize("fixed_capacity", [False, True]) +def test_custom_timing_recommendation_preserves_local_roots(local_profiles, timing, fixed_capacity, tmp_path): + root = local_profiles[0] + raw = _local_request(root) + raw["engine"]["systems_paths"] = [raw["engine"].pop("systems_path")] + worker = raw["engine"]["workers"]["aggregated"] + worker["timing"] = timing + worker["parallelism"] = { + "preset": [{"replicas": 1, "tensor": 1, "pipeline": 1, "attention_data": 1, "moe_tensor": 1, "moe_expert": 1}] + } + if fixed_capacity: + worker["kv_cache"]["capacity"] = {"type": "fixed", "blocks": 128} + raw["optimization"] = {"target": "throughput", "constraints": {"max_candidate_gpus": 1}} + raw["optimizer"] = {"algorithm": "random", "max_trials": 1, "parallelism": 1} + + result = run_recommendation( + CoreRecommendationConfig.model_validate(raw), + stack="engine", + runner_factory=EngineReplayRunnerFactory(), + show_progress=False, + ) + + assert result.counts.feasible == 1, [record.reason for record in result.candidates] + [candidate] = result.selected_candidates + assert candidate.config["systems_paths"] == [str(root)] + assert candidate.prediction_config["engine"]["systems_paths"] == [str(root)] + for metadata in result.candidates[0].provenance.performance_data: + assert metadata["config"]["systems_paths"] == [str(root)] + [saved] = write_recommendations(tmp_path / "custom-timing-export", [candidate.prediction_config]) + prediction = CorePredictionConfig.from_yaml(saved) + assert prediction.engine.workers.aggregated.timing.model_dump(exclude_none=True) == timing + assert _predict(prediction.model_dump(mode="json", exclude_none=True)).metrics["completed_requests"] == 1 + + +@pytest.mark.parametrize("explicit_version", [False, True]) +def test_worker_hardware_override_resolves_in_local_systems_root(local_profiles, explicit_version): + raw = _local_request(local_profiles[0], "disaggregated") + raw["engine"]["hardware"] = "h200_sxm" + for worker in raw["engine"]["workers"].values(): + worker["hardware"] = _SYSTEM + if not explicit_version: + raw["engine"].pop("backend_version") + + report = _predict(raw) + + assert report.metrics["completed_requests"] == 1 + # The first token needs a 20 ms prefill and a 20 ms decode; the next needs only decode. + assert report.metrics["mean_ttft_ms"] == pytest.approx(40.0) + assert report.metrics["mean_itl_ms"] == pytest.approx(20.0) + deployment = prediction_to_replay_spec(CorePredictionConfig.model_validate(raw)).backend_deployment + assert deployment.backend_version == _VERSION + + +def test_default_prediction_still_uses_bundled_op_level_data_after_a_local_request(local_profiles): + local = _local_request(local_profiles[0]) + default = deepcopy(local) + default["engine"].pop("systems_path") + default["engine"]["hardware"] = "h200_sxm" + default["engine"]["backend_version"] = perf_database.get_latest_database_version("h200_sxm", "vllm") + default["engine"]["workers"]["aggregated"].pop("timing") + + before = _predict(default) + assert _predict(local).metrics["completed_requests"] == 1 + after = _predict(default) + + assert before.metrics["completed_requests"] == after.metrics["completed_requests"] == 1 + assert after.metrics["duration_ms"] == before.metrics["duration_ms"] + config = CorePredictionConfig.model_validate(default) + assert "systems_path" not in config.model_dump(mode="json", exclude_none=True)["engine"] + assert config.engine.workers.aggregated.timing.forward_model == "op_level" + + +@pytest.mark.parametrize("mode", ["aggregated", "disaggregated"]) +def test_recommendation_resolves_local_capacity_and_exports_replayable_profiles(local_profiles, mode, tmp_path): + default_roots = perf_database.get_systems_paths() + candidates = [] + for root in local_profiles: + raw = _local_request(root, mode) + del raw["engine"]["backend_version"] + raw["traffic"]["load"] = {"type": "kv_capacity_fraction", "fraction": 0.000001} + raw["optimization"] = { + "target": "throughput", + "constraints": {"max_candidate_gpus": 1 if mode == "aggregated" else 2}, + } + raw["optimizer"] = {"algorithm": "random", "max_trials": 1, "parallelism": 1} + for worker in raw["engine"]["workers"].values(): + worker["kv_cache"].update(block_size=64, capacity={"memory_fraction": 0.9}) + worker["parallelism"] = { + "preset": [ + {"replicas": 1, "tensor": 1, "pipeline": 1, "attention_data": 1, "moe_tensor": 1, "moe_expert": 1} + ] + } + + result = run_recommendation( + CoreRecommendationConfig.model_validate(raw), + stack="engine", + runner_factory=EngineReplayRunnerFactory(), + show_progress=False, + ) + + assert result.counts.feasible == 1, [record.reason for record in result.candidates] + candidate = result.selected_candidates[0] + assert candidate.config["systems_paths"] == [str(root)] + assert candidate.prediction_config["engine"]["systems_paths"] == [str(root)] + assert candidate.prediction_config["engine"]["backend_version"] == _VERSION + assert result.candidates[0].provenance.performance_data + for metadata in result.candidates[0].provenance.performance_data: + assert metadata["config"]["systems_paths"] == [str(root)] + [saved] = write_recommendations(tmp_path / f"export-{root.name}", [candidate.prediction_config]) + rerun = subprocess.run( + [ + sys.executable, + "-m", + "aisimulate", + "predict", + "--config", + str(saved), + "--output-dir", + str(tmp_path / f"predict-{root.name}"), + "--format", + "json", + ], + cwd=tmp_path, + capture_output=True, + text=True, + timeout=60, + check=False, + ) + assert rerun.returncode == 0, rerun.stderr + assert json.loads(rerun.stdout)["mean_ttft_ms"] == candidate.metrics["mean_ttft_ms"] + candidates.append(candidate) + + fast, slow = candidates + assert slow.metrics["mean_ttft_ms"] > 5 * fast.metrics["mean_ttft_ms"] + assert slow.config["kv_load_capacity_tokens"] > fast.config["kv_load_capacity_tokens"] + assert perf_database.get_systems_paths() == default_roots diff --git a/python/aisimulate/tools/support_matrix/compare_support_matrix.py b/python/aisimulate/tools/support_matrix/compare_support_matrix.py index 6cceaf231..585e643c8 100644 --- a/python/aisimulate/tools/support_matrix/compare_support_matrix.py +++ b/python/aisimulate/tools/support_matrix/compare_support_matrix.py @@ -27,9 +27,10 @@ from pathlib import Path # Ensure local repo paths are importable when running as a standalone script. -_REPO_ROOT = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) -sys.path.insert(0, os.path.join(_REPO_ROOT, "src")) -sys.path.insert(0, _REPO_ROOT) +if __name__ == "__main__": + _REPO_ROOT = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) + sys.path.insert(0, os.path.join(_REPO_ROOT, "src")) + sys.path.insert(0, _REPO_ROOT) from tools.support_matrix.support_matrix import ( ENCODER_DATA_UNAVAILABLE_PREFIX, diff --git a/python/aisimulate/tools/support_matrix/generate_support_matrix.py b/python/aisimulate/tools/support_matrix/generate_support_matrix.py index b71e32b66..0578c85cc 100644 --- a/python/aisimulate/tools/support_matrix/generate_support_matrix.py +++ b/python/aisimulate/tools/support_matrix/generate_support_matrix.py @@ -15,9 +15,10 @@ import sys # Ensure local repo paths are importable when running as a standalone script. -_REPO_ROOT = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) -sys.path.insert(0, os.path.join(_REPO_ROOT, "src")) -sys.path.insert(0, _REPO_ROOT) +if __name__ == "__main__": + _REPO_ROOT = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) + sys.path.insert(0, os.path.join(_REPO_ROOT, "src")) + sys.path.insert(0, _REPO_ROOT) from tools.support_matrix.support_matrix import ( VALID_STATUSES, diff --git a/tests/sweeper/test_epd.py b/tests/sweeper/test_epd.py index 3c406c5e4..0522f633e 100644 --- a/tests/sweeper/test_epd.py +++ b/tests/sweeper/test_epd.py @@ -285,7 +285,7 @@ def test_complete_sweeper_selection_and_serialization(monkeypatch, gpu_budget, f catalog = {"one": _encoder(), "two": _encoder(workers=2)} monkeypatch.setattr(search_mod, "enumerate_branches", lambda *a, **kw: [branch]) monkeypatch.setattr(search_mod, "resolve_encoder_catalog", lambda c: catalog) - monkeypatch.setattr(search_mod, "resolve_backend_version", lambda *a: "0.5.14") + monkeypatch.setattr(search_mod, "resolve_backend_version", lambda *a, systems_paths=None: "0.5.14") if failure in {"build", "version", "kv"}: def fail_build(*args, **kwargs): @@ -443,10 +443,10 @@ def test_encoder_catalog_preserves_version_precedence( version_calls = [] database_calls = [] - def latest_version(system, backend): + def latest_version(system, backend, systems_paths=None): assert implicit, "explicit encoder versions must not depend on latest-version availability" version_calls.append((system, backend)) - return original_version(system, backend) + return original_version(system, backend, systems_paths=systems_paths) def database(system, backend, version, **kwargs): database_calls.append((system, backend, version)) diff --git a/tests/sweeper/test_model_hw.py b/tests/sweeper/test_model_hw.py index c1cb64f24..629151045 100644 --- a/tests/sweeper/test_model_hw.py +++ b/tests/sweeper/test_model_hw.py @@ -59,7 +59,7 @@ def test_aic_core_system_spec_contract(monkeypatch): monkeypatch.setattr( mh_mod.perf_database, "load_system_spec", - lambda hardware: { + lambda hardware, systems_paths=None: { "gpu": {"mem_capacity": 80}, "node": {"num_gpus_per_node": 8}, }, diff --git a/tests/sweeper/test_result.py b/tests/sweeper/test_result.py index 8ac3c1468..c2d906a25 100644 --- a/tests/sweeper/test_result.py +++ b/tests/sweeper/test_result.py @@ -552,7 +552,7 @@ def test_optimizer_guided_run_emits_complete_ledger_and_top_n(monkeypatch): monkeypatch.setattr( search_module, "resolve_backend_version", - lambda hardware, backend: "1.0", + lambda hardware, backend, systems_paths=None: "1.0", ) result = Sweeper( @@ -611,7 +611,7 @@ def test_strict_sla_rejection_is_preserved_in_the_candidate_ledger(monkeypatch): monkeypatch.setattr( search_module, "resolve_backend_version", - lambda hardware, backend: "1.0", + lambda hardware, backend, systems_paths=None: "1.0", ) config_data = _config().model_dump(mode="python") config_data["goal"] = { @@ -643,7 +643,7 @@ def test_same_batch_failed_duplicates_are_counted_as_coalesced_hits(monkeypatch) knob_choices={"backend": ["trtllm"]}, ) monkeypatch.setattr(search_module, "enumerate_branches", lambda *args, **kwargs: [branch]) - monkeypatch.setattr(search_module, "resolve_backend_version", lambda *args: "1.0") + monkeypatch.setattr(search_module, "resolve_backend_version", lambda *args, systems_paths=None: "1.0") result = Sweeper( runner_factory=_FailingRunnerFactory(), @@ -669,7 +669,7 @@ def test_zero_or_missing_sample_latency_preserves_ranked_sampler_feedback( knob_choices={"backend": ["trtllm"]}, ) monkeypatch.setattr(search_module, "enumerate_branches", lambda *args, **kwargs: [branch]) - monkeypatch.setattr(search_module, "resolve_backend_version", lambda *args: "1.0") + monkeypatch.setattr(search_module, "resolve_backend_version", lambda *args, systems_paths=None: "1.0") seen = {} @@ -726,7 +726,7 @@ def test_optimizer_guided_result_separates_unsupported_and_runtime_failure(monke monkeypatch.setattr( search_module, "resolve_backend_version", - lambda hardware, backend: "1.0", + lambda hardware, backend, systems_paths=None: "1.0", ) result = Sweeper( @@ -798,7 +798,7 @@ def test_sweep_keeps_feasible_result_when_another_candidate_exceeds_resources(mo knob_choices={"backend": ["trtllm"]}, ) monkeypatch.setattr(search_module, "enumerate_branches", lambda *a, **kw: [branch]) - monkeypatch.setattr(search_module, "resolve_backend_version", lambda *a: "1.0") + monkeypatch.setattr(search_module, "resolve_backend_version", lambda *a, systems_paths=None: "1.0") class ResourceFactory(_RunnerFactory): def admit_wave(self, specs): diff --git a/tests/sweeper/test_search.py b/tests/sweeper/test_search.py index 67756fe4e..501dcbb5c 100644 --- a/tests/sweeper/test_search.py +++ b/tests/sweeper/test_search.py @@ -155,7 +155,7 @@ def _stub(monkeypatch, branch): "enumerate_branches", lambda config, *, max_seq_len=None, runner_capabilities=None: [branch], ) - monkeypatch.setattr(search_mod, "resolve_backend_version", lambda hw, be: "1.3.0rc10") + monkeypatch.setattr(search_mod, "resolve_backend_version", lambda hw, be, systems_paths=None: "1.3.0rc10") def _pc(*, tp=4, replicas=2): @@ -307,7 +307,7 @@ def test_heterogeneous_disagg_materializes_role_hardware_and_provenance(monkeypa ) version_calls = [] - def resolve_version(hardware, backend): + def resolve_version(hardware, backend, *, systems_paths=None): version_calls.append((hardware, backend)) return "0.24.0" @@ -1097,7 +1097,7 @@ def test_projection_stall_only_stops_current_branch(monkeypatch): "enumerate_branches", lambda config, *, max_seq_len=None, runner_capabilities=None: [agg, disagg], ) - monkeypatch.setattr(search_mod, "resolve_backend_version", lambda hw, be: "1.3.0rc10") + monkeypatch.setattr(search_mod, "resolve_backend_version", lambda hw, be, systems_paths=None: "1.3.0rc10") seen = [] class RepeatingSampler(_FakeSampler): diff --git a/tests/sweeper/test_search_providers.py b/tests/sweeper/test_search_providers.py index c520da847..e74a28ff3 100644 --- a/tests/sweeper/test_search_providers.py +++ b/tests/sweeper/test_search_providers.py @@ -249,7 +249,7 @@ def _stub_branch(monkeypatch) -> None: monkeypatch.setattr( search_module, "resolve_backend_version", - lambda hardware, backend: "0.11.0", + lambda hardware, backend, systems_paths=None: "0.11.0", ) @@ -377,7 +377,7 @@ def materialize_replay(self, plan, selection, context): monkeypatch.setattr( search_module, "resolve_backend_version", - lambda hardware, backend: "0.11.0", + lambda hardware, backend, systems_paths=None: "0.11.0", ) prepared, result = search_module._materialize_one( { diff --git a/tests/sweeper/test_unified_optimizer.py b/tests/sweeper/test_unified_optimizer.py index a61339539..d45a7fb9c 100644 --- a/tests/sweeper/test_unified_optimizer.py +++ b/tests/sweeper/test_unified_optimizer.py @@ -148,7 +148,7 @@ def test_global_trial_budget_is_split_across_branches(monkeypatch) -> None: _CountingSampler.suggestion_batches = [] _CountingSampler.suggested = 0 monkeypatch.setattr(search_module, "enumerate_branches", lambda *args, **kwargs: _branches()) - monkeypatch.setattr(search_module, "resolve_backend_version", lambda *args: "test") + monkeypatch.setattr(search_module, "resolve_backend_version", lambda *args, systems_paths=None: "test") config = SmartSearchConfig.model_validate( { "search_space": { @@ -193,7 +193,7 @@ def test_global_trial_budget_runs_branch_batches_round_robin(monkeypatch) -> Non _CountingSampler.suggestion_batches = [] _CountingSampler.suggested = 0 monkeypatch.setattr(search_module, "enumerate_branches", lambda *args, **kwargs: _branches()) - monkeypatch.setattr(search_module, "resolve_backend_version", lambda *args: "test") + monkeypatch.setattr(search_module, "resolve_backend_version", lambda *args, systems_paths=None: "test") config = SmartSearchConfig.model_validate( { "search_space": { @@ -243,7 +243,7 @@ def test_legacy_rounds_remain_branch_major_without_max_trials(monkeypatch) -> No _CountingSampler.suggestion_batches = [] _CountingSampler.suggested = 0 monkeypatch.setattr(search_module, "enumerate_branches", lambda *args, **kwargs: _branches()) - monkeypatch.setattr(search_module, "resolve_backend_version", lambda *args: "test") + monkeypatch.setattr(search_module, "resolve_backend_version", lambda *args, systems_paths=None: "test") config = SmartSearchConfig.model_validate( { "search_space": { @@ -288,7 +288,7 @@ def test_branch_seed_is_stable_when_branch_order_changes(monkeypatch) -> None: "enumerate_branches", lambda *args, **kwargs: list(reversed(_branches())), ) - monkeypatch.setattr(search_module, "resolve_backend_version", lambda *args: "test") + monkeypatch.setattr(search_module, "resolve_backend_version", lambda *args, systems_paths=None: "test") config = SmartSearchConfig.model_validate( { "search_space": { @@ -353,7 +353,7 @@ def test_candidate_timeout_applies_with_parallelism_one(monkeypatch) -> None: "enumerate_branches", lambda *args, **kwargs: [_branches()[0]], ) - monkeypatch.setattr(search_module, "resolve_backend_version", lambda *args: "test") + monkeypatch.setattr(search_module, "resolve_backend_version", lambda *args, systems_paths=None: "test") config = SmartSearchConfig.model_validate( { "search_space": { diff --git a/tests/test_epd_cli.py b/tests/test_epd_cli.py index d1f5c4bf0..ffb71210c 100644 --- a/tests/test_epd_cli.py +++ b/tests/test_epd_cli.py @@ -490,7 +490,7 @@ def database(system, backend, version, **kwargs): return None if backend == "vllm" else original_database(system, backend, version, **kwargs) def latest_version(system, backend, **kwargs): - return None if backend == "vllm" else original_version(system, backend) + return None if backend == "vllm" else original_version(system, backend, **kwargs) raw = _recommendation() raw["engine"]["workers"]["encoder"]["replicas"] = 1