From ba3f23fc4990092ad581723ab7b07cb0f56c5c74 Mon Sep 17 00:00:00 2001 From: Ignatious Johnson Date: Mon, 31 Aug 2026 19:52:11 -0400 Subject: [PATCH] Reorganize Pydantic config schemas under cvs/schema/. Move config models from parsers/schemas.py and training loaders into cvs/schema/ mirroring cvs/input/, add validate_config_file for supported suites, and co-located schema unit tests. Co-authored-by: Cursor Signed-off-by: Ignatious Johnson --- AGENTS.md | 2 +- cvs/lib/inference/ADDING_A_SUITE.md | 9 +- cvs/lib/inference/atom/atom_config_loader.py | 315 +--- .../inference/sglang/sglang_config_loader.py | 71 +- .../unittests/test_atom_config_loader.py | 2 +- .../test_inferencing_config_loader.py | 407 ----- cvs/lib/inference/utils/AGENTS.md | 128 +- cvs/lib/inference/utils/accuracy_config.py | 47 +- cvs/lib/inference/utils/functional_config.py | 16 +- .../utils/inferencing_config_loader.py | 224 --- .../utils/long_context_accuracy_config.py | 42 +- cvs/lib/inference/utils/platform_config.py | 16 +- cvs/lib/inference/utils/vllm_config_loader.py | 278 +-- .../unittests/test_ifoe_l2_connectivity.py | 2 +- .../preflight/unittests/test_node_smoke.py | 4 +- .../unittests/test_rdma_connectivity.py | 2 +- .../unittests/test_scaleup_fabric.py | 2 +- .../unittests/test_transferbench_smoke.py | 2 +- .../utils/training_config_loader.py | 344 +--- .../megatron/utils/training_config_loader.py | 211 +-- .../utils/training_config_loader.py | 208 +-- cvs/lib/unittests/test_utils_lib.py | 2 +- cvs/lib/utils/AGENTS.md | 3 +- cvs/lib/utils/config_loader.py | 113 +- cvs/parsers/__init__.py | 32 +- cvs/parsers/schemas.py | 1548 +---------------- cvs/schema/__init__.py | 40 +- cvs/schema/base.py | 11 + cvs/schema/cluster_file/__init__.py | 0 cvs/schema/cluster_file/cluster.py | 171 ++ cvs/schema/cluster_file/unittests/__init__.py | 0 .../cluster_file/unittests/test_cluster.py | 51 + cvs/schema/common/__init__.py | 17 + cvs/schema/common/base.py | 78 + cvs/schema/config_file/__init__.py | 0 cvs/schema/config_file/aorta/__init__.py | 0 cvs/schema/config_file/aorta/benchmark.py | 207 +++ .../config_file/aorta/unittests/__init__.py | 0 .../aorta/unittests/test_benchmark.py | 33 + cvs/schema/config_file/inference/__init__.py | 0 .../config_file/inference/atom/__init__.py | 0 .../inference/atom/unittests/__init__.py | 0 .../inference/atom/unittests/test_variant.py | 582 +++++++ .../config_file/inference/atom/variant.py | 244 +++ .../config_file/inference/common/__init__.py | 0 .../config_file/inference/common/accuracy.py | 34 + .../inference/common/functional.py | 10 + .../inference/common/long_context_accuracy.py | 31 + .../config_file/inference/common/platform.py | 10 + .../config_file/inference/common/sweep.py | 104 ++ .../inference/common/unittests/__init__.py | 0 .../common/unittests/test_accuracy.py | 496 ++++++ .../common/unittests/test_helpers.py | 41 + .../inference/common/unittests/test_sweep.py | 116 ++ .../inference/pytorch_xdit/__init__.py | 0 .../inference/pytorch_xdit/config.py | 266 +++ .../pytorch_xdit/unittests/__init__.py | 0 .../pytorch_xdit/unittests/test_config.py | 79 + .../config_file/inference/sglang/__init__.py | 0 .../inference/sglang/unittests/__init__.py | 0 .../sglang/unittests/test_variant.py | 80 + .../config_file/inference/sglang/variant.py | 66 + .../config_file/inference/vllm/__init__.py | 0 .../inference/vllm/unittests/__init__.py | 0 .../inference/vllm/unittests/test_variant.py | 168 ++ .../config_file/inference/vllm/variant.py | 139 ++ cvs/schema/config_file/preflight/__init__.py | 0 cvs/schema/config_file/preflight/config.py | 637 +++++++ .../preflight/unittests/__init__.py | 0 .../preflight/unittests/test_config.py | 317 ++++ cvs/schema/config_file/training/__init__.py | 0 .../training/jaxmaxtext/__init__.py | 0 .../training/jaxmaxtext/unittests/__init__.py | 0 .../jaxmaxtext/unittests/test_variant.py | 103 ++ .../training/jaxmaxtext/variant.py | 180 ++ .../config_file/training/megatron/__init__.py | 0 .../training/megatron/unittests/__init__.py | 0 .../megatron/unittests/test_variant.py | 95 + .../config_file/training/megatron/variant.py | 141 ++ .../training/torchtitan/__init__.py | 0 .../training/torchtitan/unittests/__init__.py | 0 .../torchtitan/unittests/test_variant.py | 80 + .../training/torchtitan/variant.py | 141 ++ cvs/schema/unittests/__init__.py | 0 cvs/schema/unittests/test_validate.py | 92 + cvs/schema/validate.py | 147 ++ cvs/tests/benchmark/test_aorta.py | 9 +- .../xdit/pytorch_xdit_flux_dev_single.py | 3 +- .../xdit/pytorch_xdit_wan22_14b_single.py | 3 +- cvs/tests/preflight/README.md | 2 +- cvs/tests/preflight/preflight_checks.py | 4 +- 91 files changed, 5229 insertions(+), 3829 deletions(-) delete mode 100644 cvs/lib/inference/unittests/test_inferencing_config_loader.py delete mode 100644 cvs/lib/inference/utils/inferencing_config_loader.py create mode 100644 cvs/schema/base.py create mode 100644 cvs/schema/cluster_file/__init__.py create mode 100644 cvs/schema/cluster_file/cluster.py create mode 100644 cvs/schema/cluster_file/unittests/__init__.py create mode 100644 cvs/schema/cluster_file/unittests/test_cluster.py create mode 100644 cvs/schema/common/__init__.py create mode 100644 cvs/schema/common/base.py create mode 100644 cvs/schema/config_file/__init__.py create mode 100644 cvs/schema/config_file/aorta/__init__.py create mode 100644 cvs/schema/config_file/aorta/benchmark.py create mode 100644 cvs/schema/config_file/aorta/unittests/__init__.py create mode 100644 cvs/schema/config_file/aorta/unittests/test_benchmark.py create mode 100644 cvs/schema/config_file/inference/__init__.py create mode 100644 cvs/schema/config_file/inference/atom/__init__.py create mode 100644 cvs/schema/config_file/inference/atom/unittests/__init__.py create mode 100644 cvs/schema/config_file/inference/atom/unittests/test_variant.py create mode 100644 cvs/schema/config_file/inference/atom/variant.py create mode 100644 cvs/schema/config_file/inference/common/__init__.py create mode 100644 cvs/schema/config_file/inference/common/accuracy.py create mode 100644 cvs/schema/config_file/inference/common/functional.py create mode 100644 cvs/schema/config_file/inference/common/long_context_accuracy.py create mode 100644 cvs/schema/config_file/inference/common/platform.py create mode 100644 cvs/schema/config_file/inference/common/sweep.py create mode 100644 cvs/schema/config_file/inference/common/unittests/__init__.py create mode 100644 cvs/schema/config_file/inference/common/unittests/test_accuracy.py create mode 100644 cvs/schema/config_file/inference/common/unittests/test_helpers.py create mode 100644 cvs/schema/config_file/inference/common/unittests/test_sweep.py create mode 100644 cvs/schema/config_file/inference/pytorch_xdit/__init__.py create mode 100644 cvs/schema/config_file/inference/pytorch_xdit/config.py create mode 100644 cvs/schema/config_file/inference/pytorch_xdit/unittests/__init__.py create mode 100644 cvs/schema/config_file/inference/pytorch_xdit/unittests/test_config.py create mode 100644 cvs/schema/config_file/inference/sglang/__init__.py create mode 100644 cvs/schema/config_file/inference/sglang/unittests/__init__.py create mode 100644 cvs/schema/config_file/inference/sglang/unittests/test_variant.py create mode 100644 cvs/schema/config_file/inference/sglang/variant.py create mode 100644 cvs/schema/config_file/inference/vllm/__init__.py create mode 100644 cvs/schema/config_file/inference/vllm/unittests/__init__.py create mode 100644 cvs/schema/config_file/inference/vllm/unittests/test_variant.py create mode 100644 cvs/schema/config_file/inference/vllm/variant.py create mode 100644 cvs/schema/config_file/preflight/__init__.py create mode 100644 cvs/schema/config_file/preflight/config.py create mode 100644 cvs/schema/config_file/preflight/unittests/__init__.py create mode 100644 cvs/schema/config_file/preflight/unittests/test_config.py create mode 100644 cvs/schema/config_file/training/__init__.py create mode 100644 cvs/schema/config_file/training/jaxmaxtext/__init__.py create mode 100644 cvs/schema/config_file/training/jaxmaxtext/unittests/__init__.py create mode 100644 cvs/schema/config_file/training/jaxmaxtext/unittests/test_variant.py create mode 100644 cvs/schema/config_file/training/jaxmaxtext/variant.py create mode 100644 cvs/schema/config_file/training/megatron/__init__.py create mode 100644 cvs/schema/config_file/training/megatron/unittests/__init__.py create mode 100644 cvs/schema/config_file/training/megatron/unittests/test_variant.py create mode 100644 cvs/schema/config_file/training/megatron/variant.py create mode 100644 cvs/schema/config_file/training/torchtitan/__init__.py create mode 100644 cvs/schema/config_file/training/torchtitan/unittests/__init__.py create mode 100644 cvs/schema/config_file/training/torchtitan/unittests/test_variant.py create mode 100644 cvs/schema/config_file/training/torchtitan/variant.py create mode 100644 cvs/schema/unittests/__init__.py create mode 100644 cvs/schema/unittests/test_validate.py create mode 100644 cvs/schema/validate.py diff --git a/AGENTS.md b/AGENTS.md index 0667a159f..167ffd715 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -43,7 +43,7 @@ make ut # sdist -> .test_venv -> run_all_unittests.py ### Test Structure - Test functions use `test_` prefix (PyTest requirement) - Customer test suites live under `cvs/tests/`; library code under `cvs/lib/` -- Configuration files in `cvs/input/config_file/` (JSON-structured, explicit parameters). Sample configs mark mandatory user fields with `` in value strings — users must replace every `` before running tests; unresolved placeholders hard-exit at startup via `_resolve_placeholders_in_dict` (`cvs/lib/utils_lib.py`), reached by `resolve_test_config_placeholders`, which nearly every test module calls. The `` validators in `cvs/parsers/schemas.py` cover only the aorta and pytorch-xdit configs. +- Configuration files in `cvs/input/config_file/` (JSON-structured, explicit parameters). Sample configs mark mandatory user fields with `` in value strings — users must replace every `` before running tests; unresolved placeholders hard-exit at startup via `_resolve_placeholders_in_dict` (`cvs/lib/utils_lib.py`), reached by `resolve_test_config_placeholders`, which nearly every test module calls. The `` validators in `cvs/schema/config_file/` (aorta and pytorch_xdit) cover only those config types. - Tests require `--cluster_file` and `--config_file` CLI args (wired via `cvs/conftest.py`; the `orch` fixture is in `cvs/tests/conftest.py`) ### Orchestrator Patterns (Recommended) diff --git a/cvs/lib/inference/ADDING_A_SUITE.md b/cvs/lib/inference/ADDING_A_SUITE.md index 4200152dc..04cbfe12b 100644 --- a/cvs/lib/inference/ADDING_A_SUITE.md +++ b/cvs/lib/inference/ADDING_A_SUITE.md @@ -65,8 +65,10 @@ Create `cvs/lib//utils/_config_loader.py`. from pydantic import model_validator from typing_extensions import Literal -from cvs.lib.utils.config_loader import BaseVariantConfig, _Forbid, substitute_config -from cvs.lib.inference.utils.inferencing_config_loader import ( +from cvs.schema.base import _Forbid +from cvs.schema.common.base import BaseVariantConfig +from cvs.lib.utils.config_loader import substitute_config +from cvs.schema.config_file.inference.common.sweep import ( GoodputSlo, Roles, Run, Sweep, SeqCombo, validate_sweep_selector, ) from cvs.lib..utils._parsing import GATED_METRICS @@ -100,7 +102,8 @@ class VariantConfig(BaseVariantConfig): @model_validator(mode="after") def _check_thresholds_cover_sweep(self): - # Copy the two-axis check from inferencing_config_loader.py: + # Copy the two-axis check from cvs/schema/config_file/inference/common/sweep.py + # (`validate_thresholds_cover_sweep`): # Axis 1: every sweep cell has a threshold entry; no key names a phantom cell. # Axis 2: every present cell has a spec for every GATED_METRICS member. # When enforce_thresholds=False: warn instead of raise. diff --git a/cvs/lib/inference/atom/atom_config_loader.py b/cvs/lib/inference/atom/atom_config_loader.py index dc28485d8..08354970c 100644 --- a/cvs/lib/inference/atom/atom_config_loader.py +++ b/cvs/lib/inference/atom/atom_config_loader.py @@ -2,277 +2,50 @@ Copyright 2025 Advanced Micro Devices, Inc. All rights reserved. -ATOM suite config schema (``atom``). +ATOM suite config loader (``atom``). -Generic paths/model/container/threshold plumbing lives in -:mod:`cvs.lib.utils.config_loader`. Sweep selector types are shared with -:mod:`cvs.lib.inference.utils.inferencing_config_loader`. +Pydantic models live in ``cvs.schema.config_file.inference.atom.variant``. ''' from __future__ import annotations -import re from typing import Any -from pydantic import Field, field_validator, model_validator -from typing_extensions import Literal - -from cvs.lib import globals -from cvs.lib.inference.atom.atom_parsing import GATED_METRICS -from cvs.lib.inference.utils.accuracy_config import AccuracyConfig -from cvs.lib.inference.utils.functional_config import FunctionalConfig -from cvs.lib.inference.utils.inferencing_config_loader import ( - RoleServer, - Sweep, - validate_sweep_selector, - validate_thresholds_cover_sweep, +from cvs.lib.utils.config_loader import substitute_config +from cvs.schema.config_file.inference.atom.variant import ( + ATOM_DRIVERS, + ATOM_PP_DRIVERS, + AtomParams, + AtomRoleServer, + AtomRoles, + AtomRunCard, + AtomVariantConfig, + MtpQualityConfig, + QuantParityConfig, + merge_mxfp4_triton_env, ) -from cvs.lib.inference.utils.long_context_accuracy_config import LongContextAccuracyConfig -from cvs.lib.inference.utils.platform_config import PlatformConfig -from cvs.lib.utils.config_loader import BaseVariantConfig, _Forbid, substitute_config - -ATOM_DRIVERS = ("atom", "vllm", "vllm_atom", "sglang") -ATOM_PP_DRIVERS = ("vllm", "vllm_atom", "sglang") -# MI300X MXFP4 MoE + A4W4 GEMM require Triton; aiter A4W4 is unsupported on gfx942. -# atom.utils.envs treats only "1" as true — "true" is ignored. -_MXFP4_TRITON_ENV = { - "ATOM_USE_TRITON_MOE": "1", - "ATOM_USE_TRITON_GEMM": "1", -} - - -def merge_mxfp4_triton_env(precision: str, env: dict[str, str]) -> dict[str, str]: - """Return server env with MXFP4 Triton defaults applied when unset.""" - merged = dict(env or {}) - if (precision or "").lower() == "mxfp4": - for key, value in _MXFP4_TRITON_ENV.items(): - merged.setdefault(key, value) - for key in _MXFP4_TRITON_ENV: - if str(merged.get(key, "")).lower() == "true": - merged[key] = "1" - return merged - - -log = globals.log - -# Written by test_discover_topology / resolve_multinode_fabric — not user env. -_ORCH_MANAGED_NETWORK_ENV = frozenset({"NCCL_SOCKET_IFNAME", "GLOO_SOCKET_IFNAME", "TP_SOCKET_IFNAME", "NCCL_IB_HCA"}) -_IB_HCA_NETDEV_RE = re.compile(r"^mlx5_\d+$", re.IGNORECASE) - - -class AtomRoleServer(RoleServer): - # Extra CLI tokens for ``python -m atom.entrypoints.openai_server`` after - # ``--model`` / ``--server-port`` (e.g. ``-tp``, ``--kv_cache_dtype``). - atom_args: list[str] = [] - # Extra CLI tokens appended to ``python3 -m sglang.launch_server`` (driver=sglang). - sglang_args: list[str] = [] - # IB HCA devices for NCCL_IB_HCA (multinode only). - # absent or "auto" -> use whatever ibv_devinfo -l reports (test_discover_topology). - # explicit list -> validated at preflight against ibv_devinfo output. - ib_hca_devices: Literal["auto"] | list[str] | None = None - # Linux netdev for NCCL_SOCKET_IFNAME / GLOO_SOCKET_IFNAME on multinode PP runs. - # absent or "auto" -> resolved at runtime by test_discover_topology from cluster IPs. - ib_netdev: Literal["auto"] | str | None = None - - @field_validator("ib_netdev", mode="after") - @classmethod - def _normalize_ib_netdev(cls, v): - raw = (v or "").strip() - if raw and raw.lower() != "auto" and _IB_HCA_NETDEV_RE.match(raw): - log.warning( - "roles.server.ib_netdev=%r looks like an IB HCA name; coercing to 'auto' " - "(socket netdev is discovered from cluster IPs at runtime)", - raw, - ) - return "auto" - return v - - @model_validator(mode="after") - def _strip_orchestrator_managed_network_env(self): - if not self.env: - return self - dropped = sorted(k for k in self.env if k in _ORCH_MANAGED_NETWORK_ENV) - if not dropped: - return self - log.warning( - "roles.server.env drops orchestrator-managed keys %s " - "(set by test_discover_topology / build_server_cmd instead)", - dropped, - ) - self.env = {k: v for k, v in self.env.items() if k not in _ORCH_MANAGED_NETWORK_ENV} - return self - - -class AtomRoles(_Forbid): - server: AtomRoleServer = AtomRoleServer() - - -class AtomParams(_Forbid): - # ``atom`` = standalone ATOM openai_server + benchmark_serving. - # ``vllm_atom`` = vLLM coordinator + ATOM local kernels (true multinode PP). - # ``vllm`` = interim ROCm vLLM uplift (vllm serve + vllm bench serve). - # ``sglang`` = SGLang coordinator (launch_server + bench_serving) for PP runs. - driver: Literal["atom", "vllm", "vllm_atom", "sglang"] = "vllm" - backend: str = "vllm" - base_url: str = "http://0.0.0.0" - port_no: str = "8000" - dataset_name: str = "random" - burstiness: str = "1.0" - seed: str = "0" - request_rate: str = "inf" - random_range_ratio: str = "0.8" - random_prefix_len: str = "0" - tensor_parallelism: str = "8" - tokenizer_mode: str = "auto" - percentile_metrics: str = "ttft,tpot,itl,e2el" - metric_percentiles: str = "95,99" - num_prompts: str = "1000" - max_model_length: str = "8192" - client_poll_count: str = "50" - client_poll_wait_time: str = "60" - client_initial_wait_s: str = "120" - server_precheck_wait_s: str = "30" - server_warmup_wait_s: str = "330" - server_poll_count: str = "60" - server_poll_wait_time: str = "60" - reuse_server_across_sweep: str = "false" - bench_max_failed_requests: str = "0" - bench_extra_args: str = "" - result_filename: str = "results" - # Multinode (M5): omit or set nnodes=1 for single-node runs. When nnodes>1, - # cluster node_dict must list the same number of hosts and test_setup_sshd runs. - nnodes: str = "1" - pipeline_parallel_size: str = "1" - master_addr: str = "" - master_port: str = "29501" - # Optional single-node reference output_throughput for scaling.efficiency_pct. - scaling_baseline_output_throughput: str = "" - - -class AtomRunCard(_Forbid): - upstream_run_url: str = "" - atom_image_pin: str = "" - notes: str = "" - - -class MtpQualityConfig(_Forbid): - enabled: bool = False - chat_template_prompt: str = "Say hello in one short sentence." - chat_template_expected_sha256: str = "" - - -class QuantParityConfig(_Forbid): - enabled: bool = False - probe_prompt: str = "The capital of France is" - reference_config_stem: str = "" - - -ATOM_FRAMEWORKS = ("atom",) - - -class AtomVariantConfig(BaseVariantConfig): - framework: Literal["atom"] - - gpu_arch: str - run_card: AtomRunCard = AtomRunCard() - roles: AtomRoles = AtomRoles() - params: AtomParams - sweep: Sweep - accuracy: AccuracyConfig = Field(default_factory=AccuracyConfig) - mtp_quality: MtpQualityConfig = Field(default_factory=MtpQualityConfig) - quant_parity: QuantParityConfig = Field(default_factory=QuantParityConfig) - functional: FunctionalConfig = Field(default_factory=FunctionalConfig) - long_context_accuracy: LongContextAccuracyConfig = Field(default_factory=LongContextAccuracyConfig) - platform: PlatformConfig = Field(default_factory=PlatformConfig) - - def cell_key(self, isl, osl, concurrency): - p = self.params - key = f"ISL={isl},OSL={osl},TP={p.tensor_parallelism}" - nnodes = int(p.nnodes) - pp = int(p.pipeline_parallel_size) - if p.driver == "atom": - if nnodes > 1: - key += f",DP={nnodes},NNODES={nnodes}" - elif p.driver in ATOM_PP_DRIVERS: - if pp > 1 or nnodes > 1: - key += f",PP={p.pipeline_parallel_size}" - if nnodes > 1: - key += f",NNODES={p.nnodes}" - return f"{key},CONC={concurrency}" - - def expected_cells(self) -> list[str]: - by_name = {c.name: c for c in self.sweep.sequence_combinations} - return [self.cell_key(by_name[r.combo].isl, by_name[r.combo].osl, r.concurrency) for r in self.sweep.runs] - - @model_validator(mode="after") - def _apply_mxfp4_triton_env_defaults(self): - self.roles.server.env = merge_mxfp4_triton_env(self.model.precision, self.roles.server.env) - return self - - @model_validator(mode="after") - def _check_thresholds_cover_sweep(self): - validate_thresholds_cover_sweep( - expected_cells=self.expected_cells(), - thresholds=self.thresholds, - enforce_thresholds=self.enforce_thresholds, - gated_metrics=GATED_METRICS, - ) - if int(self.params.nnodes) > 1 and (self.params.scaling_baseline_output_throughput or "").strip(): - missing = [] - for cell in self.expected_cells(): - specs = self.thresholds.get(cell) or {} - if "scaling.efficiency_pct" not in specs: - missing.append(cell) - if missing: - msg = ( - "multinode variant with scaling_baseline_output_throughput requires " - f"scaling.efficiency_pct in every cell; missing: {missing}" - ) - if self.enforce_thresholds: - raise ValueError(msg) - import warnings - - warnings.warn(f"{msg} (enforce_thresholds=false -> record-only)", stacklevel=2) - return self - - @model_validator(mode="after") - def _atom_multinode_uses_dp_not_pp(self): - if self.params.driver == "atom" and int(self.params.nnodes) > 1: - if int(self.params.pipeline_parallel_size) > 1: - raise ValueError( - "params.driver='atom' with nnodes>1 uses ATOM SPMD data parallel (-dp); " - "standalone ATOM cannot execute pipeline parallel. For true PP>1 use " - "params.driver='vllm_atom' or 'sglang'." - ) - return self - - @model_validator(mode="after") - def _pp_driver_distributed_consistency(self): - driver = self.params.driver - if driver not in ATOM_PP_DRIVERS: - return self - nn = int(self.params.nnodes) - pp = int(self.params.pipeline_parallel_size) - is_ray = self.roles.server.serve_args.get("distributed-executor-backend") == "ray" - if nn > 1 and pp == 1 and not is_ray: - raise ValueError( - f"params.driver={driver!r} with nnodes={nn} requires pipeline_parallel_size>1 " - f"(got pp={pp}) for multinode pipeline parallel" - ) - if pp > 1 and nn == 1: - raise ValueError( - f"pipeline_parallel_size={pp} > 1 requires nnodes > 1 (got nnodes={nn}) for params.driver={driver!r}" - ) - return self - - @model_validator(mode="after") - def _atom_driver_requires_inline_server_args(self): - if self.params.driver == "atom" and not self.roles.server.atom_args: - raise ValueError( - "params.driver='atom' requires roles.server.atom_args " - "(inline ATOM openai_server CLI tokens, vLLM-style)" - ) - return self +from cvs.schema.config_file.inference.common.sweep import validate_sweep_selector + +__all__ = [ + "ATOM_DRIVERS", + "ATOM_PP_DRIVERS", + "AtomParams", + "AtomRoleServer", + "AtomRoles", + "AtomRunCard", + "AtomVariantConfig", + "MtpQualityConfig", + "QuantParityConfig", + "expand_sweep", + "expand_sweep_parametrize", + "load_variant", + "merge_mxfp4_triton_env", + "orchestrator_container_from_variant", + "placeholder_gated_threshold_cell", + "reuse_server_flag", + "server_session_key", + "validate_sweep_selector", +] def expand_sweep(sweep): @@ -296,13 +69,11 @@ def expand_sweep(sweep): def reuse_server_flag(params) -> bool: - """Return True when ``params.reuse_server_across_sweep`` is a truthy string.""" raw = str(getattr(params, "reuse_server_across_sweep", "false")).strip().lower() return raw in ("true", "1", "yes") def server_session_key(variant_config, isl, osl): - """Stable key for server reuse across sweep cells with identical model/shape.""" p = variant_config.params roles = variant_config.roles.server if p.driver == "atom": @@ -326,7 +97,6 @@ def server_session_key(variant_config, isl, osl): def expand_sweep_parametrize(sweep, fixturenames): - """Build pytest parametrize args for inference or metric-tier collection.""" from cvs.lib.inference.atom.atom_parsing import METRIC_TIER_ORDER cases, ids = expand_sweep(sweep) @@ -364,9 +134,11 @@ def placeholder_gated_threshold_cell( failed_max: int = 1_000_000_000, success_rate_min: float = 0, ) -> dict[str, Any]: - """Return one sweep cell's ``client.*`` specs covering every ``GATED_METRICS`` member.""" + """Return one sweep cell's ``client.*`` specs covering every gated metric.""" + from cvs.lib.inference.atom.atom_parsing import GATED_METRICS + loose_ms = {"kind": "max_ms", "value": 1_000_000} - return { + out = { "client.total_token_throughput": {"kind": "min_tok_s", "value": total_token_throughput_min}, "client.output_throughput": {"kind": "min_tok_s", "value": output_throughput_min}, "client.per_gpu_throughput": {"kind": "min_tok_s", "value": per_gpu_throughput_min}, @@ -393,10 +165,15 @@ def placeholder_gated_threshold_cell( "client.success_rate": {"kind": "min", "value": success_rate_min}, "client.failed": {"kind": "max", "value": failed_max}, } + for m in GATED_METRICS: + key = f"client.{m}" + if key not in out: + kind = "max_ms" if m.endswith("_ms") else "max" if m == "failed" else "min" + out[key] = {"kind": kind, "value": 0 if kind == "min" else 1_000_000} + return out def orchestrator_container_from_variant(variant: AtomVariantConfig) -> dict[str, Any]: - """``container`` block for :class:`OrchestratorConfig` (includes server env).""" block = variant.container.model_dump() server_env = variant.roles.server.env if server_env: diff --git a/cvs/lib/inference/sglang/sglang_config_loader.py b/cvs/lib/inference/sglang/sglang_config_loader.py index 5e267172b..f4192ee97 100644 --- a/cvs/lib/inference/sglang/sglang_config_loader.py +++ b/cvs/lib/inference/sglang/sglang_config_loader.py @@ -27,16 +27,13 @@ from pathlib import Path from typing import Any, Dict, Mapping -from pydantic import Field, model_validator -from typing_extensions import Literal - from cvs.lib import globals -from cvs.lib.utils.config_loader import ( - BaseVariantConfig, - _Forbid, - substitute_config, -) +from cvs.lib.utils.config_loader import substitute_config from cvs.lib.utils_lib import resolve_test_config_placeholders +from cvs.schema.config_file.inference.sglang.variant import ( + SglangSingleVariantConfig, + perf_cell_key, +) log = globals.log @@ -94,17 +91,6 @@ def flat_expected_from_specs(specs: Mapping[str, Any]) -> dict[str, float]: return out -def perf_cell_key(bp_dict: Mapping[str, Any]) -> str: - bench = (bp_dict.get("inference_tests") or {}).get("bench_serv_random") or {} - return ( - f"ISL={bench.get('input_length', '-')}," - f"OSL={bench.get('output_length', '-')}," - f"TP={bp_dict.get('tensor_parallelism', '8')}," - f"PP={bp_dict.get('pipeline_parallelism', '1')}," - f"CONC={bp_dict.get('max_concurrency', '-')}" - ) - - def bench_cell_key(bench_name: str) -> str: return f"BENCH={bench_name}" @@ -297,53 +283,6 @@ def _is_legacy_root(raw: Mapping[str, Any]) -> bool: return "benchmark_params" in raw and ("config" in raw or "container_image" in raw) -# ---------- typed config ---------- - - -class SglangRoleServer(_Forbid): - env: Dict[str, str] = Field(default_factory=dict) - serve_port: str = "8000" - - -class SglangRoles(_Forbid): - server: SglangRoleServer = Field(default_factory=SglangRoleServer) - - -class SglangSingleVariantConfig(BaseVariantConfig): - """Typed config for ``sglang_single`` + ContainerOrchestrator.""" - - framework: Literal["sglang_single"] - gpu_arch: str - variant_key: str = "" - config_path: str = "" - - # Legacy blocks kept for ``SglangSingle`` until that lib is refactored. - inference: Dict[str, Any] = Field(default_factory=dict) - benchmark_params: Dict[str, Any] = Field(default_factory=dict) - - roles: SglangRoles = Field(default_factory=SglangRoles) - - def cell_key(self, isl, osl, concurrency) -> str: - tp = self.benchmark_params.get("tensor_parallelism", "-") - pp = self.benchmark_params.get("pipeline_parallelism", "-") - return f"ISL={isl},OSL={osl},TP={tp},PP={pp},CONC={concurrency}" - - def perf_cell_key(self) -> str: - return perf_cell_key(self.benchmark_params) - - @property - def hf_token_file(self) -> str: - return self.paths.hf_token_file - - @model_validator(mode="after") - def _sync_legacy_inference_container_name(self): - """Keep legacy inference dict aligned with orchestrator container name.""" - if self.inference and self.container.name: - self.inference["container_name"] = self.container.name - self.inference["container_image"] = self.container.image - return self - - # ---------- public API ---------- diff --git a/cvs/lib/inference/unittests/test_atom_config_loader.py b/cvs/lib/inference/unittests/test_atom_config_loader.py index 5e205edcc..6a1d87927 100644 --- a/cvs/lib/inference/unittests/test_atom_config_loader.py +++ b/cvs/lib/inference/unittests/test_atom_config_loader.py @@ -18,7 +18,7 @@ reuse_server_flag, server_session_key, ) -from cvs.lib.inference.utils.inferencing_config_loader import Run, SeqCombo, Sweep +from cvs.schema.config_file.inference.common.sweep import Run, SeqCombo, Sweep def _cluster_dict(): diff --git a/cvs/lib/inference/unittests/test_inferencing_config_loader.py b/cvs/lib/inference/unittests/test_inferencing_config_loader.py deleted file mode 100644 index 0e3dc67f2..000000000 --- a/cvs/lib/inference/unittests/test_inferencing_config_loader.py +++ /dev/null @@ -1,407 +0,0 @@ -''' -Copyright 2025 Advanced Micro Devices, Inc. -All rights reserved. - -Unit tests for cvs.lib.utils.config_loader (ModelSpec, BaseVariantConfig, -substitute_config) and cvs.lib.inference.utils.inferencing_config_loader -(Sweep, SeqCombo, GoodputSlo, Run, VariantConfig.expected_cells, -_check_thresholds_cover_sweep). No hardware. -''' - -import unittest -import warnings - -from pydantic import ValidationError - -from cvs.lib.inference.utils.inferencing_config_loader import ( - GoodputSlo, - Run, - SeqCombo, - Sweep, - VariantConfig, -) -from cvs.lib.utils.config_loader import ModelSpec -from cvs.lib.inference.utils.vllm_parsing import GATED_METRICS - - -def _combo(name, isl="128", osl="2048"): - return SeqCombo(name=name, isl=isl, osl=osl) - - -def _full_gated_specs(): - """A spec for every gated metric -- the minimum that satisfies coverage. - - Values are inert (a 0 floor / huge ceiling) so the set passes without - asserting anything; these tests pin the coverage gate, not the numbers. - """ - out = {} - for m in GATED_METRICS: - kind = "max_ms" if m.endswith("_ms") else "max" if m == "failed" else "min" - out[f"client.{m}"] = {"kind": kind, "value": 0 if kind == "min" else 1e12} - return out - - -def _variant(sweep, tp="8", thresholds=None, enforce_thresholds=False): - """A minimal VariantConfig carrying just enough to exercise expected_cells. - - remote=0 (the remote guard would otherwise reject it) and - enforce_thresholds=False so the empty threshold dict does not trip the - coverage check -- this test pins the selector expansion, not the gate. - """ - return VariantConfig( - schema_version=1, - framework="vllm_single", - gpu_arch="mi300x", - enforce_thresholds=enforce_thresholds, - threshold_json="", - paths={ - "shared_fs": "/home/x", - "models_dir": "/home/x/models", - "log_dir": "/home/x/LOGS", - "hf_token_file": "/home/x/.hf", - }, - model={"id": "amd/Llama-3.1-70B-Instruct-FP8-KV", "remote": 0}, - container={ - "name": "c", - "image": "rocm/vllm-dev:nightly-sshd", - "runtime": {"name": "docker"}, - }, - params={"tensor_parallelism": tp}, - sweep=sweep, - thresholds=thresholds or {}, - ) - - -class TestSweepValidator(unittest.TestCase): - def test_valid_runs_selector_constructs(self): - sw = Sweep( - sequence_combinations=[_combo("a"), _combo("b", osl="4096")], - runs=[Run(combo="a", concurrency=16), Run(combo="b", concurrency=32)], - ) - self.assertEqual([r.combo for r in sw.runs], ["a", "b"]) - - def test_unknown_run_combo_raises(self): - with self.assertRaises(ValidationError) as ctx: - Sweep( - sequence_combinations=[_combo("a")], - runs=[Run(combo="typo", concurrency=16)], - ) - self.assertIn("names no sequence_combination", str(ctx.exception)) - - def test_duplicate_combo_names_raise(self): - with self.assertRaises(ValidationError) as ctx: - Sweep( - sequence_combinations=[_combo("a"), _combo("a", osl="4096")], - runs=[Run(combo="a", concurrency=16)], - ) - self.assertIn("duplicate sequence_combination names", str(ctx.exception)) - - def test_concurrency_levels_is_rejected(self): - # The old cartesian key must be gone (extra=forbid): a config still - # carrying concurrency_levels should fail loudly, not silently ignore it. - with self.assertRaises(ValidationError): - Sweep( - sequence_combinations=[_combo("a")], - runs=[Run(combo="a", concurrency=16)], - concurrency_levels=[16], - ) - - -class TestExpectedCells(unittest.TestCase): - def test_runs_expand_to_exactly_their_cells(self): - sw = Sweep( - sequence_combinations=[_combo("a", isl="128", osl="2048"), _combo("b", isl="256", osl="4096")], - runs=[ - Run(combo="a", concurrency=16), - Run(combo="b", concurrency=32), - Run(combo="a", concurrency=64), - ], - ) - vc = _variant(sw) - self.assertEqual( - vc.expected_cells(), - [ - "ISL=128,OSL=2048,TP=8,CONC=16", - "ISL=256,OSL=4096,TP=8,CONC=32", - "ISL=128,OSL=2048,TP=8,CONC=64", - ], - ) - - def test_no_cartesian_blowup(self): - # Two combos + two runs must yield TWO cells, not 2x2=4 (the old bug). - sw = Sweep( - sequence_combinations=[_combo("a"), _combo("b", osl="4096")], - runs=[Run(combo="a", concurrency=16), Run(combo="b", concurrency=16)], - ) - self.assertEqual(len(_variant(sw).expected_cells()), 2) - - -class TestGatedMetricCoverage(unittest.TestCase): - """The gated-metric axis of _check_thresholds_cover_sweep.""" - - _CELL = "ISL=128,OSL=2048,TP=8,CONC=16" - - def _variant_with(self, thresholds, enforce): - sw = Sweep( - sequence_combinations=[_combo("a")], - runs=[Run(combo="a", concurrency=16)], - ) - return VariantConfig( - schema_version=1, - framework="vllm_single", - gpu_arch="mi300x", - enforce_thresholds=enforce, - threshold_json="", - paths={ - "shared_fs": "/home/x", - "models_dir": "/home/x/models", - "log_dir": "/home/x/LOGS", - "hf_token_file": "/home/x/.hf", - }, - model={"id": "amd/Llama-3.1-70B-Instruct-FP8-KV", "remote": 0}, - container={"name": "c", "image": "rocm/vllm-dev:nightly-sshd", "runtime": {"name": "docker"}}, - params={"tensor_parallelism": "8"}, - sweep=sw, - thresholds=thresholds, - ) - - def test_full_gated_set_constructs(self): - vc = self._variant_with({self._CELL: _full_gated_specs()}, enforce=True) - self.assertEqual(vc.enforce_thresholds, True) - - def test_missing_gated_metric_raises_when_enforced(self): - specs = _full_gated_specs() - del specs["client.p99_ttft_ms"] # drop one gated metric - with self.assertRaises(ValidationError) as ctx: - self._variant_with({self._CELL: specs}, enforce=True) - self.assertIn("missing gated-metric specs", str(ctx.exception)) - self.assertIn("client.p99_ttft_ms", str(ctx.exception)) - - def test_missing_gated_metric_warns_when_record_only(self): - specs = _full_gated_specs() - del specs["client.failed"] - - with warnings.catch_warnings(record=True) as caught: - warnings.simplefilter("always") - self._variant_with({self._CELL: specs}, enforce=False) - self.assertTrue(any("missing gated-metric specs" in str(x.message) for x in caught)) - - def test_extra_non_gated_spec_is_allowed(self): - # A spec for a non-gated metric (record-only display extra) must not - # trip coverage -- gating is a floor, not an allow-list. - specs = _full_gated_specs() - specs["client.num_prompts"] = {"kind": "min", "value": 0} - vc = self._variant_with({self._CELL: specs}, enforce=True) - self.assertIn("client.num_prompts", vc.thresholds[self._CELL]) - - -class TestModelSpecPrecision(unittest.TestCase): - """precision is an accepted optional field on ModelSpec (default '').""" - - def test_precision_field_is_accepted(self): - ms = ModelSpec(id="amd/Llama-3.1-70B", remote=0, precision="fp8") - self.assertEqual(ms.precision, "fp8") - - def test_valid_model_spec_without_precision(self): - ms = ModelSpec(id="amd/Llama-3.1-70B", remote=0) - self.assertEqual(ms.id, "amd/Llama-3.1-70B") - self.assertEqual(ms.remote, 0) - # precision is optional and defaults to empty. - self.assertEqual(ms.precision, "") - - def test_unknown_field_is_rejected(self): - # ModelSpec is _Forbid: a truly unknown field still fails validation. - with self.assertRaises(ValidationError): - ModelSpec(id="amd/Llama-3.1-70B", remote=0, bogus="x") - - -class TestThresholdJsonField(unittest.TestCase): - """threshold_json is an optional field on BaseVariantConfig / VariantConfig - (default ''); when absent, threshold discovery falls back to the sibling - *threshold.json next to the config.""" - - def _base_kwargs(self): - sw = Sweep( - sequence_combinations=[_combo("a")], - runs=[Run(combo="a", concurrency=16)], - ) - return dict( - schema_version=1, - framework="vllm_single", - gpu_arch="mi300x", - enforce_thresholds=False, - paths={ - "shared_fs": "/home/x", - "models_dir": "/home/x/models", - "log_dir": "/home/x/LOGS", - "hf_token_file": "/home/x/.hf", - }, - model={"id": "amd/Llama-3.1-70B", "remote": 0}, - container={"name": "c", "image": "img", "runtime": {"name": "docker"}}, - params={"tensor_parallelism": "8"}, - sweep=sw, - thresholds={}, - ) - - def test_missing_threshold_json_defaults_to_empty(self): - kwargs = self._base_kwargs() - # threshold_json deliberately absent -> optional, defaults to "". - vc = VariantConfig(**kwargs) - self.assertEqual(vc.threshold_json, "") - - def test_threshold_json_present_constructs(self): - kwargs = self._base_kwargs() - kwargs["threshold_json"] = "/some/absolute/path/threshold.json" - vc = VariantConfig(**kwargs) - self.assertEqual(vc.threshold_json, "/some/absolute/path/threshold.json") - - -class TestCellCoverageAxis(unittest.TestCase): - """Axis-1 of _check_thresholds_cover_sweep: cell vs threshold key mismatch.""" - - _CELL = "ISL=128,OSL=2048,TP=8,CONC=16" - - def _variant_with(self, thresholds, enforce=True): - sw = Sweep( - sequence_combinations=[_combo("a")], - runs=[Run(combo="a", concurrency=16)], - ) - return VariantConfig( - schema_version=1, - framework="vllm_single", - gpu_arch="mi300x", - enforce_thresholds=enforce, - threshold_json="", - paths={ - "shared_fs": "/home/x", - "models_dir": "/home/x/models", - "log_dir": "/home/x/LOGS", - "hf_token_file": "/home/x/.hf", - }, - model={"id": "amd/Llama-3.1-70B", "remote": 0}, - container={"name": "c", "image": "img", "runtime": {"name": "docker"}}, - params={"tensor_parallelism": "8"}, - sweep=sw, - thresholds=thresholds, - ) - - def test_sweep_cell_with_no_threshold_entry_raises(self): - # thresholds is empty -> sweep cell has no entry -> axis-1 fires - with self.assertRaises(ValidationError) as ctx: - self._variant_with(thresholds={}, enforce=True) - self.assertIn("sweep cells with no threshold entry", str(ctx.exception)) - self.assertIn(self._CELL, str(ctx.exception)) - - def test_threshold_key_matching_no_sweep_cell_raises(self): - # thresholds has the real cell PLUS a bogus key -> extra set is non-empty - specs = _full_gated_specs() - with self.assertRaises(ValidationError) as ctx: - self._variant_with( - thresholds={self._CELL: specs, "ISL=999,OSL=999,TP=8,CONC=99": specs}, - enforce=True, - ) - self.assertIn("threshold keys matching no sweep cell", str(ctx.exception)) - self.assertIn("ISL=999,OSL=999,TP=8,CONC=99", str(ctx.exception)) - - def test_cell_mismatch_warns_when_record_only(self): - with warnings.catch_warnings(record=True) as caught: - warnings.simplefilter("always") - self._variant_with(thresholds={}, enforce=False) - self.assertTrue(any("sweep cells with no threshold entry" in str(w.message) for w in caught)) - - def test_accuracy_key_does_not_trip_extra_key_check(self): - # "accuracy" is a top-level threshold key for lm-eval gating, not a - # sweep cell -- it must not be flagged as an unrecognized extra key. - specs = _full_gated_specs() - vc = self._variant_with( - thresholds={self._CELL: specs, "accuracy": {"mmlu": {"mmlu.acc__none": {"kind": "min", "value": 0.5}}}}, - enforce=True, - ) - self.assertIn("accuracy", vc.thresholds) - - def test_unrecognized_key_still_raises_alongside_accuracy(self): - # The "accuracy" exclusion must be narrow: a genuinely unrecognized - # key (typo'd or bogus) alongside a valid "accuracy" block still trips - # the extra-key check. - specs = _full_gated_specs() - with self.assertRaises(ValidationError) as ctx: - self._variant_with( - thresholds={ - self._CELL: specs, - "accuracy": {"mmlu": {"mmlu.acc__none": {"kind": "min", "value": 0.5}}}, - "acuracy": {}, - }, - enforce=True, - ) - self.assertIn("threshold keys matching no sweep cell", str(ctx.exception)) - self.assertIn("acuracy", str(ctx.exception)) - - -class TestExpectedCellsBoundaries(unittest.TestCase): - """Boundary cases for VariantConfig.expected_cells.""" - - def test_empty_runs_yields_empty_cells(self): - sw = Sweep(sequence_combinations=[_combo("a")], runs=[]) - self.assertEqual(_variant(sw).expected_cells(), []) - - def test_unreferenced_combo_not_in_expected_cells(self): - # combo 'unused' is declared but never referenced by any run - sw = Sweep( - sequence_combinations=[_combo("a"), _combo("unused", isl="999", osl="999")], - runs=[Run(combo="a", concurrency=16)], - ) - cells = _variant(sw).expected_cells() - self.assertEqual(len(cells), 1) - self.assertNotIn("ISL=999", cells[0]) - - -class TestGoodputSlo(unittest.TestCase): - """GoodputSlo is a _Forbid model with three required float fields.""" - - def test_valid_goodput_slo_constructs(self): - slo = GoodputSlo(ttft_ms=100.0, tpot_ms=50.0, e2el_ms=5000.0) - self.assertEqual(slo.ttft_ms, 100.0) - self.assertEqual(slo.tpot_ms, 50.0) - self.assertEqual(slo.e2el_ms, 5000.0) - - def test_missing_required_field_raises(self): - for missing in ("ttft_ms", "tpot_ms", "e2el_ms"): - with self.subTest(missing=missing): - kwargs = {"ttft_ms": 1.0, "tpot_ms": 1.0, "e2el_ms": 1.0} - del kwargs[missing] - with self.assertRaises(ValidationError): - GoodputSlo(**kwargs) - - def test_extra_key_raises(self): - with self.assertRaises(ValidationError): - GoodputSlo(ttft_ms=1.0, tpot_ms=1.0, e2el_ms=1.0, ttft_msec=1.0) - - def test_seq_combo_with_goodput_slo(self): - slo = GoodputSlo(ttft_ms=1000.0, tpot_ms=50.0, e2el_ms=10000.0) - combo = SeqCombo(name="a", isl="128", osl="2048", goodput_slo=slo) - self.assertIsNotNone(combo.goodput_slo) - self.assertEqual(combo.goodput_slo.e2el_ms, 10000.0) - - def test_seq_combo_without_goodput_slo(self): - combo = SeqCombo(name="a", isl="128", osl="2048") - self.assertIsNone(combo.goodput_slo) - - -class TestSeqComboForbid(unittest.TestCase): - """SeqCombo is _Forbid: missing required fields and extra keys must raise.""" - - def test_missing_required_field_raises(self): - for missing in ("name", "isl", "osl"): - with self.subTest(missing=missing): - kwargs = {"name": "a", "isl": "128", "osl": "2048"} - del kwargs[missing] - with self.assertRaises(ValidationError): - SeqCombo(**kwargs) - - def test_extra_key_raises(self): - with self.assertRaises(ValidationError): - SeqCombo(name="a", isl="128", osl="2048", unknown_field="x") - - -if __name__ == "__main__": - unittest.main() diff --git a/cvs/lib/inference/utils/AGENTS.md b/cvs/lib/inference/utils/AGENTS.md index 28aadada7..f8490e0ad 100644 --- a/cvs/lib/inference/utils/AGENTS.md +++ b/cvs/lib/inference/utils/AGENTS.md @@ -1,130 +1,24 @@ # cvs/lib/inference/utils — inference-specific config and parsing **Boundary**: this is the serving/inference half of the config machinery. -The generic half (`BaseVariantConfig`, `substitute_config`, `evaluate_all`, `Paths`, `ContainerSpec`) -lives in `cvs/lib/utils/` — import from there, never duplicate it here. +Shared schema (`BaseVariantConfig`, `Paths`, `ContainerSpec`) lives in +`cvs/schema/common/base.py`; I/O (`substitute_config`) lives in `cvs/lib/utils/config_loader.py`. +Sweep primitives live in `cvs/schema/config_file/inference/common/sweep.py`. --- ## Files -### `inferencing_config_loader.py` +### Shared inference schema (`cvs/schema/...`) -#### Schema classes +Sweep types (`RoleServer`, `Roles`, `GoodputSlo`, `SeqCombo`, `Run`, `Sweep`), +`validate_sweep_selector`, and `validate_thresholds_cover_sweep` live in +`cvs/schema/config_file/inference/common/sweep.py`. Per-framework variant models +live under `cvs/schema/config_file/inference//variant.py`. -**`RoleServer`** (`_Forbid`): per-model server overrides. -- `serve_args: Dict[str, Any]` — extra `vllm serve` flags; scalar → `--flag value`, - `True` → bare `--flag`, list → flag repeated per element -- `env: Dict[str, str]` — env vars merged over orchestrator defaults -- Both default empty; fp8-kv cells set `--kv-cache-dtype` here to keep the generic driver model-agnostic - -**`Roles`** (`_Forbid`): wraps `RoleServer`. -- `server: RoleServer` — defaults to empty `RoleServer()` - -**`GoodputSlo`** (`_Forbid`): per-combo goodput gate, in milliseconds. -- `ttft_ms: float`, `tpot_ms: float`, `e2el_ms: float` -- **INPUT to the run** (passed to `vllm bench serve --goodput`), NOT a threshold to assert. - Lives in the sweep, not `threshold.json`. `_Forbid` ensures a typo'd SLO key fails load - rather than silently drop the SLO and run with the wrong gate on hardware. - -**`SeqCombo`** (`_Forbid`): one named sequence-length combination. -- `name: str` — the join key referenced by `Run.combo` -- `isl: str`, `osl: str` -- `goodput_slo: Optional[GoodputSlo]` — omit when no goodput gate is needed - -**`Run`** (`_Forbid`): one sweep cell — a named combo at a single concurrency. -- `combo: str` — references a `SeqCombo.name` -- `concurrency: int` -- Explicit `runs[]` replaces the old NxM cartesian (`sequence_combinations × concurrency_levels`); - you enumerate exactly the cells you want - -**`Sweep`** (`_Forbid`): the full sweep selector. -- `sequence_combinations: List[SeqCombo]` -- `runs: List[Run]` -- `@model_validator(mode="after")` delegates to `validate_sweep_selector` - -**`Params`** (`_Forbid`): `vllm bench serve` CLI flags; all fields are `str`. - -| Field | Default | Notes | -|---|---|---| -| `backend` | `"vllm"` | | -| `base_url` | `"http://0.0.0.0"` | | -| `port_no` | `"8888"` | | -| `dataset_name` | `"random"` | | -| `burstiness` | `"1.0"` | | -| `seed` | `"0"` | | -| `request_rate` | `"inf"` | | -| `random_range_ratio` | `"0.8"` | | -| `random_prefix_len` | `"0"` | | -| `tensor_parallelism` | `"1"` | used in `cell_key` and `per_gpu_throughput` | -| `tokenizer_mode` | `"auto"` | | -| `percentile_metrics` | `"ttft,tpot,itl,e2el"` | | -| `metric_percentiles` | `"50,90,95,99"` | | -| `num_prompts` | `"3200"` | overridden per-cell by `_num_prompts_for` | -| `client_poll_count` | `"20"` | see below | - -`client_poll_count` semantics: total client wait budget = -**(client_poll_count × 60 s) + 120 s initial wait**. -The poll loop exits as soon as the client finishes, so raising this never slows down fast cells. -Raise it for high-osl cells where large-output runs take longer to complete. (Regression: REG-20260609-001) - -**`VariantConfig(BaseVariantConfig)`**: the full typed config. -- Adds: `framework: Literal["vllm_single"]`, `gpu_arch: str`, `roles: Roles = Roles()`, - `params: Params`, `sweep: Sweep` -- Implements: `cell_key(isl, osl, concurrency)`, `expected_cells()` -- Has: `@model_validator(mode="after") _check_thresholds_cover_sweep` - ---- - -#### Public functions - -**`load_variant(config_path, cluster_dict) -> VariantConfig`** - -The function a suite's `variant_config` fixture calls. - -1. Delegates file read + 3-pass placeholder substitution to `substitute_config` -2. Attaches thresholds returned by `substitute_config` to the raw dict -3. Builds and returns a typed, validated `VariantConfig` - -Does not reimplement file reading or substitution — always calls `substitute_config`. -See `cvs/lib/utils/AGENTS.md` for the full `substitute_config` contract. - -**`validate_sweep_selector(combo_names, run_combo_refs)`** — PUBLIC ENTRY POINT - -Shared rule called by **both**: -- the typed `Sweep` validator at load time -- `pytest_generate_tests` at collection time (reads raw JSON before the loader runs) - -Checks: -- combo names are unique (duplicate → `ValueError`) -- every `run.combo` names a known `sequence_combination` (unknown → `ValueError`) - -Operates on plain `list[str]` so both call sites feed it without the full typed schema. -If you add a sweep check, add it here so both paths enforce it without drift. - ---- - -#### `_check_thresholds_cover_sweep` — two-axis coverage check - -`@model_validator(mode="after")` on `VariantConfig`. Fails at load time if the threshold file -does not match the sweep matrix. - -**Axis 1 — cell coverage** -- Every sweep cell produced by `expected_cells()` has an entry in `threshold.json` -- No threshold key names a non-existent cell (catches typos in threshold key names) - -**Axis 2 — gated-metric coverage** -- Every cell present in both sets has a spec for every `client.` key - (e.g. `client.total_token_throughput`, not `total_token_throughput`) -- Without this, a gated metric with no spec falls through `test_metric`'s `spec is None` - record-only branch and reports PASS with zero assertions even under `enforce_thresholds=true` -- Only checked for cells present in both expected and threshold sets (missing cells are already - reported by axis 1; no double-reporting) - -When `enforce_thresholds=false`: both failures become warnings, not errors. -The config loads as a record-only scaffold (metrics captured, nothing asserted). - -See `docs/cell-key-format.md` for the exact key format used in `threshold.json`. +Loaders in this directory call `substitute_config` then build the typed schema model. +See `cvs/schema/config_file/inference/common/sweep.py` for sweep validation and the +two-axis threshold-coverage check (`validate_thresholds_cover_sweep`). --- diff --git a/cvs/lib/inference/utils/accuracy_config.py b/cvs/lib/inference/utils/accuracy_config.py index 5f22c5faf..b6bc3848f 100644 --- a/cvs/lib/inference/utils/accuracy_config.py +++ b/cvs/lib/inference/utils/accuracy_config.py @@ -1,46 +1,5 @@ -''' -Copyright 2025 Advanced Micro Devices, Inc. -All rights reserved. +'''Copyright 2025 Advanced Micro Devices, Inc. All rights reserved.''' -Accuracy-evaluation config schema, shared across inference suites. +from cvs.schema.config_file.inference.common.accuracy import AccuracyConfig, AccuracyTask -`AccuracyTask`/`AccuracyConfig` define the `config.json`-side selection schema -for lm-eval-harness based accuracy tasks (see cvs/lib/inference/utils/AGENTS.md -for the broader accuracy-evaluation design). This module holds selection only --- no threshold/gating values, which live in the sibling threshold.json file -and are joined against `AccuracyConfig.tasks` at runtime by a later unit. -''' - -from __future__ import annotations - -from typing import Any, Dict, List - -from pydantic import model_validator - -from cvs.lib.utils.config_loader import _Forbid - - -class AccuracyTask(_Forbid): - id: str - task: str - num_fewshot: int = 0 - metadata: Dict[str, Any] = {} - include_path: str = "" - num_concurrent: int = 8 - apply_chat_template: bool = False - gen_kwargs: Dict[str, Any] = {} - - -class AccuracyConfig(_Forbid): - tasks: List[AccuracyTask] = [] - - @model_validator(mode="after") - def _check_unique_task_ids(self): - from collections import Counter - - counts = Counter(t.id for t in self.tasks) - dupes = sorted(i for i, n in counts.items() if n > 1) - if dupes: - rendered = ", ".join(repr(d) for d in dupes) - raise ValueError(f"duplicate task id(s): {rendered}") - return self +__all__ = ["AccuracyConfig", "AccuracyTask"] diff --git a/cvs/lib/inference/utils/functional_config.py b/cvs/lib/inference/utils/functional_config.py index 6965d1d70..1ea71bbc4 100644 --- a/cvs/lib/inference/utils/functional_config.py +++ b/cvs/lib/inference/utils/functional_config.py @@ -1,15 +1,5 @@ -''' -Copyright 2025 Advanced Micro Devices, Inc. -All rights reserved. +'''Copyright 2025 Advanced Micro Devices, Inc. All rights reserved.''' -Optional functional smoke gates for inference suites (FUNC-1/2). -''' +from cvs.schema.config_file.inference.common.functional import FunctionalConfig -from __future__ import annotations - -from cvs.lib.utils.config_loader import _Forbid - - -class FunctionalConfig(_Forbid): - api_smoke: bool = False - health_check: bool = False +__all__ = ["FunctionalConfig"] diff --git a/cvs/lib/inference/utils/inferencing_config_loader.py b/cvs/lib/inference/utils/inferencing_config_loader.py deleted file mode 100644 index a48926a46..000000000 --- a/cvs/lib/inference/utils/inferencing_config_loader.py +++ /dev/null @@ -1,224 +0,0 @@ -''' -Copyright 2025 Advanced Micro Devices, Inc. -All rights reserved. - -Inference-specific config schema for the vllm_single suite. - -The framework-agnostic machinery (paths/model/image/container schema, the -3-pass placeholder substitution, the `enforce_thresholds` gate, and the -`substitute_config` file-read helper) lives in `cvs.lib.utils.config_loader`. -This module holds the inference half: the sweep selector -(`SeqCombo`/`Run`/`Sweep`), the goodput SLO, the framework `Params`, the -`server` role, and `VariantConfig(BaseVariantConfig)` with the -ISL/OSL/TP/CONC `cell_key` and its sweep-coverage check. - -`Sweep`/`SeqCombo`/`GoodputSlo`/`Roles`/`cell_key` are inference-generic (any -serving framework sweeps sequence shapes at concurrencies); only `Params` (the -`vllm bench serve` flags) is framework-flavored and is the seam to subclass -when a second serving framework lands. -''' - -from __future__ import annotations - -import warnings -from collections import Counter -from typing import Any, Dict, List, Optional - -from pydantic import model_validator -from typing_extensions import Literal - -from cvs.lib.utils.config_loader import BaseVariantConfig, _Forbid, substitute_config -from cvs.lib.inference.utils.vllm_parsing import GATED_METRICS - - -# ---------- pydantic models (inference) ---------- - - -class RoleServer(_Forbid): - # The `vllm serve` command is built in Python (cvs.lib.inference.vllm_single), - # not cloned from a `.sh` script, so a run is self-contained. Per-model - # server quirks live here: extra `vllm serve` flags (serve_args) and env vars - # merged over the defaults the orchestrator sets. Both default empty; the - # fp8-kv cell sets its --kv-cache-dtype via serve_args, kept out of the - # generic driver so it stays model-agnostic. A {flag: value} map (flag - # without the leading --): a scalar renders `--flag value`, True renders a - # bare `--flag`, and a list renders the flag once per element. Cleaner than - # a flat [flag, value, flag, value] list and still covers vllm's bare and - # repeatable flags. - serve_args: Dict[str, Any] = {} - env: Dict[str, str] = {} - - -class Roles(_Forbid): - server: RoleServer = RoleServer() - - -class GoodputSlo(_Forbid): - # Per-cell SLOs for the goodput gate, in milliseconds. An INPUT to the run - # (passed to `vllm bench serve --goodput`), NOT a threshold to assert -- so - # it lives in the sweep, not threshold.json. Attached per seq_combo because - # e2el scales ~linearly with osl. _Forbid: a typo'd key fails load, not a - # silently-dropped SLO. - ttft_ms: float - tpot_ms: float - e2el_ms: float - - -class SeqCombo(_Forbid): - # `name` is the join key the `runs` selector references; required. - name: str - isl: str - osl: str - goodput_slo: Optional[GoodputSlo] = None - - -class Run(_Forbid): - # One sweep cell: a named combo run at a single concurrency. The explicit - # list of Runs replaces the old `sequence_combinations x concurrency_levels` - # cartesian -- you enumerate exactly the cells you want, no NxM explosion. - combo: str - concurrency: int - - -NON_SWEEP_THRESHOLD_KEYS = {"accuracy", "mtp_quality", "long_context_accuracy", "quant_parity"} - - -def validate_thresholds_cover_sweep( - *, - expected_cells, - thresholds, - enforce_thresholds: bool, - gated_metrics=None, - gated_gpu_metrics=None, -) -> None: - """Shared sweep/threshold coverage check for inference variant configs.""" - expected = set(expected_cells) - present = set(thresholds.keys()) - NON_SWEEP_THRESHOLD_KEYS - missing = sorted(expected - present) - extra = sorted(present - expected) - problems = [] - if missing: - problems.append(f"sweep cells with no threshold entry: {missing}") - if extra: - problems.append(f"threshold keys matching no sweep cell (typo?): {extra}") - gated = gated_metrics if gated_metrics is not None else GATED_METRICS - gated_keys = [f"client.{m}" for m in sorted(gated)] - if gated_gpu_metrics: - gated_keys += [f"gpu.{m}" for m in sorted(gated_gpu_metrics)] - gated_gaps = {} - for cell in sorted(expected & present): - specs = thresholds.get(cell) or {} - absent = [k for k in gated_keys if k not in specs] - if absent: - gated_gaps[cell] = absent - if gated_gaps: - problems.append(f"cells missing gated-metric specs: {gated_gaps}") - if problems: - msg = "threshold.json does not match the sweep matrix; " + "; ".join(problems) - if enforce_thresholds: - raise ValueError(msg) - warnings.warn(f"{msg} (enforce_thresholds=false -> record-only)", stacklevel=3) - - -def validate_sweep_selector(combo_names, run_combo_refs): - """The sweep-selector rule: combo names unique, every run.combo names one. - - The single home for this check, shared by the typed `Sweep` validator (load - time) and `pytest_generate_tests` (collection time, which reads raw JSON - before the loader runs) so the two can never drift. Operates on plain lists - of strings so both call sites can feed it. - - Without it a duplicate name silently shadows a combo and a typo'd - `run.combo` is a silently-dropped cell -- either way the sweep runs a - different matrix than the config reads. - """ - counts = Counter(combo_names) - dupes = sorted(name for name, count in counts.items() if count > 1) - if dupes: - raise ValueError(f"duplicate sequence_combination names: {dupes}") - known = set(counts) - unknown = sorted({r for r in run_combo_refs if r not in known}) - if unknown: - raise ValueError(f"run.combo names no sequence_combination: {unknown} (known: {sorted(known)})") - - -class Sweep(_Forbid): - sequence_combinations: List[SeqCombo] - runs: List[Run] - - @model_validator(mode="after") - def _check_runs_reference_known_combos(self): - validate_sweep_selector( - [c.name for c in self.sequence_combinations], - [r.combo for r in self.runs], - ) - return self - - -class Params(_Forbid): - backend: str = "vllm" - base_url: str = "http://0.0.0.0" - port_no: str = "8888" - dataset_name: str = "random" - burstiness: str = "1.0" - seed: str = "0" - request_rate: str = "inf" - random_range_ratio: str = "0.8" - random_prefix_len: str = "0" - tensor_parallelism: str = "1" - tokenizer_mode: str = "auto" - percentile_metrics: str = "ttft,tpot,itl,e2el" - metric_percentiles: str = "50,90,95,99" - num_prompts: str = "3200" - # Completion-poll budget for the bench client = client_poll_count * 60s - # (plus a 120s initial wait). Large-output cells (high osl) need a bigger - # budget; the poll loop exits as soon as the client finishes, so raising - # this never slows down fast cells. See regressions REG-20260609-001. - client_poll_count: str = "20" - - -class VariantConfig(BaseVariantConfig): - framework: Literal["vllm_single"] - gpu_arch: str - roles: Roles = Roles() - params: Params - sweep: Sweep - - def cell_key(self, isl, osl, concurrency): - """The canonical threshold key for one sweep cell. - - Single source of truth shared by the loader's coverage check and the - test's verdict lookup -- so the two can never drift on whitespace, - ordering, or field names. - """ - return f"ISL={isl},OSL={osl},TP={self.params.tensor_parallelism},CONC={concurrency}" - - def expected_cells(self): - """Every (isl, osl, conc) cell the sweep's `runs` selector picks.""" - by_name = {c.name: c for c in self.sweep.sequence_combinations} - return [self.cell_key(by_name[r.combo].isl, by_name[r.combo].osl, r.concurrency) for r in self.sweep.runs] - - @model_validator(mode="after") - def _check_thresholds_cover_sweep(self): - """Fail at load time if any sweep cell lacks a threshold entry.""" - validate_thresholds_cover_sweep( - expected_cells=self.expected_cells(), - thresholds=self.thresholds, - enforce_thresholds=self.enforce_thresholds, - ) - return self - - -# ---------- public API (inference) ---------- - - -def load_variant(config_path, cluster_dict): - """Load and validate a vllm_single variant config + its sibling threshold file. - - Delegates the file read + placeholder substitution + threshold discovery to - the generic `substitute_config`, then attaches the thresholds and builds the - typed `VariantConfig`. - """ - raw, thresholds = substitute_config(config_path, cluster_dict) - raw["thresholds"] = thresholds - return VariantConfig(**raw) diff --git a/cvs/lib/inference/utils/long_context_accuracy_config.py b/cvs/lib/inference/utils/long_context_accuracy_config.py index e87b2d606..b08ddf49b 100644 --- a/cvs/lib/inference/utils/long_context_accuracy_config.py +++ b/cvs/lib/inference/utils/long_context_accuracy_config.py @@ -1,38 +1,8 @@ -''' -Copyright 2025 Advanced Micro Devices, Inc. -All rights reserved. +'''Copyright 2025 Advanced Micro Devices, Inc. All rights reserved.''' -Long-context accuracy cell selection (ACC-12 NIAH) for inference suites. -Threshold/gating values live in threshold.json under ``long_context_accuracy``. -''' +from cvs.schema.config_file.inference.common.long_context_accuracy import ( + LongContextAccCell, + LongContextAccuracyConfig, +) -from __future__ import annotations - -from typing import List - -from pydantic import model_validator - -from cvs.lib.utils.config_loader import _Forbid - - -class LongContextAccCell(_Forbid): - id: str - isl: int - osl: int = 32 - num_prompts: int = 8 - seed: int = 42 - - -class LongContextAccuracyConfig(_Forbid): - cells: List[LongContextAccCell] = [] - - @model_validator(mode="after") - def _check_unique_cell_ids(self): - from collections import Counter - - counts = Counter(c.id for c in self.cells) - dupes = sorted(i for i, n in counts.items() if n > 1) - if dupes: - rendered = ", ".join(repr(d) for d in dupes) - raise ValueError(f"duplicate long-context cell id(s): {rendered}") - return self +__all__ = ["LongContextAccCell", "LongContextAccuracyConfig"] diff --git a/cvs/lib/inference/utils/platform_config.py b/cvs/lib/inference/utils/platform_config.py index e97e1b829..bcd8af0d3 100644 --- a/cvs/lib/inference/utils/platform_config.py +++ b/cvs/lib/inference/utils/platform_config.py @@ -1,15 +1,5 @@ -''' -Copyright 2025 Advanced Micro Devices, Inc. -All rights reserved. +'''Copyright 2025 Advanced Micro Devices, Inc. All rights reserved.''' -Optional platform / post-run checks for inference suites (INF-6/7). -''' +from cvs.schema.config_file.inference.common.platform import PlatformConfig -from __future__ import annotations - -from cvs.lib.utils.config_loader import _Forbid - - -class PlatformConfig(_Forbid): - dmesg_scan: bool = False - gpu_metrics_poll: bool = False +__all__ = ["PlatformConfig"] diff --git a/cvs/lib/inference/utils/vllm_config_loader.py b/cvs/lib/inference/utils/vllm_config_loader.py index 2314774b9..42cd85766 100644 --- a/cvs/lib/inference/utils/vllm_config_loader.py +++ b/cvs/lib/inference/utils/vllm_config_loader.py @@ -2,273 +2,45 @@ Copyright 2025 Advanced Micro Devices, Inc. All rights reserved. -Unified config schema for the vllm suite (single-node and distributed). +Load and validate vLLM inference variant configs. -Replaces inferencing_config_loader.py (single-node) and -vllm_distributed_config_loader.py (distributed) with a single schema. -Distributed params (pipeline_parallel_size, master_addr, master_port, -nnodes) default to single-node values so the same VariantConfig works for -both topologies. - -cell_key format: - Single-node (pp=1): ISL=,OSL=,TP=,CONC= - Distributed (pp>1): ISL=,OSL=,TP=,PP=,CONC= - -IB device config: - roles.server.ib_hca_devices: list[str] | "auto" | absent - If absent or "auto", use everything ibv_devinfo -l reports. - If an explicit list, validate at preflight (test_discover_topology). - roles.server.ib_netdev: str (required for distributed runs) - Linux network interface name for NCCL_SOCKET_IFNAME / GLOO_SOCKET_IFNAME. - Not derivable from HCA names. Operator sets it explicitly. - Optional for single-node (NCCL socket selection not critical). +Pydantic models live in ``cvs.schema.config_file.inference.vllm.variant``. ''' -from __future__ import annotations - -from collections import Counter -from typing import Any, Dict, List, Optional, Union - -from pydantic import BaseModel, ConfigDict, Field, field_validator, model_validator -from typing_extensions import Literal - -from cvs.lib.inference.utils.accuracy_config import AccuracyConfig -from cvs.lib.inference.utils.inferencing_config_loader import validate_thresholds_cover_sweep from cvs.lib.inference.utils.vllm_server_metrics import PROM_METRICS from cvs.lib.utils.config_loader import substitute_config from cvs.lib.utils.gpu import GPU_METRICS +from cvs.schema.config_file.inference.common.sweep import ( + GoodputSlo, + Run, + SeqCombo, + Sweep, + validate_sweep_selector, +) +from cvs.schema.config_file.inference.vllm.variant import VariantConfig, VllmRoleServer GATED_GPU_METRICS = {k for k, _unit in GPU_METRICS} -# A fully separate, parallel gated family, following GPU_METRICS's precedent -# rather than joining vllm_parsing.GATED_METRICS/METRIC_TIERS -- prom.* must -# not be mixed into the client.* tiering machinery (a locked invariant test -# partitions that set exactly). GATED_PROM_METRICS = {k for k, _unit in PROM_METRICS} +# Backward-compat alias used by server-reuse tests. +RoleServer = VllmRoleServer -class _Forbid(BaseModel): - model_config = ConfigDict(extra="forbid") - - -class _Allow(BaseModel): - model_config = ConfigDict(extra="allow") - - -# ---------- sub-models ---------- - - -class ContainerConfig(_Allow): - lifetime: str = "per_run" - name: str = "" - image: str = "" - - -class Paths(_Forbid): - shared_fs: str - models_dir: str - log_dir: str - hf_token_file: str - - -class ModelSpec(_Forbid): - id: str - remote: Literal[0, 1] - - -_VLLM_LOG_LEVELS = {"debug", "info", "warning", "error", "critical"} - - -class RoleServer(_Forbid): - serve_args: Dict[str, Any] = {} - env: Dict[str, str] = {} - # IB HCA devices for NCCL_IB_HCA. - # absent or "auto" -> use whatever ibv_devinfo -l reports. - # explicit list -> validated at preflight against ibv_devinfo output. - ib_hca_devices: Union[Literal["auto"], List[str], None] = None - # Linux netdev for NCCL_SOCKET_IFNAME / GLOO_SOCKET_IFNAME. - # Required when nnodes > 1. No "auto" — not reliably derivable from HCA names. - ib_netdev: Optional[str] = None - - @field_validator("serve_args", mode="after") - @classmethod - def _check_log_level(cls, v): - level = v.get("log-level") - if level is not None and level not in _VLLM_LOG_LEVELS: - raise ValueError(f"serve_args.log-level must be one of {sorted(_VLLM_LOG_LEVELS)}, got: {level!r}") - return v - - -class Roles(_Forbid): - server: RoleServer = Field(default_factory=RoleServer) - - -class GoodputSlo(_Forbid): - ttft_ms: float - tpot_ms: float - e2el_ms: float - - -class SeqCombo(_Forbid): - name: str - isl: str - osl: str - goodput_slo: Optional[GoodputSlo] = None - - -class Run(_Forbid): - combo: str - concurrency: int - - -def validate_sweep_selector(combo_names, run_combo_refs): - """Single home for the sweep-selector rule: names unique, every run.combo known. - - Called both at load time (via Sweep model_validator) and at collection time - (pytest_generate_tests reads raw JSON before load_variant runs) so the two - paths cannot drift. - """ - counts = Counter(combo_names) - dupes = sorted(name for name, count in counts.items() if count > 1) - if dupes: - raise ValueError(f"duplicate sequence_combination names: {dupes}") - known = set(counts) - unknown = sorted({r for r in run_combo_refs if r not in known}) - if unknown: - raise ValueError(f"run.combo names no sequence_combination: {unknown} (known: {sorted(known)})") - - -class Sweep(_Forbid): - sequence_combinations: List[SeqCombo] - runs: List[Run] - - @model_validator(mode="after") - def _check_runs_reference_known_combos(self): - validate_sweep_selector( - [c.name for c in self.sequence_combinations], - [r.combo for r in self.runs], - ) - return self - - -class Params(_Forbid): - backend: str = "vllm" - base_url: str = "http://0.0.0.0" - port_no: str = "8888" - dataset_name: str = "random" - burstiness: str = "1.0" - seed: str = "0" - request_rate: str = "inf" - random_range_ratio: str = "0.8" - random_prefix_len: str = "0" - tensor_parallelism: str = "8" - # Distributed params. Defaults encode single-node (no PP, one node, localhost). - pipeline_parallel_size: str = "1" - master_addr: str = "localhost" - master_port: str = "29501" - nnodes: str = "1" - tokenizer_mode: str = "auto" - percentile_metrics: str = "ttft,tpot,itl,e2el" - metric_percentiles: str = "50,90,95,99" - num_prompts: str = "3200" - client_poll_count: str = "20" - - -class VariantConfig(_Forbid): - """Unified typed config for both single-node and distributed vllm runs. - - Standalone (does not extend BaseVariantConfig) so it can be constructed - without the threshold_json field the base requires — absent from unit-test - fixtures. Production configs always supply it via substitute_config. - - ``container`` is optional at model-level: unit-test fixtures omit it. - The conftest ``orch`` fixture accesses ``variant_config.container.model_dump()``. - """ - - schema_version: Literal[1] - framework: Literal["vllm"] - gpu_arch: str - enforce_thresholds: bool = True - container: ContainerConfig = Field(default_factory=ContainerConfig) - paths: Paths - model: ModelSpec - roles: Roles = Field(default_factory=Roles) - params: Params = Field(default_factory=Params) - sweep: Sweep - thresholds: Dict[str, Dict[str, Any]] = Field(default_factory=dict) - accuracy: AccuracyConfig = Field(default_factory=AccuracyConfig) - - @model_validator(mode="after") - def _check_distributed_consistency(self): - nn = int(self.params.nnodes) - pp = int(self.params.pipeline_parallel_size) - # Ray backend uses its own distributed orchestration and does not require - # pipeline parallelism (pp=1 is the expected ray multi-node configuration). - # Only the exact lowercase string "ray" triggers this relaxation (AC6). - is_ray = self.roles.server.serve_args.get("distributed-executor-backend") == "ray" - if nn > 1 and pp == 1 and not is_ray: - raise ValueError(f"nnodes={nn} > 1 requires pipeline_parallel_size > 1 (got pp={pp})") - if pp > 1 and nn == 1: - raise ValueError(f"pipeline_parallel_size={pp} > 1 requires nnodes > 1 (got nnodes={nn})") - if nn > 1 and not self.roles.server.ib_netdev: - raise ValueError( - "ib_netdev is required in roles.server when nnodes > 1. " - "Set it to the Linux network interface name for NCCL_SOCKET_IFNAME " - "(e.g. \"ens51f1np1\"). Cannot be auto-derived from HCA names." - ) - return self - - @model_validator(mode="after") - def _check_remote_not_implemented(self): - if self.model.remote == 1: - raise NotImplementedError("model.remote=1 (remote model download) is not implemented.") - return self - - def cell_key(self, isl, osl, concurrency): - """Canonical threshold key for one sweep cell. - - Emits PP= segment only for distributed runs (pp > 1), preserving - backward-compatible single-node keys (no PP= segment). - - Single-node: ISL=,OSL=,TP=,CONC= - Distributed: ISL=,OSL=,TP=,PP=,CONC= - """ - base = f"ISL={isl},OSL={osl},TP={self.params.tensor_parallelism}," - if int(self.params.pipeline_parallel_size) > 1: - base += f"PP={self.params.pipeline_parallel_size}," - return base + f"CONC={concurrency}" - - def expected_cells(self): - by_name = {c.name: c for c in self.sweep.sequence_combinations} - return [self.cell_key(by_name[r.combo].isl, by_name[r.combo].osl, r.concurrency) for r in self.sweep.runs] - - @model_validator(mode="after") - def _check_thresholds_cover_sweep(self): - """Every sweep cell must have a threshold entry; no metric within it is - mandatory. Evaluation (``test_metric``/``test_gpu_metric``/ - ``test_prom_metric``) already treats an absent ``client.*``/``gpu.*``/ - ``prom.*`` spec as "don't gate this metric" (skips the assertion), so - a threshold.json is free to gate only the handful of metrics an - operator cares about instead of every member of every family. - """ - validate_thresholds_cover_sweep( - expected_cells=self.expected_cells(), - thresholds=self.thresholds, - enforce_thresholds=self.enforce_thresholds, - gated_metrics=set(), - ) - return self - - -# ---------- public API ---------- +__all__ = [ + "GATED_GPU_METRICS", + "GATED_PROM_METRICS", + "GoodputSlo", + "RoleServer", + "Run", + "SeqCombo", + "Sweep", + "VariantConfig", + "VllmRoleServer", + "load_variant", + "validate_sweep_selector", +] def load_variant(config_path, cluster_dict): - """Load and validate a vllm variant config + its sibling threshold file. - - Strips fields unknown to VariantConfig before construction so production - configs (which carry extra keys like threshold_json) and unit-test fixtures - (which omit optional fields) are handled identically. - """ raw, thresholds = substitute_config(config_path, cluster_dict) known = {k: v for k, v in raw.items() if k in VariantConfig.model_fields} known["thresholds"] = thresholds diff --git a/cvs/lib/preflight/unittests/test_ifoe_l2_connectivity.py b/cvs/lib/preflight/unittests/test_ifoe_l2_connectivity.py index e419e8b37..19a3d5590 100644 --- a/cvs/lib/preflight/unittests/test_ifoe_l2_connectivity.py +++ b/cvs/lib/preflight/unittests/test_ifoe_l2_connectivity.py @@ -20,7 +20,7 @@ parse_afmctl_show_device_json, ) from cvs.lib.preflight.report import PreflightReportGenerator -from cvs.parsers.schemas import PreflightConfigFile +from cvs.schema.config_file.preflight.config import PreflightConfigFile PASSING_OUTPUT = """\ diff --git a/cvs/lib/preflight/unittests/test_node_smoke.py b/cvs/lib/preflight/unittests/test_node_smoke.py index 340e4919f..8ff4b6c39 100644 --- a/cvs/lib/preflight/unittests/test_node_smoke.py +++ b/cvs/lib/preflight/unittests/test_node_smoke.py @@ -276,7 +276,7 @@ def test_node_smoke_tier_summaries_use_tier_labels(self): class TestLegacyNodeSmokeConfigNormalization(unittest.TestCase): def test_legacy_node_smoke_copied_to_tier1(self): - from cvs.parsers.schemas import normalize_legacy_preflight_node_smoke_sections + from cvs.schema.config_file.preflight.config import normalize_legacy_preflight_node_smoke_sections cfg = {"node_smoke": {"connectivity_mode": "run", "primus_dir": "/home/user/Primus"}} normalized, warning = normalize_legacy_preflight_node_smoke_sections(cfg) @@ -285,7 +285,7 @@ def test_legacy_node_smoke_copied_to_tier1(self): self.assertEqual(normalized["node_smoke"]["primus_dir"], "/home/user/Primus") def test_canonical_tier1_not_overwritten_by_legacy(self): - from cvs.parsers.schemas import normalize_legacy_preflight_node_smoke_sections + from cvs.schema.config_file.preflight.config import normalize_legacy_preflight_node_smoke_sections cfg = { "node_smoke_tier1": {"primus_dir": "/tier1/Primus"}, diff --git a/cvs/lib/preflight/unittests/test_rdma_connectivity.py b/cvs/lib/preflight/unittests/test_rdma_connectivity.py index d9d402bdf..d278094e1 100644 --- a/cvs/lib/preflight/unittests/test_rdma_connectivity.py +++ b/cvs/lib/preflight/unittests/test_rdma_connectivity.py @@ -15,7 +15,7 @@ from cvs.lib.preflight.base import partition_nodes_into_groups from cvs.lib.preflight.rdma_connectivity import RdmaConnectivityCheck from cvs.lib.preflight.report import PreflightReportGenerator -from cvs.parsers.schemas import PreflightConfigFile, normalize_legacy_preflight_rdma_config +from cvs.schema.config_file.preflight.config import PreflightConfigFile, normalize_legacy_preflight_rdma_config def _make_checker( diff --git a/cvs/lib/preflight/unittests/test_scaleup_fabric.py b/cvs/lib/preflight/unittests/test_scaleup_fabric.py index fa1aa2e25..fb11648b8 100644 --- a/cvs/lib/preflight/unittests/test_scaleup_fabric.py +++ b/cvs/lib/preflight/unittests/test_scaleup_fabric.py @@ -19,7 +19,7 @@ parse_station_masks, ) from cvs.lib.preflight.report import PreflightReportGenerator # noqa: E402 -from cvs.parsers.schemas import PreflightConfigFile # noqa: E402 +from cvs.schema.config_file.preflight.config import PreflightConfigFile # noqa: E402 def _afm_devices(*, phase="ACTIVE", vpod=None): diff --git a/cvs/lib/preflight/unittests/test_transferbench_smoke.py b/cvs/lib/preflight/unittests/test_transferbench_smoke.py index 54e4a339b..87f32de81 100644 --- a/cvs/lib/preflight/unittests/test_transferbench_smoke.py +++ b/cvs/lib/preflight/unittests/test_transferbench_smoke.py @@ -26,7 +26,7 @@ reconcile_cluster_vpod, ) from cvs.lib.preflight.report import PreflightReportGenerator -from cvs.parsers.schemas import PreflightConfigFile +from cvs.schema.config_file.preflight.config import PreflightConfigFile # --------------------------------------------------------------------------- diff --git a/cvs/lib/training/jaxmaxtext/utils/training_config_loader.py b/cvs/lib/training/jaxmaxtext/utils/training_config_loader.py index 5a328f5e7..00b52a06e 100644 --- a/cvs/lib/training/jaxmaxtext/utils/training_config_loader.py +++ b/cvs/lib/training/jaxmaxtext/utils/training_config_loader.py @@ -2,225 +2,46 @@ Copyright 2025 Advanced Micro Devices, Inc. All rights reserved. -Training-specific config schema for the jaxmaxtext suite. +Load and validate JAX MaxText training variant configs from +``cvs/input/config_file/training/jaxmaxtext/``. -The framework-agnostic machinery (paths/model/container schema, the 3-pass -placeholder substitution, the `enforce_thresholds` gate, and the -`substitute_config` file-read helper) lives in `cvs.lib.utils.config_loader`. -This module holds the training half: the MaxText config, tokenizer, NCCL, -JAX distributed settings, RDMA lib, and `TrainingVariantConfig(BaseVariantConfig)`. - -A training suite does not sweep cells the way inference does (no NxM matrix of -ISL/OSL/concurrency). Instead each declared `sweep` is one full training run and -its `name` IS the threshold-file key (also the key `metric()` looks up at -runtime). `expected_cells()` therefore returns the declared sweep names, and the -coverage check validates the threshold file against those names directly. +Pydantic models live in ``cvs.schema.config_file.training.jaxmaxtext.variant``. ''' -from __future__ import annotations - -import warnings -from typing import Any, Dict, List, Literal - -from pydantic import field_validator - -from cvs.lib.utils.config_loader import BaseVariantConfig, _Allow, _Forbid, substitute_config from cvs.lib.training.jaxmaxtext.utils.maxtext_parsing import GATED_METRICS - - -class Tokenizer(_Forbid): - hf_model_id: str - tokenizer_path: str - - -class NcclConfig(_Allow): - ib_hca_list: str = "" - ib_hca: str = "" - socket_ifname: str = "" - gloo_socket_ifname: str = "" - ib_gid_index: str = "3" - - @field_validator("ib_hca_list", "ib_hca", "socket_ifname", "gloo_socket_ifname", "ib_gid_index") - @classmethod - def _reject_changeme(cls, v, info): - """Hard-exit when a cluster-specific RDMA/NIC field is left as ''. - - These device/interface names are cluster-specific and shipped as - '' placeholders (see the sibling _example_* values). Running a - distributed job with them unresolved would silently use the wrong - NIC/RDMA devices, so fail loudly at config load instead. - """ - if isinstance(v, str) and "" in v.lower(): - raise ValueError( - f"nccl.{info.field_name} is still ''. Set your cluster's RDMA/NIC " - "device/interface (see the sibling _example_* value) before running distributed training." - ) - return v - - -class JaxDistributed(_Forbid): - coordinator_ip: str = "auto" - coordinator_port: str = "12346" - initialization_timeout_seconds: str = "1800" - heartbeat_timeout_seconds: str = "900" - - -class RdmaLib(_Allow): - host_source_file: str = "" - container_mount_file: str = "" - container_dest_file: str = "" - - -class ScalingBaseline(_Allow): - """Reference (typically 1-node) throughput for scaling-efficiency %. - - `tokens_per_sec_total` is the TOTAL tokens/sec measured on a prior run of - `num_nodes` nodes (source it from a previous single-node run log). Scaling - efficiency % = throughput_N / ((N / num_nodes) * tokens_per_sec_total) * 100. - - Leave `tokens_per_sec_total` at 0.0 to disable the metric (it then reports - record-only as None instead of gating on an uncalibrated baseline). - """ - - tokens_per_sec_total: float = 0.0 - num_nodes: int = 1 - - -class Convergence(_Allow): - """Target for convergence / time-to-target-accuracy (row 33). - - `target_metric` selects the loss series to converge on: - - "eval_loss" : validation loss (requires eval enabled + parseable) - - "train_loss" : per-step training loss - - "auto" : eval loss when eval points exist, else training loss - - `target_value` is the loss threshold to reach; <= 0 disables the metric - (steps_to_target / time_to_target_seconds report record-only as None). - """ - - target_metric: Literal["auto", "train_loss", "eval_loss"] = "auto" - target_value: float = 0.0 - - -class LossCurve(_Allow): - """Loss-curve (row 32) sampling + pass/fail settings. - - `sample_every` and `milestone_steps` control which per-step losses are kept - for the plotted/asserted curve (keeps short runs non-empty). The verdict is - the least-squares slope of the sampled curve: the run passes when - `slope < max_slope` (default 0.0 = strictly decreasing). `enforce` gates the - test (fail on a non-decreasing curve); set False for record-only. - """ - - sample_every: int = 10 - milestone_steps: List[int] = [100, 500, 1000, 5000] - max_slope: float = 0.0 - enforce: bool = True - - -class SmokeTest(_Allow): - """Smoke test (ENABLED by default). Loads the model and runs `steps` steps - with a small fixed batch/seqlen in BF16, passing only if no error signature - fires (no metric/threshold checks). A failure gates the rest of the suite. - - Set `enabled=false` to SKIP it -- e.g. during iterative experiments where you - don't want the smoke run every time (mirrors checkpoint_resume, but opt-OUT - rather than opt-in). `steps`/`per_device_batch_size`/`max_target_length` tune - the smoke run itself. - """ - - enabled: bool = True - steps: int = 5 - per_device_batch_size: int = 1 - max_target_length: int = 2048 - - -class CheckpointResume(_Allow): - """Checkpoint save + resume test (opt-in; off by default). - - Runs ONE sweep twice: Phase 1 trains `steps_before_ckpt` steps with - checkpointing on (a checkpoint is written at `checkpoint_period`); Phase 2 - resumes from that checkpoint and trains `steps_after_resume` more. Passes - when the resumed run restarts at the checkpoint step and the loss at the - resume boundary matches Phase 1 within `loss_tolerance` (state restored, not - reinitialized). Also benchmarks checkpoint I/O: `checkpoint_save_seconds` / - `checkpoint_load_seconds` are gated against `max_save_seconds` / - `max_load_seconds` when those are > 0 (else record-only). - - `sweep` selects which sweep to use ("" -> first enabled). `smoke_model_overrides` - optionally shrinks the model for a fast run WITHOUT changing the tokenizer/ - vocab (e.g. {"base_num_decoder_layers": 4}); empty -> the config's full model - (real-size checkpoint I/O). - - `delete_ckpt_dir` (default true) removes the checkpoint directory after the - test to free disk space; set it false to keep the checkpoint files for - inspection. - """ - - enabled: bool = False - sweep: str = "" - steps_before_ckpt: int = 6 - steps_after_resume: int = 6 - checkpoint_period: int = 5 - loss_tolerance: float = 0.1 - max_save_seconds: float = 0.0 - max_load_seconds: float = 0.0 - delete_ckpt_dir: bool = True - smoke_model_overrides: Dict[str, Any] = {} - - -class Sweep(_Allow): - """One sweep entry = one full training run with per-run maxtext overrides. - - `name` is the canonical cell key (also the threshold-file key), e.g. - "NNODES=2,STEPS=30,PRECISION=BF16,BATCH=3,GBS=48,SEQLEN=8192". Only the - parameters that actually vary need a `maxtext_overrides` entry (for now just - precision, e.g. FP8 sets `quantization`); everything else falls back to the - base `maxtext_config`. - """ - - name: str - maxtext_overrides: Dict[str, Any] = {} - - -class TrainingConfig(_Allow): - distributed: bool = True - gpus_per_node: int = 8 # do not assume a uniform topology; override per cluster - # Scan host dmesg (all nodes) for GPU/HW/kernel faults over the training - # window. Set false on clusters without passwordless sudo for `dmesg`. - verify_dmesg: bool = True - steps: int = 30 - enable_checkpointing: bool = False - # MaxText moved the train entrypoint across versions; list candidates and the - # job picks whichever exists in the running container (first match wins). - # v26.4+: .../src/maxtext/trainers/pre_train/train.py - # v26.3 and earlier: .../src/MaxText/train.py - train_script_paths: List[str] = [ - "/workspace/maxtext/src/maxtext/trainers/pre_train/train.py", - "/workspace/maxtext/src/MaxText/train.py", - ] - # Deprecated single-path form; kept for backward compatibility and used as a - # final fallback candidate when train_script_paths is empty. - train_script: str = "/workspace/maxtext/src/MaxText/train.py" - maxtext_config: Dict[str, Any] = {} - tokenizer: Tokenizer - nic_type: str = "thor2" - rdma_lib: RdmaLib = RdmaLib() - env_vars: Dict[str, str] = {} - xla_flags: Dict[str, str] = {} - # {name: regex} error signatures scanned in the training log during polling. - # Empty -> the driver falls back to its built-in default set. Lets users - # add/remove signatures per config without touching code. - error_patterns: Dict[str, str] = {} - nccl: NcclConfig = NcclConfig() - jax_distributed: JaxDistributed = JaxDistributed() - scaling_baseline: ScalingBaseline = ScalingBaseline() - convergence: Convergence = Convergence() - loss_curve: LossCurve = LossCurve() - smoke: SmokeTest = SmokeTest() - checkpoint_resume: CheckpointResume = CheckpointResume() - sweeps: List[Sweep] = [] - enabled_sweep_list: List[str] = [] +from cvs.lib.utils.config_loader import substitute_config +from cvs.schema.config_file.training.jaxmaxtext.variant import ( + CheckpointResume, + Convergence, + JaxDistributed, + LossCurve, + NcclConfig, + RdmaLib, + ScalingBaseline, + SmokeTest, + Sweep, + Tokenizer, + TrainingConfig, + TrainingVariantConfig, + validate_thresholds_cover_training as _validate_thresholds_cover_training, +) + +__all__ = [ + "CheckpointResume", + "Convergence", + "JaxDistributed", + "LossCurve", + "NcclConfig", + "RdmaLib", + "ScalingBaseline", + "SmokeTest", + "Sweep", + "Tokenizer", + "TrainingConfig", + "TrainingVariantConfig", + "load_training_variant", + "validate_thresholds_cover_training", +] def validate_thresholds_cover_training( @@ -230,86 +51,19 @@ def validate_thresholds_cover_training( enforce_thresholds: bool, gated_metrics=None, ) -> None: - """Shared training threshold/cell coverage check.""" - expected = set(expected_cells) - # Skip "_"-prefixed metadata keys (e.g. "_comment") so they are not mistaken - # for a threshold cell that matches no training sweep. - present = {k for k in thresholds.keys() if not str(k).startswith("_")} - missing = sorted(expected - present) - extra = sorted(present - expected) - problems = [] - if missing: - problems.append(f"training cells with no threshold entry: {missing}") - if extra: - problems.append(f"threshold keys matching no training cell (typo?): {extra}") - gated = gated_metrics if gated_metrics is not None else GATED_METRICS - gated_keys = [f"training.{m}" for m in sorted(gated)] - gated_gaps = {} - for cell in sorted(expected & present): - specs = thresholds.get(cell) or {} - absent = [k for k in gated_keys if k not in specs] - if absent: - gated_gaps[cell] = absent - if gated_gaps: - problems.append(f"cells missing gated-metric specs: {gated_gaps}") - if problems: - msg = "threshold.json does not match the training config; " + "; ".join(problems) - if enforce_thresholds: - raise ValueError(msg) - warnings.warn(f"{msg} (enforce_thresholds=false -> record-only)", stacklevel=3) - - -class TrainingVariantConfig(BaseVariantConfig): - framework: Literal["jaxmaxtext"] - gpu_arch: str - training: TrainingConfig - - def expected_cells(self): - """Threshold cell keys this config expects: one per declared sweep. - - The sweep `name` IS the threshold-file key and the key `metric()` looks - up at runtime (see cvs/tests/training/jaxmaxtext/_common.py::metric), so - coverage is checked against the declared sweep names directly -- not a - synthesized key. `enabled_sweep_list` only selects which of these - actually run; the threshold file still carries an entry per declared - sweep. A config with no sweeps degrades to a single implicit "default" - cell. - """ - names = [s.name for s in self.training.sweeps] - return names or ["default"] - - def enabled_sweeps(self): - """Return the Sweep objects selected to run. - - `enabled_sweep_list` (if non-empty) selects a subset by name; otherwise - every declared sweep runs. A config with no `sweeps` degrades to a single - implicit sweep named "default" (its threshold cell, if any, is keyed - "default"), so the suite still runs unparametrized. - """ - sweeps = self.training.sweeps - if not sweeps: - return [Sweep(name="default")] - by_name = {s.name: s for s in sweeps} - names = self.training.enabled_sweep_list or [s.name for s in sweeps] - selected = [] - for n in names: - if n in by_name: - selected.append(by_name[n]) - else: - warnings.warn(f"enabled_sweep_list references unknown sweep '{n}'", stacklevel=2) - return selected - - -# ---------- public API (training) ---------- - - -def load_training_variant(config_path, cluster_dict): - """Load and validate a jaxmaxtext variant config + its sibling threshold file. - - Delegates the file read + placeholder substitution + threshold discovery to - the generic `substitute_config`, then attaches the thresholds and builds the - typed `TrainingVariantConfig`. - """ + """Training threshold coverage check; defaults ``gated_metrics`` to MaxText gated set.""" + if gated_metrics is None: + gated_metrics = GATED_METRICS + return _validate_thresholds_cover_training( + expected_cells=expected_cells, + thresholds=thresholds, + enforce_thresholds=enforce_thresholds, + gated_metrics=gated_metrics, + ) + + +def load_training_variant(config_path, cluster_dict) -> TrainingVariantConfig: + """Load and validate a jaxmaxtext variant config + its sibling threshold file.""" raw, thresholds = substitute_config(config_path, cluster_dict) raw["thresholds"] = thresholds return TrainingVariantConfig(**raw) diff --git a/cvs/lib/training/megatron/utils/training_config_loader.py b/cvs/lib/training/megatron/utils/training_config_loader.py index 877a69591..fd12b8deb 100644 --- a/cvs/lib/training/megatron/utils/training_config_loader.py +++ b/cvs/lib/training/megatron/utils/training_config_loader.py @@ -2,201 +2,18 @@ Copyright 2025 Advanced Micro Devices, Inc. All rights reserved. -Training-specific config schema for Megatron suites (single-node and distributed). +Load and validate Megatron training variant configs from +``cvs/input/config_file/training/megatron/``. -The framework-agnostic machinery (ContainerSpec, RuntimeSpec, placeholder -substitution, threshold file discovery) lives in `cvs.lib.utils.config_loader`. -This module holds the training half: MegatronSweepCombo, MegatronSweep, -MegatronVariantConfig, and load_training_variant. - -Thresholds live in a sibling *threshold.json file (not inline in result_dict). -The threshold file is discovered via the `threshold_json` field in the config or -auto-discovered as the sole *threshold.json sibling. Cell keys in the threshold -file must match the combination keys in sweep.combinations exactly. - -enforce_thresholds gates whether threshold specs are asserted in test_metric. - -Both megatron_single and megatron_distributed are covered by MegatronVariantConfig -via the framework field, which is a validated schema tag / config discriminator. +Pydantic models live in ``cvs.schema.config_file.training.megatron.variant``. ''' -from __future__ import annotations - -import warnings -from collections import Counter -from typing import Any, Dict, List - -from pydantic import Field, model_validator -from typing_extensions import Literal - -from cvs.lib.utils.config_loader import ( - ContainerSpec, - _Forbid, - substitute_config, -) - - -# ---------- pydantic models (training) ---------- - - -class MegatronSweepCombo(_Forbid): - name: str - micro_batch_size: str - global_batch_size: str - precision: str = "" - - -def validate_sweep_selector(combo_keys, run_refs): - """The sweep-selector rule: combination keys unique, every run references one. - - Single home for this check, shared by the typed MegatronSweep validator - (load time) and pytest_generate_tests (collection time, which reads raw - JSON before the loader runs) so the two can never drift. - - Without it a typo'd run key is a silently-dropped cell — the sweep runs - a different matrix than the config reads. - """ - counts = Counter(combo_keys) - dupes = sorted(k for k, count in counts.items() if count > 1) - if dupes: - raise ValueError(f"duplicate sweep.combinations keys: {dupes}") - known = set(counts) - unknown = sorted(r for r in run_refs if r not in known) - if unknown: - raise ValueError(f"sweep.runs references unknown combinations: {unknown} (known: {sorted(known)})") - - -def validate_thresholds_cover_sweep( - *, - expected_cells, - thresholds, - enforce_thresholds: bool, - gated_metrics=None, -) -> None: - """Shared sweep/threshold coverage check for training variant configs. - - Checks every sweep cell has a threshold entry and no threshold key is - orphaned. Individual metrics within a cell are optional — absent specs - are skipped in test_metric (record-only for that metric). - """ - expected = set(expected_cells) - present = set(thresholds.keys()) - missing = sorted(expected - present) - extra = sorted(present - expected) - problems = [] - if missing: - problems.append(f"sweep cells with no threshold entry: {missing}") - if extra: - problems.append(f"threshold keys matching no sweep cell (typo?): {extra}") - gated = gated_metrics if gated_metrics is not None else set() - gated_keys = [f"training.{m}" for m in sorted(gated)] - gated_gaps = {} - for cell in sorted(expected & present): - specs = thresholds.get(cell) or {} - absent = [k for k in gated_keys if k not in specs] - if absent: - gated_gaps[cell] = absent - if gated_gaps: - problems.append(f"cells missing gated-metric specs: {gated_gaps}") - if problems: - msg = "threshold.json does not match the sweep matrix; " + "; ".join(problems) - if enforce_thresholds: - raise ValueError(msg) - warnings.warn(f"{msg} (enforce_thresholds=false -> record-only)", stacklevel=3) - - -class MegatronSweep(_Forbid): - combinations: Dict[str, MegatronSweepCombo] - runs: List[str] - - @model_validator(mode="after") - def _check_runs_reference_known_combos(self): - validate_sweep_selector( - list(self.combinations.keys()), - self.runs, - ) - return self - - -class ScalingBaseline(_Forbid): - tokens_per_sec_total: float = 0.0 - num_nodes: int = 1 - - -class LossCurveConfig(_Forbid): - sample_every: int = 10 - milestone_steps: List[int] = Field(default_factory=lambda: [100, 500, 1000, 5000]) - max_slope: float = 0.0 - enforce: bool = True - - -class ConvergenceConfig(_Forbid): - target_metric: Literal["auto", "train_loss", "eval_loss"] = "auto" - target_value: float = 0.0 - - -class CheckpointConfig(_Forbid): - enforce: bool = False # if False, test_checkpoint is skipped entirely - save_interval: int = 20 # checkpoint written every N steps - save_iters: int = 21 # save phase total; last checkpoint = floor(save/interval)*interval - resume_iters: int = 25 # load phase total (must be > last_ckpt_step) - loss_rtol: float = 0.05 # max allowed fractional loss increase across boundary - checkpoint_dir: str = "" # shared path for distributed; empty = derive from log_dir (single-node) - - -class MegatronVariantConfig(_Forbid): - schema_version: Literal[1] - framework: Literal["megatron_single", "megatron_distributed"] - gpu_arch: str - enforce_thresholds: bool = True - threshold_json: str = "" - scaling_baseline: ScalingBaseline = Field(default_factory=ScalingBaseline) - loss_curve: LossCurveConfig = Field(default_factory=LossCurveConfig) - convergence: ConvergenceConfig = Field(default_factory=ConvergenceConfig) - checkpoint: CheckpointConfig = Field(default_factory=CheckpointConfig) - config: Dict[str, Any] # training knobs: megatron_root, nccl_*, nic_type, ... - model_params: Dict[str, Any] # model knobs: model_name, precision, tp, pp, ... - container: ContainerSpec - sweep: MegatronSweep - thresholds: Dict[str, Dict[str, Any]] = Field(default_factory=dict) - - def cell_key(self, combo_key: str) -> str: - """Canonical threshold lookup key for a sweep combo. - - Constructs a key from the combo's micro_batch_size, global_batch_size, - and precision — must match the top-level keys in the threshold file exactly. - """ - combo = self.sweep.combinations[combo_key] - return f"MBS={combo.micro_batch_size},GBS={combo.global_batch_size},PRECISION={combo.precision}" - - def expected_cells(self) -> List[str]: - """Return the threshold cell key for every run in sweep.runs.""" - return [self.cell_key(k) for k in self.sweep.runs] - - @model_validator(mode="after") - def _check_thresholds_cover_sweep(self): - """Every sweep cell must have a threshold entry; no metric within it is - mandatory. test_metric treats an absent ``training.*`` spec as - "don't gate this metric" (skips the assertion), so a threshold.json - is free to gate only the metrics an operator cares about. - """ - validate_thresholds_cover_sweep( - expected_cells=self.expected_cells(), - thresholds=self.thresholds, - enforce_thresholds=self.enforce_thresholds, - gated_metrics=set(), - ) - return self - - -# ---------- public API (training) ---------- +from cvs.lib.utils.config_loader import substitute_config +from cvs.schema.config_file.training.megatron.variant import MegatronVariantConfig def _check_no_changeme(node, path="", _offenders=None): - """Recursively collect config fields whose value still contains ''. - - Collects all offending dotted paths so the caller can report them all at once. - """ + """Recursively collect config fields whose value still contains ''.""" if _offenders is None: _offenders = [] if isinstance(node, dict): @@ -213,23 +30,9 @@ def _check_no_changeme(node, path="", _offenders=None): def load_training_variant(config_path, cluster_dict) -> MegatronVariantConfig: - """Load and validate a Megatron training variant config + its threshold file. - - Delegates file read, placeholder substitution, and threshold file discovery - to the generic substitute_config. The threshold file is located via the - threshold_json field in the config (relative to the config file's directory) - or auto-discovered as the sole *threshold.json sibling. - - Cell keys in the threshold file must match MegatronVariantConfig.cell_key() - output exactly — MBS=,GBS=,PRECISION=. A load-time - validator checks that every sweep cell has a threshold entry and no key is - orphaned. - """ + """Load and validate a Megatron training variant config + its threshold file.""" raw, thresholds = substitute_config(config_path, cluster_dict) - # When checkpoint testing is disabled, checkpoint_dir and its shared-FS - # volume mount are unused — exempt both from the check so - # operators can use the template as-is without filling in checkpoint paths. if not raw.get("checkpoint", {}).get("enforce", False): raw.get("checkpoint", {}).pop("checkpoint_dir", None) try: diff --git a/cvs/lib/training/torchtitan/utils/training_config_loader.py b/cvs/lib/training/torchtitan/utils/training_config_loader.py index ffdcea8bf..374a4020a 100644 --- a/cvs/lib/training/torchtitan/utils/training_config_loader.py +++ b/cvs/lib/training/torchtitan/utils/training_config_loader.py @@ -2,201 +2,18 @@ Copyright 2025 Advanced Micro Devices, Inc. All rights reserved. -Training-specific config schema for TorchTitan suites (single-node and distributed). +Load and validate TorchTitan training variant configs from +``cvs/input/config_file/training/torchtitan/``. -The framework-agnostic machinery (ContainerSpec, RuntimeSpec, placeholder -substitution, threshold file discovery) lives in `cvs.lib.utils.config_loader`. -This module holds the training half: TorchTitanSweepCombo, TorchTitanSweep, -TorchTitanVariantConfig, and load_training_variant. - -Thresholds live in a sibling *threshold.json file (not inline in result_dict). -The threshold file is discovered via the `threshold_json` field in the config or -auto-discovered as the sole *threshold.json sibling. Cell keys in the threshold -file must match the combination keys in sweep.combinations exactly. - -enforce_thresholds gates whether threshold specs are asserted in test_metric. - -Both torchtitan_single and torchtitan_distributed are covered by TorchTitanVariantConfig -via the framework field, which is a validated schema tag / config discriminator. +Pydantic models live in ``cvs.schema.config_file.training.torchtitan.variant``. ''' -from __future__ import annotations - -import warnings -from collections import Counter -from typing import Any, Dict, List - -from pydantic import Field, model_validator -from typing_extensions import Literal - -from cvs.lib.utils.config_loader import ( - ContainerSpec, - _Forbid, - substitute_config, -) - - -# ---------- pydantic models (training) ---------- - - -class TorchTitanSweepCombo(_Forbid): - name: str - micro_batch_size: str - global_batch_size: str - precision: str = "" - - -def validate_sweep_selector(combo_keys, run_refs): - """The sweep-selector rule: combination keys unique, every run references one. - - Single home for this check, shared by the typed TorchTitanSweep validator - (load time) and pytest_generate_tests (collection time, which reads raw - JSON before the loader runs) so the two can never drift. - - Without it a typo'd run key is a silently-dropped cell — the sweep runs - a different matrix than the config reads. - """ - counts = Counter(combo_keys) - dupes = sorted(k for k, count in counts.items() if count > 1) - if dupes: - raise ValueError(f"duplicate sweep.combinations keys: {dupes}") - known = set(counts) - unknown = sorted(r for r in run_refs if r not in known) - if unknown: - raise ValueError(f"sweep.runs references unknown combinations: {unknown} (known: {sorted(known)})") - - -def validate_thresholds_cover_sweep( - *, - expected_cells, - thresholds, - enforce_thresholds: bool, - gated_metrics=None, -) -> None: - """Shared sweep/threshold coverage check for training variant configs. - - Checks every sweep cell has a threshold entry and no threshold key is - orphaned. Individual metrics within a cell are optional — absent specs - are skipped in test_metric (record-only for that metric). - """ - expected = set(expected_cells) - present = set(thresholds.keys()) - missing = sorted(expected - present) - extra = sorted(present - expected) - problems = [] - if missing: - problems.append(f"sweep cells with no threshold entry: {missing}") - if extra: - problems.append(f"threshold keys matching no sweep cell (typo?): {extra}") - gated = gated_metrics if gated_metrics is not None else set() - gated_keys = [f"training.{m}" for m in sorted(gated)] - gated_gaps = {} - for cell in sorted(expected & present): - specs = thresholds.get(cell) or {} - absent = [k for k in gated_keys if k not in specs] - if absent: - gated_gaps[cell] = absent - if gated_gaps: - problems.append(f"cells missing gated-metric specs: {gated_gaps}") - if problems: - msg = "threshold.json does not match the sweep matrix; " + "; ".join(problems) - if enforce_thresholds: - raise ValueError(msg) - warnings.warn(f"{msg} (enforce_thresholds=false -> record-only)", stacklevel=3) - - -class TorchTitanSweep(_Forbid): - combinations: Dict[str, TorchTitanSweepCombo] - runs: List[str] - - @model_validator(mode="after") - def _check_runs_reference_known_combos(self): - validate_sweep_selector( - list(self.combinations.keys()), - self.runs, - ) - return self - - -class ScalingBaseline(_Forbid): - tokens_per_sec_total: float = 0.0 - num_nodes: int = 1 - - -class LossCurveConfig(_Forbid): - sample_every: int = 10 - milestone_steps: List[int] = Field(default_factory=lambda: [100, 500, 1000, 5000]) - max_slope: float = 0.0 - enforce: bool = True - - -class ConvergenceConfig(_Forbid): - target_metric: Literal["auto", "train_loss", "eval_loss"] = "auto" - target_value: float = 0.0 - - -class CheckpointConfig(_Forbid): - enforce: bool = False # if False, test_checkpoint is skipped entirely - save_interval: int = 20 # checkpoint written every N steps - save_iters: int = 21 # save phase total; last checkpoint = floor(save/interval)*interval - resume_iters: int = 25 # load phase total (must be > last_ckpt_step) - loss_rtol: float = 0.05 # max allowed fractional loss increase across boundary - checkpoint_dir: str = "" # shared path for distributed; empty = derive from log_dir (single-node) - - -class TorchTitanVariantConfig(_Forbid): - schema_version: Literal[1] - framework: Literal["torchtitan_single", "torchtitan_distributed"] - gpu_arch: str - enforce_thresholds: bool = True - threshold_json: str = "" - scaling_baseline: ScalingBaseline = Field(default_factory=ScalingBaseline) - loss_curve: LossCurveConfig = Field(default_factory=LossCurveConfig) - convergence: ConvergenceConfig = Field(default_factory=ConvergenceConfig) - checkpoint: CheckpointConfig = Field(default_factory=CheckpointConfig) - config: Dict[str, Any] # training knobs: torchtitan_root, nccl_*, nic_type, ... - model_params: Dict[str, Any] # model knobs: model_name, precision, tp, pp, ... - container: ContainerSpec - sweep: TorchTitanSweep - thresholds: Dict[str, Dict[str, Any]] = Field(default_factory=dict) - - def cell_key(self, combo_key: str) -> str: - """Canonical threshold lookup key for a sweep combo. - - Constructs a key from the combo's micro_batch_size, global_batch_size, - and precision — must match the top-level keys in the threshold file exactly. - """ - combo = self.sweep.combinations[combo_key] - return f"MBS={combo.micro_batch_size},GBS={combo.global_batch_size},PRECISION={combo.precision}" - - def expected_cells(self) -> List[str]: - """Return the threshold cell key for every run in sweep.runs.""" - return [self.cell_key(k) for k in self.sweep.runs] - - @model_validator(mode="after") - def _check_thresholds_cover_sweep(self): - """Every sweep cell must have a threshold entry; no metric within it is - mandatory. test_metric treats an absent ``training.*`` spec as - "don't gate this metric" (skips the assertion), so a threshold.json - is free to gate only the metrics an operator cares about. - """ - validate_thresholds_cover_sweep( - expected_cells=self.expected_cells(), - thresholds=self.thresholds, - enforce_thresholds=self.enforce_thresholds, - gated_metrics=set(), - ) - return self - - -# ---------- public API (training) ---------- +from cvs.lib.utils.config_loader import substitute_config +from cvs.schema.config_file.training.torchtitan.variant import TorchTitanVariantConfig def _check_no_changeme(node, path="", _offenders=None): - """Recursively collect config fields whose value still contains ''. - - Collects all offending dotted paths so the caller can report them all at once. - """ + """Recursively collect config fields whose value still contains ''.""" if _offenders is None: _offenders = [] if isinstance(node, dict): @@ -213,18 +30,7 @@ def _check_no_changeme(node, path="", _offenders=None): def load_training_variant(config_path, cluster_dict) -> TorchTitanVariantConfig: - """Load and validate a TorchTitan training variant config + its threshold file. - - Delegates file read, placeholder substitution, and threshold file discovery - to the generic substitute_config. The threshold file is located via the - threshold_json field in the config (relative to the config file's directory) - or auto-discovered as the sole *threshold.json sibling. - - Cell keys in the threshold file must match TorchTitanVariantConfig.cell_key() - output exactly — MBS=,GBS=,PRECISION=. A load-time - validator checks that every sweep cell has a threshold entry and no key is - orphaned. - """ + """Load and validate a TorchTitan training variant config + its threshold file.""" raw, thresholds = substitute_config(config_path, cluster_dict) _check_no_changeme(raw) diff --git a/cvs/lib/unittests/test_utils_lib.py b/cvs/lib/unittests/test_utils_lib.py index 78fc42146..e1cd3a329 100644 --- a/cvs/lib/unittests/test_utils_lib.py +++ b/cvs/lib/unittests/test_utils_lib.py @@ -7,7 +7,7 @@ import cvs.lib.utils_lib as utils_lib from cvs.core.run_layout import RunLayout -from cvs.parsers.schemas import AortaBenchmarkConfigFile +from cvs.schema.config_file.aorta.benchmark import AortaBenchmarkConfigFile class TestUtilsLib(unittest.TestCase): diff --git a/cvs/lib/utils/AGENTS.md b/cvs/lib/utils/AGENTS.md index 7eeb70d64..6d5666c7e 100644 --- a/cvs/lib/utils/AGENTS.md +++ b/cvs/lib/utils/AGENTS.md @@ -219,7 +219,8 @@ Contract for new suite authors: **Must add:** - A `@model_validator(mode="after")` that performs threshold-coverage checking - (equivalent to `_check_thresholds_cover_sweep` in `inferencing_config_loader.py`). + (equivalent to `validate_thresholds_cover_sweep` in + `cvs/schema/config_file/inference/common/sweep.py`). The check must cover **two axes**: 1. **Cell coverage** — sweep cells with no threshold entry AND threshold keys that match no sweep cell (both directions; a one-way check silently skips diff --git a/cvs/lib/utils/config_loader.py b/cvs/lib/utils/config_loader.py index e56b44d5c..ea3724a6e 100644 --- a/cvs/lib/utils/config_loader.py +++ b/cvs/lib/utils/config_loader.py @@ -4,13 +4,9 @@ Framework-agnostic config machinery shared by every CVS suite. -Holds the generic half of what used to be one monolithic loader: the -container/paths/model/image schema, the 3-pass placeholder substitution, the -`enforce_thresholds` gate carried on `BaseVariantConfig`, and the -`substitute_config` helper that reads a variant `config.json` + sibling -`*threshold.json` and resolves placeholders. A per-framework module subclasses -`BaseVariantConfig` and adds its own `Params`/`Sweep`/`cell_key` (see -`cvs.lib.inference.utils.vllm_config_loader` for the vllm flavour). +Holds placeholder substitution, threshold file discovery, and +``substitute_config`` — the I/O half of variant loading. Pydantic models +(``Paths``, ``BaseVariantConfig``, etc.) live in ``cvs.schema.common.base``. 3-pass placeholder substitution: 1. cluster placeholders (`{user-id}`) anywhere @@ -37,89 +33,26 @@ import re import warnings from pathlib import Path -from typing import Any, Dict -from pydantic import BaseModel, ConfigDict, Field, model_validator -from typing_extensions import Literal - - -# ---------- pydantic models (generic) ---------- - - -class _Forbid(BaseModel): - model_config = ConfigDict(extra="forbid") - - -class _Allow(BaseModel): - model_config = ConfigDict(extra="allow") - - -class Paths(_Forbid): - shared_fs: str - models_dir: str - log_dir: str - hf_token_file: str - # Host-user-namespaced scratch (jaxmaxtext launchers/yml). Optional so - # inference configs that omit it still load; jaxmaxtext configs set it to - # /tmp/{user-id}/jaxmaxtext so container-root /tmp/root is never used. - temp_dir: str = "" - - -class ModelSpec(_Forbid): - id: str - remote: Literal[0, 1] - precision: str = "" - - -class RuntimeSpec(_Allow): - name: str - args: Dict[str, Any] = Field(default_factory=dict) - - -class ContainerSpec(_Forbid): - lifetime: Literal["no_launch", "per_run", "persistent"] = "per_run" - name: str - image: str - runtime: RuntimeSpec - - -class BaseVariantConfig(_Forbid): - """The framework-agnostic skeleton of a variant config. - - Carries the fields every suite shares (schema/paths/model/image/container/ - thresholds + the enforce gate) and the remote-not-implemented guard. - Per-framework subclasses add their own `framework`/`Params`/`Sweep` and the - `cell_key`/coverage-check pair that depend on them. - """ - - schema_version: Literal[1] - # When false, the threshold-coverage gate warns instead of raising and the - # test records metrics without asserting pass/fail (record-only). Use for - # un-calibrated shapes (e.g. a throughput characterization whose published - # numbers are curves, not tabulated values). Default true keeps the gate - # strict for calibrated configs -- no regression to the remediation work. - enforce_thresholds: bool = True - threshold_json: str = "" - paths: Paths - model: ModelSpec - # The container image is declared once, on container.image (ContainerSpec). - # There is no separate top-level image block. - container: ContainerSpec - thresholds: Dict[str, Dict[str, Any]] = Field(default_factory=dict) - - # pydantic runs @model_validator(mode="after") hooks in definition order, - # parent-class hooks before subclass hooks. This remote check is intentionally - # the first to run: an unimplemented remote config fails fast - # (NotImplementedError) before any subclass's threshold-coverage check runs, - # which is meaningless for a config we are going to reject anyway. - @model_validator(mode="after") - def _check_remote_not_implemented(self): - if self.model.remote == 1: - raise NotImplementedError( - "model.remote=1 (remote model download) is not implemented in the PoC. " - "Port from cvs-dtni-v1/resource_resolver.py before enabling." - ) - return self +from cvs.schema.base import _Allow, _Forbid +from cvs.schema.common.base import ( + BaseVariantConfig, + ContainerSpec, + ModelSpec, + Paths, + RuntimeSpec, +) + +__all__ = [ + "BaseVariantConfig", + "ContainerSpec", + "ModelSpec", + "Paths", + "RuntimeSpec", + "_Allow", + "_Forbid", + "substitute_config", +] # ---------- placeholder substitution ---------- @@ -191,7 +124,7 @@ def substitute_config(config_path, cluster_dict): and the parsed, comment-stripped threshold dict. Threshold discovery supports both layouts: - - ``threshold_json`` in the config (literal path; vllm_single style), or + - ``threshold_json`` in the config (literal path), or - a sole ``*threshold.json`` sibling next to the config (atom style). This is the framework-neutral body of the old `load_variant`: file read + diff --git a/cvs/parsers/__init__.py b/cvs/parsers/__init__.py index de96b70c3..706ef4cca 100644 --- a/cvs/parsers/__init__.py +++ b/cvs/parsers/__init__.py @@ -5,7 +5,9 @@ - Transforming raw benchmark outputs into structured data - Validating results against Pydantic schemas - Aggregating metrics across runs/ranks -- Validating configuration files (fail fast) + +Configuration file schemas live under ``cvs/schema/`` (mirroring ``cvs/input/``); use +``cvs.schema.validate.validate_config_file`` to validate configs before running benchmarks. Parsers should NOT: - Execute benchmarks @@ -14,46 +16,20 @@ """ from cvs.parsers.schemas import ( - # Result schemas - AortaTraceMetrics, AortaBenchmarkResult, + AortaTraceMetrics, ParseResult, ParseStatus, - # Config file schemas - ClusterConfigFile, - ClusterNodeConfig, - AortaBenchmarkConfigFile, - AortaDockerConfigFile, - AortaRcclConfigFile, - AortaEnvironmentConfigFile, - AortaExpectedResultsConfigFile, - AortaAnalysisConfigFile, - # Validation helper - validate_config_file, ) -# Parser implementations from cvs.parsers.aorta_report import AortaReportParser from cvs.parsers.tracelens import TraceLensParser __all__ = [ - # Result schemas "AortaTraceMetrics", "AortaBenchmarkResult", "ParseResult", "ParseStatus", - # Config file schemas - "ClusterConfigFile", - "ClusterNodeConfig", - "AortaBenchmarkConfigFile", - "AortaDockerConfigFile", - "AortaRcclConfigFile", - "AortaEnvironmentConfigFile", - "AortaExpectedResultsConfigFile", - "AortaAnalysisConfigFile", - # Validation helper - "validate_config_file", - # Parser implementations "AortaReportParser", "TraceLensParser", ] diff --git a/cvs/parsers/schemas.py b/cvs/parsers/schemas.py index ec9ece90e..df9c6ca77 100644 --- a/cvs/parsers/schemas.py +++ b/cvs/parsers/schemas.py @@ -1,26 +1,19 @@ """ -Pydantic schemas for ALL benchmark results AND configuration files. +Pydantic schemas for benchmark result parsing. -This is the single source of truth for: -- Result data structures (parsed benchmark output) -- Configuration file schemas (validated before running benchmarks) - -All parsers produce instances of these models. -Config validation happens early to fail fast with clear errors. +Configuration file schemas live under ``cvs/schema/`` (mirroring ``cvs/input/``); use +``cvs.schema.validate.validate_config_file`` to load and validate configs. Copyright 2025 Advanced Micro Devices, Inc. All rights reserved. """ -from copy import deepcopy from dataclasses import dataclass, field from enum import Enum -from pathlib import Path -from typing import Any, Dict, Generic, List, Optional, TypeVar, Union import math -import warnings +from typing import Any, Dict, Generic, List, Optional, TypeVar -from pydantic import BaseModel, Field, field_validator, model_validator, ConfigDict +from pydantic import BaseModel, ConfigDict, Field, field_validator # ============================================================================= @@ -218,1534 +211,3 @@ def from_rank_metrics( per_rank_metrics=metrics, **kwargs, ) - - -# ============================================================================= -# RCCL Schemas (for future use - mirrors existing models/rccl.py patterns) -# ============================================================================= - -# Note: RCCL schemas already exist in models/rccl.py -# When porting RCCL tests to this architecture, we can either: -# 1. Move those schemas here -# 2. Re-export them from here -# 3. Keep them separate and import as needed - - -# ============================================================================= -# Configuration File Schemas (Input Validation - Fail Fast) -# ============================================================================= - - -class ClusterNodeConfig(BaseModel): - """Schema for a single node entry in cluster.json node_dict.""" - - model_config = ConfigDict(extra="allow") # Allow extra fields like bmc_ip, rack_id - - vpc_ip: str = Field(description="VPC IP or hostname for inter-node communication") - bmc_ip: Optional[str] = Field(default=None, description="BMC IP for out-of-band management") - - -class HeadNodeConfig(BaseModel): - """Schema for head_node_dict in cluster.json.""" - - model_config = ConfigDict(extra="allow") - - mgmt_ip: str = Field(description="Management IP of head node") - - -class RackConfig(BaseModel): - """ - Schema for a single rack entry inside the 'racks' block of cluster.json. - - A rack groups compute trays (referenced via node_dict rack_id) and the - switch trays physically associated with that rack. - """ - - model_config = ConfigDict(extra="allow") - - platform: Optional[str] = Field(default=None, description="ARC platform name, e.g. 'HeliosP' or 'HeliosR'") - arc_controller: Optional[str] = Field( - default=None, - description="IP of the ARC controller node. Defaults to first sorted node_dict entry with matching rack_id.", - ) - switch_trays: List[str] = Field( - default_factory=list, - description="IPs of switch trays in this rack", - ) - rmc: Optional[str] = Field(default=None, description="IP of the Rack Management Controller") - - -class RacksBlock(BaseModel): - """ - Schema for the top-level 'racks' field in cluster.json. - - Holds optional global switch credentials and one RackConfig entry per rack - (keyed by rack ID, e.g. 'rack-01'). Extra keys (rack IDs) are accepted via - extra='allow' and retrieved via get_racks(). - - Switch credentials are fleet-wide (homogeneous across all racks). Per-rack - overrides are not supported in the current exec path; add them to RackConfig - when that need arises. - """ - - model_config = ConfigDict(extra="allow") - - switch_ssh_user: Optional[str] = Field( - default=None, - description="SSH username for all switch trays in every rack.", - ) - switch_ssh_password: Optional[str] = Field( - default=None, - description="SSH password for all switch trays. Ignored when switch_ssh_key_file is set.", - ) - switch_ssh_key_file: Optional[str] = Field( - default=None, - description="Path to SSH private key for all switch trays. Takes priority over switch_ssh_password when set.", - ) - - def get_racks(self) -> Dict[str, RackConfig]: - """Return only the rack entries, excluding credential fields.""" - skip = {'switch_ssh_user', 'switch_ssh_password', 'switch_ssh_key_file'} - result = {} - for key, value in (self.__pydantic_extra__ or {}).items(): - if key not in skip and isinstance(value, dict): - result[key] = RackConfig(**value) - return result - - -class ClusterConfigFile(BaseModel): - """ - Schema for cluster.json configuration file. - - Validates the cluster configuration before running benchmarks. - Fails fast with clear error messages if required fields are missing. - """ - - model_config = ConfigDict(extra="allow") - - username: str = Field(description="SSH username for cluster nodes") - priv_key_file: Optional[str] = Field(default=None, description="Path to SSH private key") - password: Optional[str] = Field(default=None, description="SSH password (if not using key)") - - node_dict: Dict[str, ClusterNodeConfig] = Field( - description="Dictionary mapping node hostname/IP to node configuration" - ) - head_node_dict: Optional[HeadNodeConfig] = Field(default=None, description="Head node configuration") - - racks: Optional[RacksBlock] = Field( - default=None, - description=( - "Rack topology block. Contains optional global switch credentials and one entry per rack " - "(keyed by rack ID) listing switch_trays and platform." - ), - ) - rack_groups: Optional[RacksBlock] = Field( - default=None, - description="Deprecated alias for 'racks'. Use 'racks' instead.", - ) - - # Optional fields that may be present - home_mount_dir_name: Optional[str] = Field(default="home") - node_dir_name: Optional[str] = Field(default="root") - - @model_validator(mode='after') - def validate_auth_method(self): - """Ensure at least one authentication method is provided.""" - if not self.priv_key_file and not self.password: - raise ValueError("Authentication required: provide either 'priv_key_file' or 'password' in cluster config") - return self - - @model_validator(mode='after') - def validate_nodes_exist(self): - """Ensure at least one node is configured.""" - if not self.node_dict: - raise ValueError("No nodes configured in 'node_dict' - at least one node is required") - return self - - @model_validator(mode='after') - def warn_rack_groups_deprecated(self): - """Emit a deprecation warning when the old 'rack_groups' key is used.""" - import warnings - - if self.rack_groups is not None and self.racks is None: - warnings.warn( - "'rack_groups' in cluster.json is deprecated. Rename it to 'racks'.", - DeprecationWarning, - stacklevel=2, - ) - return self - - def get_racks_block(self) -> Optional[RacksBlock]: - """Return the active racks block, preferring 'racks' over the deprecated 'rack_groups'.""" - return self.racks if self.racks is not None else self.rack_groups - - @field_validator('username') - @classmethod - def validate_username_not_placeholder(cls, v: str) -> str: - """Check that username is not still a placeholder.""" - if '' in v.lower(): - raise ValueError( - "Username contains placeholder ''. Please set a valid username in cluster config." - ) - return v - - -class AortaDockerConfigFile(BaseModel): - """Schema for docker section in aorta_benchmark.yaml.""" - - model_config = ConfigDict(extra="forbid") # Catch typos - - image: str = Field( - default="jeffdaily/pytorch:torchrec-dlrm-complete", description="Docker image for Aorta container" - ) - container_name: str = Field(default="aorta-benchmark", description="Name for the Docker container") - shm_size: str = Field(default="17G", description="Shared memory size") - network_mode: str = Field(default="host", description="Docker network mode") - privileged: bool = Field(default=True, description="Run container in privileged mode") - - -class AortaRcclConfigFile(BaseModel): - """Schema for rccl section in aorta_benchmark.yaml.""" - - model_config = ConfigDict(extra="forbid") - - clone_url: str = Field( - default="https://github.com/ROCmSoftwarePlatform/rccl.git", description="RCCL git repository URL" - ) - branch: str = Field(default="develop", description="RCCL branch to build") - build_path: str = Field(default="/mnt/rccl", description="Path inside container for RCCL build") - - -class AortaEnvironmentConfigFile(BaseModel): - """Schema for environment section in aorta_benchmark.yaml.""" - - model_config = ConfigDict(extra="allow") # Allow custom env vars - - NCCL_MAX_NCHANNELS: int = Field(default=112, ge=1, le=256, description="Maximum NCCL channels") - NCCL_MAX_P2P_NCHANNELS: int = Field(default=112, ge=1, le=256, description="Maximum NCCL P2P channels") - NCCL_DEBUG: str = Field(default="VERSION", description="NCCL debug level") - TORCH_NCCL_HIGH_PRIORITY: int = Field(default=1, ge=0, le=1, description="Enable high priority NCCL streams") - OMP_NUM_THREADS: int = Field(default=1, ge=1, description="OpenMP thread count") - RCCL_MSCCL_ENABLE: int = Field(default=0, ge=0, le=1, description="Enable MSCCL") - - -class AortaExpectedResultsConfigFile(BaseModel): - """Schema for expected_results section in aorta_benchmark.yaml.""" - - model_config = ConfigDict(extra="allow") # Allow custom thresholds - - max_avg_iteration_ms: Optional[float] = Field( - default=None, ge=0, description="Maximum acceptable average iteration time in ms" - ) - min_compute_ratio: Optional[float] = Field(default=None, ge=0, le=1, description="Minimum acceptable compute ratio") - min_overlap_ratio: Optional[float] = Field( - default=None, ge=0, le=1, description="Minimum acceptable compute-comm overlap ratio" - ) - max_time_variance_ratio: Optional[float] = Field( - default=None, ge=0, description="Maximum acceptable iteration time variance" - ) - - -class AortaAnalysisConfigFile(BaseModel): - """Schema for analysis section in aorta_benchmark.yaml.""" - - model_config = ConfigDict(extra="forbid") - - enable_tracelens: bool = Field(default=True, description="Run Aorta's TraceLens analysis after benchmark") - enable_gemm_analysis: bool = Field(default=False, description="Run Aorta's GEMM analysis (for sweep experiments)") - tracelens_script: str = Field( - default="scripts/tracelens_single_config/run_tracelens_single_config.sh", - description="Path to TraceLens analysis script relative to aorta_path", - ) - gemm_script: str = Field( - default="scripts/gemm_analysis/run_tracelens_analysis.sh", - description="Path to GEMM analysis script relative to aorta_path", - ) - skip_if_exists: bool = Field( - default=False, description="Skip analysis if tracelens_analysis directory already exists" - ) - - -class AortaBenchmarkConfigFile(BaseModel): - """ - Schema for the entire aorta_benchmark.yaml configuration file. - - Validates structure and provides sensible defaults. - Fails fast with clear error messages if configuration is invalid. - - For ``test_aorta``, load YAML, apply ``resolve_test_config_placeholders`` with the resolved - cluster dict (same as other CVS suites), then ``model_validate``. Standalone tools may validate - raw YAML without placeholder resolution if paths are already absolute. - """ - - model_config = ConfigDict(extra="forbid") # Catch typos in top-level keys - - # Path to Aorta repository on host (will be bind-mounted). If missing and aorta_auto_clone is true, it is cloned. - aorta_path: str = Field(description="Path to Aorta repository on host (will be bind-mounted)") - - # Optional: clone Aorta repo when aorta_path does not exist - aorta_auto_clone: bool = Field( - default=False, description="If true and aorta_path missing, clone from aorta_clone_url" - ) - aorta_clone_url: Optional[str] = Field(default=None, description="Git URL to clone when aorta_auto_clone is true") - - # Container settings - container_mount_path: str = Field(default="/mnt", description="Mount point inside container for aorta_path") - - # Aorta config - base_config: str = Field(default="config/distributed.yaml", description="Aorta config file relative to aorta_path") - - # Nested configuration sections - docker: AortaDockerConfigFile = Field( - default_factory=AortaDockerConfigFile, description="Docker container configuration" - ) - rccl: AortaRcclConfigFile = Field(default_factory=AortaRcclConfigFile, description="RCCL build configuration") - environment: AortaEnvironmentConfigFile = Field( - default_factory=AortaEnvironmentConfigFile, description="Environment variables for RCCL/NCCL" - ) - - # Training overrides - training_overrides: Dict[str, Any] = Field( - default_factory=dict, description="Overrides passed to Aorta via --override flag" - ) - - # Scripts - build_script: str = Field( - default="scripts/build_rccl.sh", description="RCCL build script relative to container mount" - ) - experiment_script: str = Field( - default="scripts/rccl_exp.sh", description="Experiment script relative to container mount" - ) - - # Hardware - gpus_per_node: int = Field(default=8, ge=1, description="Number of GPUs per node") - - # Execution settings - timeout_seconds: int = Field(default=10800, ge=60, description="Benchmark timeout in seconds") - skip_rccl_build: bool = Field(default=False, description="Skip RCCL build if already built") - - # Validation thresholds - expected_results: AortaExpectedResultsConfigFile = Field( - default_factory=AortaExpectedResultsConfigFile, description="Expected results for validation" - ) - - # Analysis configuration (use Aorta's built-in analysis scripts) - analysis: AortaAnalysisConfigFile = Field( - default_factory=AortaAnalysisConfigFile, description="Post-benchmark analysis configuration" - ) - - @field_validator('aorta_path') - @classmethod - def validate_aorta_path_not_placeholder(cls, v: str) -> str: - """Check that aorta_path is not a placeholder.""" - if '' in v.lower(): - raise ValueError( - "aorta_path contains placeholder ''. Please set the actual path to your Aorta installation." - ) - return v - - def validate_paths_exist(self) -> List[str]: - """ - Validate that referenced paths exist on the filesystem. - - Call this after loading config to check paths. - Returns list of error messages (empty if all valid). - """ - errors = [] - - aorta = Path(self.aorta_path) - if not aorta.exists(): - if self.aorta_auto_clone and self.aorta_clone_url: - # Runner will clone in setup(); skip path checks here - return errors - errors.append(f"aorta_path does not exist: {self.aorta_path}") - else: - # Check internal paths - base_cfg = aorta / self.base_config - if not base_cfg.exists(): - errors.append(f"base_config does not exist: {base_cfg}") - - build_script = aorta / self.build_script - if not build_script.exists(): - errors.append(f"build_script does not exist: {build_script}") - - exp_script = aorta / self.experiment_script - if not exp_script.exists(): - errors.append(f"experiment_script does not exist: {exp_script}") - - # Check analysis scripts if enabled - if self.analysis.enable_tracelens: - tracelens_script = aorta / self.analysis.tracelens_script - if not tracelens_script.exists(): - errors.append(f"tracelens_script does not exist: {tracelens_script}") - - if self.analysis.enable_gemm_analysis: - gemm_script = aorta / self.analysis.gemm_script - if not gemm_script.exists(): - errors.append(f"gemm_script does not exist: {gemm_script}") - - return errors - - -# ============================================================================= -# PyTorch XDit (WAN/Flux) Schemas -# ============================================================================= - - -class PytorchXditDistributedNcclExamples(BaseModel): - """Documentation-only example values shipped beside ```` in sample JSON.""" - - model_config = ConfigDict(extra="forbid", populate_by_name=True) - - example_nccl_ib_hca: Optional[str] = Field( - default=None, - alias="_example_nccl_ib_hca", - description="Documentation only: example nccl_ib_hca value for this cluster", - ) - example_nccl_socket_ifname: Optional[str] = Field( - default=None, - alias="_example_nccl_socket_ifname", - description="Documentation only: example nccl_socket_ifname value", - ) - example_gloo_socket_ifname: Optional[str] = Field( - default=None, - alias="_example_gloo_socket_ifname", - description="Documentation only: example gloo_socket_ifname value", - ) - - -class PytorchXditContainerConfig(BaseModel): - """Schema for container_config section in pytorch-xdit configs.""" - - model_config = ConfigDict(extra="allow") - - device_list: List[str] = Field( - default=["/dev/dri", "/dev/kfd"], description="List of device paths to mount in container" - ) - volume_dict: Dict[str, str] = Field(default_factory=dict, description="Host:container volume mount mappings") - env_dict: Dict[str, str] = Field(default_factory=dict, description="Environment variables for container") - - -class PytorchXditExpectedResults(BaseModel): - """Schema for expected_results in pytorch-xdit WAN benchmark params.""" - - model_config = ConfigDict(extra="forbid") - - max_avg_total_time_s: Optional[float] = Field( - default=None, - gt=0, - description="Maximum acceptable average total_time in seconds (native/packaged Wan)", - ) - max_avg_pipe_time_s: Optional[float] = Field( - default=None, - gt=0, - description="Maximum acceptable average pipe_time in seconds (xFuser Wan I2V)", - ) - - @model_validator(mode="after") - def validate_threshold_present(self) -> "PytorchXditExpectedResults": - if self.max_avg_total_time_s is None and self.max_avg_pipe_time_s is None: - raise ValueError("expected_results entry must include max_avg_total_time_s and/or max_avg_pipe_time_s") - return self - - -class PytorchXditWan22Benchmarks(BaseModel): - """Schema for wan22_i2v_a14b benchmark parameters.""" - - model_config = ConfigDict(extra="forbid") - - prompt: str = Field(description="Text prompt for image-to-video generation") - model_format: Optional[str] = Field( - default=None, - description=( - "WAN checkpoint layout override: native (Wan2.2-I2V-A14B) or diffusers " - "(Wan2.2-I2V-A14B-Diffusers). Auto-inferred from model_repo or model_index.json when omitted." - ), - ) - size: str = Field(default="720*1280", pattern=r"^\d+\*\d+$", description="Video resolution (format: height*width)") - frame_num: int = Field(default=81, ge=1, description="Number of frames to generate") - num_benchmark_steps: int = Field(default=5, ge=1, description="Number of benchmark iterations to run") - num_inference_steps: Optional[int] = Field( - default=None, - ge=1, - description="Diffusers denoising steps for /app/Wan/run.py (defaults to 40 when omitted).", - ) - seed: Optional[int] = Field( - default=None, - description="Random seed for Diffusers Wan runs (defaults to 42 when omitted).", - ) - wan_diffusers_run_script: Optional[str] = Field( - default=None, - description=( - "In-container Diffusers Wan launcher script. Defaults to /app/Wan/run.py " - "(shipped in amdsiloai/pytorch-xdit and rocm/pytorch-xdit benchmark images)." - ), - ) - wan_diffusers_i2v_image: Optional[str] = Field( - default=None, - description=( - "In-container input image for Diffusers Wan I2V. Omit or set 'auto' to generate a " - "placeholder JPEG in-container; otherwise bind-mount the host file via volume_dict." - ), - ) - wan_xfuser_auto_input_image: Optional[bool] = Field( - default=None, - description=( - "Generate a synthetic I2V input JPEG inside the container for xFuser runs. " - "Defaults to true when no host image is bind-mounted." - ), - ) - wan_xfuser_install_video_deps: bool = Field( - default=True, - description="Run pip install imageio imageio-ffmpeg before xFuser video export.", - ) - wan_diffusers_launcher: Optional[str] = Field( - default=None, - description=( - "Diffusers Wan launcher: packaged (/app/Wan/run.py in pytorch-xdit images) or " - "xfuser_example (mount cvs .../scripts/wan_i2v_example.py for ufb-private)." - ), - ) - warmup_steps: Optional[int] = Field( - default=None, - ge=0, - description="Warmup iterations for xfuser_example launcher (defaults to 1).", - ) - wan_xfuser_output_type: Optional[str] = Field( - default=None, - description="xFuser output_type for xfuser_example (defaults to pil for ufb-private video export).", - ) - wan_diffusers_save_video_path: Optional[str] = Field( - default=None, - description="In-container MP4 path for xfuser_example (default /outputs/results/video_i2v.mp4).", - ) - wan_diffusers_timing_json_path: Optional[str] = Field( - default=None, - description="In-container timing JSON path for xfuser_example (default results/timing.json).", - ) - wan_diffusers_video_fps: Optional[int] = Field( - default=None, - ge=1, - description="FPS passed to export_to_video for xfuser_example (defaults to 16).", - ) - require_video_artifact: bool = Field( - default=True, - description="Require video.mp4 under the output dir when parsing results.", - ) - compile: bool = Field(default=True, description="Whether to use torch.compile for optimization") - torchrun_nproc: int = Field(default=8, ge=1, description="Number of processes for torchrun (usually num GPUs)") - ulysses_size: int = Field(default=8, ge=1, description="Ulysses parallelism degree") - ring_size: int = Field(default=1, ge=1, description="Ring parallelism degree") - expected_results: Dict[str, PytorchXditExpectedResults] = Field( - description="Expected results by GPU type (auto, mi300x, mi355, etc.)" - ) - - @field_validator('expected_results') - @classmethod - def validate_has_auto_or_specific( - cls, v: Dict[str, PytorchXditExpectedResults] - ) -> Dict[str, PytorchXditExpectedResults]: - """Ensure either 'auto' or a specific GPU type is present.""" - if not v: - raise ValueError("expected_results must contain at least one GPU type threshold") - if 'auto' not in v and not any(k in v for k in ['mi300x', 'mi325', 'mi350', 'mi355']): - raise ValueError("expected_results must contain either 'auto' or a specific GPU type (mi300x, mi325, etc.)") - return v - - -class PytorchXditFluxExpectedResults(BaseModel): - """Schema for expected_results in Flux benchmark params.""" - - model_config = ConfigDict(extra="forbid") - - max_avg_pipe_time_s: float = Field(gt=0, description="Maximum acceptable average pipe_time in seconds") - - -class PytorchXditFlux1DevBenchmarks(BaseModel): - """Schema for flux1_dev_t2i benchmark parameters.""" - - model_config = ConfigDict(extra="forbid") - - prompt: str = Field(description="Text prompt for text-to-image generation") - seed: int = Field(default=42, description="Random seed for reproducibility") - num_inference_steps: int = Field(default=25, ge=1, description="Number of denoising steps") - max_sequence_length: int = Field(default=256, ge=1, description="Maximum sequence length for text encoder") - model_type: Optional[str] = Field( - default=None, - description=( - "FLUX model family override (flux2 for FLUX.2-dev, flux_kontext for FLUX.1-Kontext). " - "Auto-inferred from model_repo or model_index.json when omitted." - ), - ) - guidance_scale: Optional[float] = Field( - default=None, - gt=0, - description=( - "Classifier-free guidance scale for run_usp.py. Defaults to 4.0 for FLUX.2-dev " - "and 2.5 for FLUX.1-Kontext when omitted; not passed for FLUX.1-dev." - ), - ) - no_use_resolution_binning: bool = Field(default=True, description="Disable resolution binning") - warmup_steps: int = Field(default=1, ge=0, description="Number of warmup steps before benchmarking") - warmup_calls: int = Field(default=5, ge=0, description="Number of warmup calls") - num_repetitions: int = Field(default=25, ge=1, description="Number of benchmark repetitions") - height: int = Field(default=1024, ge=1, description="Output image height in pixels") - width: int = Field(default=1024, ge=1, description="Output image width in pixels") - ulysses_degree: int = Field(default=8, ge=1, description="Ulysses parallelism degree") - ring_degree: int = Field(default=1, ge=1, description="Ring parallelism degree") - pipefusion_parallel_degree: int = Field( - default=1, ge=1, description="PipeFusion pipeline-parallel degree (multi-node)" - ) - tensor_parallel_degree: int = Field(default=1, ge=1, description="Tensor-parallel degree (1 = disabled)") - data_parallel_degree: int = Field(default=1, ge=1, description="Data-parallel degree (1 = disabled)") - use_torch_compile: bool = Field(default=True, description="Whether to use torch.compile for optimization") - torchrun_nproc: int = Field(default=8, ge=1, description="Number of processes for torchrun (usually num GPUs)") - expected_results: Dict[str, PytorchXditFluxExpectedResults] = Field( - description="Expected results by GPU type (auto, mi300x, mi355, etc.)" - ) - - @field_validator('expected_results') - @classmethod - def validate_has_auto_or_specific( - cls, v: Dict[str, PytorchXditFluxExpectedResults] - ) -> Dict[str, PytorchXditFluxExpectedResults]: - """Ensure either 'auto' or a specific GPU type is present.""" - if not v: - raise ValueError("expected_results must contain at least one GPU type threshold") - if 'auto' not in v and not any(k in v for k in ['mi300x', 'mi325', 'mi350', 'mi355']): - raise ValueError("expected_results must contain either 'auto' or a specific GPU type (mi300x, mi325, etc.)") - return v - - -class PytorchXditBenchmarkParams(BaseModel): - """Schema for benchmark_params section in pytorch-xdit configs.""" - - model_config = ConfigDict(extra="forbid") - - wan22_i2v_a14b: Optional[PytorchXditWan22Benchmarks] = Field( - default=None, description="WAN 2.2 image-to-video A14B benchmark parameters" - ) - flux1_dev_t2i: Optional[PytorchXditFlux1DevBenchmarks] = Field( - default=None, description="FLUX.1-dev text-to-image benchmark parameters" - ) - - -class PytorchXditWanConfigFile(BaseModel): - """ - Schema for PyTorch XDit WAN microbenchmark configuration file. - - Validates WAN inference config structure and provides fail-fast validation. - - Usage: - with open("mi300x_wan22_i2v_a14b.json") as f: - raw = json.load(f) - config = PytorchXditWanConfigFile.model_validate(raw) - """ - - model_config = ConfigDict(extra="forbid") - - config: 'PytorchXditWanConfig' = Field(description="Main configuration section") - benchmark_params: PytorchXditBenchmarkParams = Field(description="Benchmark parameters section") - - @model_validator(mode='after') - def validate_benchmark_present(self): - """Ensure at least one benchmark is configured.""" - if not self.benchmark_params.wan22_i2v_a14b: - raise ValueError("No benchmarks configured in 'benchmark_params' - at least wan22_i2v_a14b is required") - return self - - @model_validator(mode='after') - def validate_distributed_parallelism(self): - """When nnodes >= 2, ensure xDiT parallel degrees match nnodes × torchrun_nproc.""" - wan = self.benchmark_params.wan22_i2v_a14b - nnodes = self.config.nnodes - if not wan or not nnodes or nnodes < 2: - return self - - world_size = nnodes * wan.torchrun_nproc - product = wan.ulysses_size * wan.ring_size - if product != world_size: - raise ValueError( - f"Parallel degree product {product} != world_size {world_size} " - f"(nnodes={nnodes} × torchrun_nproc={wan.torchrun_nproc}). " - f"Adjust ulysses_size and ring_size." - ) - return self - - -class PytorchXditWanConfig(PytorchXditDistributedNcclExamples): - """Schema for config section in pytorch-xdit WAN configs.""" - - model_config = ConfigDict(extra="forbid", populate_by_name=True) - - container_image: str = Field( - default="amdsiloai/pytorch-xdit:v25.11.2", description="Docker image for pytorch-xdit container" - ) - container_name: str = Field(default="wan22-benchmark", description="Name for the Docker container") - hf_token_file: str = Field( - default="", - description=( - "Optional path to Hugging Face token file. " - "Not required when using a pre-staged local model path (recommended) or pre-cached HF snapshots (offline)." - ), - ) - hf_home: str = Field(description="Host directory for Hugging Face cache (mounted to /hf_home)") - output_base_dir: str = Field(description="Host base directory for benchmark outputs") - model_repo: str = Field( - default="Wan-AI/Wan2.2-I2V-A14B", - description=( - "Model identifier. Prefer an explicit local filesystem path (e.g., /models/Wan-AI/Wan2.2-I2V-A14B) " - "to avoid any runtime downloads. For backward compatibility, a Hugging Face repo id may be used only if " - "the snapshot is already cached under hf_home." - ), - ) - model_rev: str = Field( - default="206a9ee1b7bfaaf8f7e4d81335650533490646a3", - description="Model revision (commit hash). Ignored if model_repo is an explicit local filesystem path.", - ) - nnodes: Optional[int] = Field( - default=None, - ge=1, - description="Distributed node count for unified multi-node torchrun (omit for single-node / scale-out)", - ) - server_node_list: Optional[List[str]] = Field( - default=None, - description="Ordered server nodes for distributed job; defaults to all cluster nodes", - ) - master_addr: str = Field( - default="", - description="Rank-0 rendezvous address; empty means first server node at runtime", - ) - master_port: int = Field( - default=29500, - ge=1, - le=65535, - description="Rank-0 rendezvous port for distributed torchrun", - ) - nccl_ib_hca: str = Field( - default="", - description="NCCL_IB_HCA for multi-node ROCm/NCCL (e.g. rdma0,...,rdma7)", - ) - nccl_socket_ifname: str = Field( - default="", - description="NCCL_SOCKET_IFNAME for multi-node jobs", - ) - gloo_socket_ifname: str = Field( - default="", - description="GLOO_SOCKET_IFNAME for multi-node jobs", - ) - nccl_ib_gid_index: int = Field( - default=1, - ge=0, - description="NCCL_IB_GID_INDEX for IB/RoCE", - ) - nccl_debug: str = Field( - default="INFO", - description="NCCL_DEBUG level (ERROR, INFO, WARN, ...)", - ) - container_config: PytorchXditContainerConfig = Field( - default_factory=PytorchXditContainerConfig, description="Container device/volume/env configuration" - ) - - @field_validator('hf_token_file', 'hf_home', 'output_base_dir') - @classmethod - def validate_path_not_placeholder(cls, v: str, info) -> str: - """Check that paths are not still placeholders.""" - if not v: - return v - if '' in v.lower(): - raise ValueError(f"{info.field_name} contains placeholder ''. Please set a valid path in config.") - return v - - -class PytorchXditFluxConfigFile(BaseModel): - """ - Schema for PyTorch XDit Flux microbenchmark configuration file. - - Validates Flux inference config structure and provides fail-fast validation. - - Usage: - with open("mi300x_flux1_dev_t2i.json") as f: - raw = json.load(f) - config = PytorchXditFluxConfigFile.model_validate(raw) - """ - - model_config = ConfigDict(extra="forbid") - - config: 'PytorchXditFluxConfig' = Field(description="Main configuration section") - benchmark_params: PytorchXditBenchmarkParams = Field(description="Benchmark parameters section") - - @model_validator(mode='after') - def validate_benchmark_present(self): - """Ensure at least one benchmark is configured.""" - if not self.benchmark_params.flux1_dev_t2i: - raise ValueError("No benchmarks configured in 'benchmark_params' - at least flux1_dev_t2i is required") - return self - - @model_validator(mode='after') - def validate_distributed_parallelism(self): - """When nnodes >= 2, ensure xDiT parallel degrees match nnodes × torchrun_nproc.""" - flux = self.benchmark_params.flux1_dev_t2i - nnodes = self.config.nnodes - if not flux or not nnodes or nnodes < 2: - return self - - world_size = nnodes * flux.torchrun_nproc - product = ( - flux.ulysses_degree - * flux.ring_degree - * flux.pipefusion_parallel_degree - * flux.tensor_parallel_degree - * flux.data_parallel_degree - ) - if product != world_size: - raise ValueError( - f"Parallel degree product {product} != world_size {world_size} " - f"(nnodes={nnodes} × torchrun_nproc={flux.torchrun_nproc}). " - f"Adjust ulysses/ring/pipefusion/tensor_parallel/data_parallel." - ) - return self - - -class PytorchXditFluxConfig(PytorchXditDistributedNcclExamples): - """Schema for config section in pytorch-xdit Flux configs.""" - - model_config = ConfigDict(extra="forbid", populate_by_name=True) - - container_image: str = Field( - default="amdsiloai/pytorch-xdit:v25.11.2", description="Docker image for pytorch-xdit container" - ) - container_name: str = Field(default="flux-benchmark", description="Name for the Docker container") - hf_token_file: str = Field( - default="", - description=( - "Optional path to Hugging Face token file. " - "Not required when using a pre-staged local model path (recommended) or pre-cached HF snapshots (offline)." - ), - ) - hf_home: str = Field(description="Host directory for Hugging Face cache (mounted to /hf_home)") - output_base_dir: str = Field(description="Host base directory for benchmark outputs") - model_repo: str = Field( - default="black-forest-labs/FLUX.1-dev", - description=( - "Model identifier. Prefer an explicit local filesystem path (e.g., /models/black-forest-labs/FLUX.1-dev) " - "to avoid any runtime downloads. For backward compatibility, a Hugging Face repo id may be used only if " - "the snapshot is already cached under hf_home." - ), - ) - model_rev: str = Field( - default="", - description=( - "Model revision (commit hash). Empty means use any available cached snapshot under hf_home. " - "Ignored if model_repo is an explicit local filesystem path." - ), - ) - nnodes: Optional[int] = Field( - default=None, - ge=1, - description="Distributed node count for unified multi-node torchrun (omit for single-node / scale-out)", - ) - server_node_list: Optional[List[str]] = Field( - default=None, - description="Ordered server nodes for distributed job; defaults to all cluster nodes", - ) - master_addr: str = Field( - default="", - description="Rank-0 rendezvous address; empty means first server node at runtime", - ) - master_port: int = Field( - default=29500, - ge=1, - le=65535, - description="Rank-0 rendezvous port for distributed torchrun", - ) - nccl_ib_hca: str = Field( - default="", - description="NCCL_IB_HCA for multi-node ROCm/NCCL (e.g. rdma0,...,rdma7)", - ) - nccl_socket_ifname: str = Field( - default="", - description="NCCL_SOCKET_IFNAME for multi-node jobs", - ) - gloo_socket_ifname: str = Field( - default="", - description="GLOO_SOCKET_IFNAME for multi-node jobs", - ) - nccl_ib_gid_index: int = Field( - default=1, - ge=0, - description="NCCL_IB_GID_INDEX for IB/RoCE", - ) - nccl_debug: str = Field( - default="INFO", - description="NCCL_DEBUG level (ERROR, INFO, WARN, ...)", - ) - container_config: PytorchXditContainerConfig = Field( - default_factory=PytorchXditContainerConfig, description="Container device/volume/env configuration" - ) - - @field_validator('hf_token_file', 'hf_home', 'output_base_dir') - @classmethod - def validate_path_not_placeholder(cls, v: str, info) -> str: - """Check that paths are not still placeholders.""" - if not v: - return v - if '' in v.lower(): - raise ValueError(f"{info.field_name} contains placeholder ''. Please set a valid path in config.") - return v - - -# ============================================================================= -# Preflight Check Configuration Schema -# ============================================================================= - - -LEGACY_PREFLIGHT_RDMA_PATHS = { - "gid_index": "gid_index", - "rdma_interfaces": "interfaces", -} - -PREFLIGHT_METADATA_PREFIXES = ("_comment", "_example") - - -def strip_preflight_metadata(value): - """Remove documentation-only pseudo-fields before schema validation. - - Preflight JSON files conventionally carry ``_comment*`` and ``_example*`` - keys so that the files remain self-documenting. They are not runtime - options. Strip only those reserved prefixes recursively, preserving strict - rejection of every other unknown customer-facing option. - """ - if isinstance(value, dict): - return { - key: strip_preflight_metadata(item) - for key, item in value.items() - if not (isinstance(key, str) and key.startswith(PREFLIGHT_METADATA_PREFIXES)) - } - if isinstance(value, list): - return [strip_preflight_metadata(item) for item in value] - return value - - -def normalize_legacy_preflight_rdma_config(value): - """Move the two deprecated node-check RDMA keys to their canonical block. - - Returns a deep-copied configuration and one consolidated warning message, - or the original value and ``None`` when no legacy keys are present. - Conflicting legacy and canonical values fail rather than silently choosing - which RDMA inventory should be tested. - """ - if not isinstance(value, dict): - return value, None - - node_check = value.get("node_check") - if not isinstance(node_check, dict): - return value, None - - legacy_keys = [key for key in LEGACY_PREFLIGHT_RDMA_PATHS if key in node_check] - if not legacy_keys: - return value, None - - normalized = deepcopy(value) - normalized_node_check = normalized["node_check"] - connectivity_check = normalized.setdefault("connectivity_check", {}) - if not isinstance(connectivity_check, dict): - raise ValueError( - "preflight.connectivity_check must be an object when deprecated node_check RDMA options are used" - ) - rdma = connectivity_check.setdefault("rdma", {}) - if not isinstance(rdma, dict): - raise ValueError( - "preflight.connectivity_check.rdma must be an object when deprecated node_check RDMA options are used" - ) - - migrations = [] - for legacy_key in legacy_keys: - canonical_key = LEGACY_PREFLIGHT_RDMA_PATHS[legacy_key] - legacy_value = normalized_node_check.pop(legacy_key) - if canonical_key in rdma and rdma[canonical_key] != legacy_value: - raise ValueError( - f"Conflicting preflight RDMA options: preflight.node_check.{legacy_key} and " - f"preflight.connectivity_check.rdma.{canonical_key} must have the same value when both are provided" - ) - rdma.setdefault(canonical_key, legacy_value) - migrations.append(f"preflight.node_check.{legacy_key} -> preflight.connectivity_check.rdma.{canonical_key}") - - warning_message = ( - "Deprecated preflight RDMA configuration detected: " - + ", ".join(migrations) - + ". Use the preflight.connectivity_check.rdma paths; legacy paths will be removed in a future release." - ) - return normalized, warning_message - - -LEGACY_PREFLIGHT_NODE_SMOKE_SECTIONS = { - "node_smoke": "node_smoke_tier1", - "tier3_info": "node_smoke_tier3", -} - - -def _preflight_section_has_values(section: dict) -> bool: - if not isinstance(section, dict): - return False - return any(value not in (None, "") for value in section.values()) - - -def normalize_legacy_preflight_node_smoke_sections(value): - """Copy legacy Node Smoke section names to their canonical tier keys. - - Returns a deep-copied configuration and one consolidated warning message, - or the original value and ``None`` when no legacy keys need migration. - Canonical sections win when both legacy and canonical blocks are populated. - """ - if not isinstance(value, dict): - return value, None - - legacy_keys = [key for key in LEGACY_PREFLIGHT_NODE_SMOKE_SECTIONS if key in value] - if not legacy_keys: - return value, None - - normalized = deepcopy(value) - migrations = [] - for legacy_key in legacy_keys: - canonical_key = LEGACY_PREFLIGHT_NODE_SMOKE_SECTIONS[legacy_key] - legacy_block = normalized.get(legacy_key) - if not isinstance(legacy_block, dict): - continue - canonical_block = normalized.get(canonical_key) - if isinstance(canonical_block, dict) and _preflight_section_has_values(canonical_block): - continue - normalized[canonical_key] = deepcopy(legacy_block) - migrations.append(f"preflight.{legacy_key} -> preflight.{canonical_key}") - - if not migrations: - return value, None - - warning_message = ( - "Deprecated preflight Node Smoke section name(s) detected: " - + ", ".join(migrations) - + ". Prefer node_smoke_tier1 and node_smoke_tier3 in new configs." - ) - return normalized, warning_message - - -class PreflightParallelismConfig(BaseModel): - """Legacy parallelism settings for preflight checks.""" - - model_config = ConfigDict(extra="allow") - - parallel_group_size: int = Field( - default=128, - ge=2, - le=512, - description=("Legacy alias for RDMA grouping. Prefer connectivity_check.rdma.nodes_per_full_mesh_group."), - ) - - -class PreflightDebugConfig(BaseModel): - """Debug and troubleshooting settings for preflight checks.""" - - model_config = ConfigDict(extra="allow") - - scriptlet: bool = Field( - default=False, - description=( - "Enable ScriptLet debug: preserve generated scripts/logs on remote nodes. " - "For RDMA connectivity, also wraps each ibv_rc_pingpong server in strace with " - "per-test traces under /tmp/preflight/strace_server__.log (expensive at scale)." - ), - ) - - -class PreflightNodeCheckConfig(BaseModel): - """Individual node validation settings.""" - - model_config = ConfigDict(extra="forbid") - - enabled: bool = Field(default=True, description="Enable generic GPU node health and ROCm validation") - gpus_per_node: int = Field(default=4, ge=1, description="Expected AMD GPU count on each node") - expected_rocm_version: str = Field(default="6.2.0", description="Expected ROCm version across all cluster nodes") - - -class PreflightRdmaConfig(BaseModel): - """RDMA connectivity testing settings.""" - - model_config = ConfigDict(extra="allow") - - connectivity_mode: str = Field(default="basic", description="RDMA connectivity testing: basic, full_mesh, or skip") - gid_index: str = Field(default="3", description="GID index to check on all RDMA interfaces (typically 3 for RoCE)") - interfaces: List[str] = Field( - default_factory=lambda: ["rocep28s0", "rocep62s0", "rocep79s0", "rocep96s0"], - min_length=1, - description="RDMA device names checked for presence, GID consistency, and connectivity", - ) - nodes_per_full_mesh_group: int = Field( - default=128, - ge=2, - le=512, - description=( - "Number of nodes in each full-mesh partition group (2-512). " - "Smaller groups use fewer resources per node but require more rounds." - ), - ) - parallel_group_size: int = Field( - default=128, - ge=2, - le=512, - description="Legacy alias for nodes_per_full_mesh_group.", - ) - ibv_test_timeout: int = Field( - default=90, - ge=1, - description="Timeout in seconds for RDMA connectivity tests using ibv_rc_pingpong", - ) - ibv_test_port_range: str = Field( - default="10000-50000", description="Port range for RDMA connectivity tests (format: start-end)" - ) - inter_full_mesh_group_pairs_per_wave: str = Field( - default="auto", description="Max ordered group-pairs per wave during inter-group testing ('auto' or integer)" - ) - inter_group_wave_pairs: str = Field( - default="auto", - description="Legacy alias for inter_full_mesh_group_pairs_per_wave.", - ) - prune_failure_threshold: float = Field( - default=0.5, - gt=0.0, - le=1.0, - description=( - "Round 1 (intra) prune before inter-group: prune nodes whose fraction of peers with ≥1 FAIL " - "intra test is ≥ this value (default 0.5). Peers counted per distinct other node in the same partition group." - ), - ) - port_retry_max: int = Field( - default=3, - ge=0, - le=10, - description=( - "After each ScriptLet wave (intra/inter), rerun only pairs whose logs show PORT_LISTEN_FAILED, " - "up to this many extra batches with new TCP ports (default 3)." - ), - ) - port_retry_gap: int = Field( - default=1000, - ge=1, - le=65535, - description=( - "When remapping ports for PORT_LISTEN_FAILED retries, start at (max port in batch) + this gap " - "to reduce overlap with ephemeral ports." - ), - ) - exclude_failed_interface_nodes: str = Field( - default="true", - description=( - "Legacy hint for reporting: preflight now prunes interface- and GID-failed nodes from the SSH " - "host list before RDMA; interface failures are not run in the mesh regardless of this flag." - ), - ) - - @field_validator('connectivity_mode') - @classmethod - def validate_connectivity_check(cls, v: str) -> str: - """Validate RDMA connectivity check setting.""" - valid_modes = ['basic', 'full_mesh', 'skip'] - if v not in valid_modes: - raise ValueError(f"connectivity_mode must be one of: {', '.join(valid_modes)}") - return v - - @field_validator('ibv_test_port_range') - @classmethod - def validate_port_range(cls, v: str) -> str: - """Validate port range format.""" - try: - start, end = map(int, v.split('-')) - if start >= end or start < 1024 or end > 65535: - raise ValueError("Invalid port range") - except (ValueError, AttributeError): - raise ValueError("ibv_test_port_range must be in format 'start-end' with valid port numbers") - return v - - -class PreflightL2PingConfig(BaseModel): - """Small customer-facing IFoE L2 ping policy.""" - - model_config = ConfigDict(extra="forbid") - - enabled: bool = Field(default=False, description="Enable the mandatory IFoE L2 connectivity gate") - pings_per_port: int = Field(default=3, ge=1, description="Ping samples per selected IFoE port pair") - - -class PreflightTransferBenchConfig(BaseModel): - """Small customer-facing TransferBench preflight policy.""" - - model_config = ConfigDict(extra="forbid") - - enabled: bool = Field(default=False, description="Enable the mandatory TransferBench preflight gate") - scope: str = Field(default="node", description="node for independent runs or cluster for one multi-rank run") - profile: str = Field(default="smoketest", description="CVS-supported TransferBench validation profile") - message_sizes: List[str] = Field( - default_factory=lambda: ["1K", "16M"], - min_length=1, - description="Message sizes exercised by the selected profile", - ) - iterations: int = Field(default=2, ge=1, description="Validated iterations per test and message size") - warmup_iterations: int = Field(default=0, ge=0, description="Warmup iterations before validation") - - @field_validator('scope') - @classmethod - def validate_transferbench_scope(cls, value: str) -> str: - normalized = value.strip().lower() - if normalized not in ('node', 'cluster'): - raise ValueError("TransferBench scope must be one of: node, cluster") - return normalized - - @field_validator('profile') - @classmethod - def validate_transferbench_profile(cls, value: str) -> str: - normalized = value.strip().lower() - if normalized != 'smoketest': - raise ValueError("TransferBench profile must be a CVS-supported profile: smoketest") - return normalized - - @field_validator('message_sizes') - @classmethod - def validate_transferbench_message_sizes(cls, value: List[str]) -> List[str]: - normalized = [str(size).strip() for size in value] - if any(not size for size in normalized): - raise ValueError("TransferBench message_sizes entries must not be empty") - return normalized - - -class PreflightIfoeConfig(BaseModel): - """MI4XX IFoE admission and data-path checks.""" - - model_config = ConfigDict(extra="forbid") - - fabric_checks: bool = Field( - default=False, - description="Enable MI4XX AIFM, AFM, vPOD, station-mask, and IFoE port admission", - ) - l2ping: PreflightL2PingConfig = Field( - default_factory=PreflightL2PingConfig, - description="Strict IFoE L2 connectivity admission", - ) - transferbench: PreflightTransferBenchConfig = Field( - default_factory=PreflightTransferBenchConfig, - description="TransferBench IFoE data-path validation", - ) - - -class PreflightConnectivityCheckConfig(BaseModel): - """Connectivity check settings by protocol.""" - - model_config = ConfigDict(extra="allow") - - rdma: PreflightRdmaConfig = Field(default_factory=PreflightRdmaConfig, description="RDMA connectivity settings") - ifoe: PreflightIfoeConfig = Field(default_factory=PreflightIfoeConfig, description="IFoE connectivity settings") - - -class PreflightNodeSmokeConfig(BaseModel): - """Primus node_smoke settings (primus-cli direct -- node_smoke).""" - - model_config = ConfigDict(extra="allow") - - connectivity_mode: str = Field( - default="skip", - description="Primus node_smoke mode: 'run' (host/GPU/RDMA roll-call) or 'skip' (default)", - ) - auto_setup: bool = Field( - default=True, - description="Clone/update Primus and prepare venv on each node before node_smoke", - ) - setup_timeout: int = Field(default=600, ge=60, description="SSH timeout in seconds for Primus auto_setup") - force_reclone: bool = Field( - default=False, - description="Remove primus_dir and clone fresh on every run (destructive)", - ) - shared_install: bool = Field( - default=True, - description=( - "When true (default), clone and venv setup run only on the first reachable node; " - "other nodes wait for the shared NFS home install. Set false only if each node has " - "a local primus_dir/venv_activate path." - ), - ) - pip_install_mode: str = Field( - default="minimal", - description="Venv deps: minimal (torch only), requirements, or skip", - ) - torch_pip_index_url: str = Field( - default="https://download.pytorch.org/whl/rocm6.2", - description="PyTorch ROCm wheel index URL for minimal pip_install_mode", - ) - primus_git_url: str = Field( - default="https://github.com/AMD-AIG-AIMA/Primus.git", - description="Primus repository URL for auto_setup clone", - ) - primus_git_branch: str = Field( - default="dev/preflight-direct-test", - description="Git branch to checkout during auto_setup", - ) - primus_git_recurse_submodules: bool = Field( - default=False, - description="Clone git submodules during auto_setup (not required for node_smoke)", - ) - primus_dir: str = Field( - default="/home/{user-id}/INSTALL/Primus", - description="Path to cloned Primus repo under the user's home directory (required when connectivity_mode is 'run')", - ) - venv_activate: str = Field( - default="/home/{user-id}/envs/preflight/.venv/bin/activate", - description="Path to Python venv activate script on each node (required when connectivity_mode is 'run')", - ) - gpus_per_node: int = Field(default=8, ge=1, description="GPUs per node for node_smoke") - master_port: int = Field(default=1234, ge=1024, le=65535, description="Distributed master port for node_smoke") - dump_path: str = Field( - default="", - description="Per-node dump directory for smoke JSON (default: /node_smoke)", - ) - expected_rdma_nics: Optional[int] = Field( - default=None, - ge=1, - description="Hard-fail when training RDMA NIC count differs (default: len(node_check.rdma_interfaces))", - ) - ulimit_l_min_gb: float = Field(default=32.0, ge=0, description="Minimum RLIMIT_MEMLOCK in GiB (0 disables)") - shm_min_gb: float = Field(default=8.0, ge=0, description="Minimum /dev/shm size in GiB (0 disables)") - skip_dmesg: bool = Field(default=False, description="Skip dmesg error scan (e.g. unprivileged containers)") - allow_foreign_procs: bool = Field( - default=False, - description="Do not FAIL nodes with foreign GPU processes (still reported)", - ) - allowed_procs: str = Field( - default="gpuagent,rocm-smi-daemon,amd-smi,dcgm-exporter", - description="Comma-separated process names allowed to hold GPUs", - ) - require_tools: str = Field( - default="", - description="Comma-separated CLI tools that must exist in PATH (empty = warn only)", - ) - nccl_socket_ifname: str = Field(default="", description="NCCL_SOCKET_IFNAME override for node_smoke") - gloo_socket_ifname: str = Field( - default="", description="GLOO_SOCKET_IFNAME override (defaults to nccl_socket_ifname)" - ) - nccl_ib_hca: str = Field(default="", description="NCCL_IB_HCA override (defaults to node_check.rdma_interfaces)") - nccl_ib_gid_index: Optional[int] = Field( - default=None, - description="NCCL_IB_GID_INDEX override (defaults to node_check.gid_index)", - ) - rdma_nic_allowlist: str = Field( - default="", - description="Training NIC allowlist for node_smoke (defaults to node_check.rdma_interfaces)", - ) - ssh_timeout: int = Field(default=300, ge=30, description="SSH timeout in seconds for each node_smoke run") - tier2_perf: bool = Field( - default=False, - description=( - "Enable Primus node_smoke Tier 2 perf sanity (--tier2-perf): " - "8192³ GEMM TFLOPS floor, HBM D2D bandwidth, local multi-GPU RCCL all-reduce" - ), - ) - gemm_tflops_min: float = Field( - default=600.0, - ge=0, - description="Tier 2 large GEMM TFLOPS floor (--gemm-tflops-min); used when tier2_perf is true", - ) - hbm_gbs_min: float = Field( - default=2000.0, - ge=0, - description="Tier 2 HBM device-to-device bandwidth floor in GB/s (--hbm-gbs-min)", - ) - rccl_gbs_min: float = Field( - default=100.0, - ge=0, - description="Tier 2 local multi-GPU RCCL all-reduce bandwidth floor in GB/s (--rccl-gbs-min)", - ) - rccl_size_mb: int = Field( - default=64, - ge=1, - description="Tier 2 local RCCL all-reduce message size in MB (--rccl-size-mb)", - ) - rccl_timeout_sec: int = Field( - default=120, - ge=30, - description="Tier 2 local RCCL all-reduce hard timeout in seconds (--rccl-timeout-sec)", - ) - extra_args: List[str] = Field( - default_factory=list, - description="Additional node_smoke CLI flags forwarded to primus-cli", - ) - - @field_validator("connectivity_mode") - @classmethod - def validate_node_smoke_mode(cls, v: str) -> str: - valid_modes = ["run", "skip"] - if v not in valid_modes: - raise ValueError(f"node_smoke.connectivity_mode must be one of: {', '.join(valid_modes)}") - return v - - -class PreflightReportingConfig(BaseModel): - """Report generation and output settings.""" - - model_config = ConfigDict(extra="allow") - - generate_html_report: bool = Field(default=True, description="Whether to generate HTML report") - artifacts_root_dir: str = Field( - default="/tmp/preflight", - description=( - "Root directory for preflight artifacts. HTML report output and RDMA full_mesh ScriptLet logs use " - "/rdma_connectivity_workspace/// on each node (NFS-friendly)." - ), - ) - generate_rdma_pairs_csv: bool = Field( - default=True, - description="If true, write preflight_report_*_rdma_pairs.csv beside the HTML report (failed pairs only)", - ) - - -class PreflightConfigFile(BaseModel): - """ - Schema for preflight check configuration file. - - Uses nested structure organized by execution phase for better organization. - """ - - model_config = ConfigDict(extra="allow") # Allow comment fields - - parallelism: PreflightParallelismConfig = Field( - default_factory=PreflightParallelismConfig, description="Parallel execution settings" - ) - debug: PreflightDebugConfig = Field( - default_factory=PreflightDebugConfig, description="Debug and troubleshooting options" - ) - node_check: PreflightNodeCheckConfig = Field( - default_factory=PreflightNodeCheckConfig, description="Individual node validation settings" - ) - connectivity_check: PreflightConnectivityCheckConfig = Field( - default_factory=PreflightConnectivityCheckConfig, description="Inter-node connectivity tests" - ) - node_smoke: PreflightNodeSmokeConfig = Field( - default_factory=PreflightNodeSmokeConfig, description="Primus node_smoke checks" - ) - reporting: PreflightReportingConfig = Field( - default_factory=PreflightReportingConfig, description="Report generation and output settings" - ) - - @model_validator(mode="before") - @classmethod - def reject_flat_preflight_checks(cls, value): - if not isinstance(value, dict): - return value - cleaned = strip_preflight_metadata(value) - removed = sorted(set(cleaned) & {"node_health", "l2ping", "transferbench"}) - if removed: - raise ValueError( - "Unsupported flat preflight block(s): " - + ", ".join(removed) - + "; use node_check and connectivity_check.ifoe" - ) - normalized, warning_message = normalize_legacy_preflight_rdma_config(cleaned) - if warning_message: - warnings.warn(warning_message, FutureWarning, stacklevel=2) - normalized, smoke_warning = normalize_legacy_preflight_node_smoke_sections(normalized) - if smoke_warning: - warnings.warn(smoke_warning, FutureWarning, stacklevel=2) - return normalized - - @model_validator(mode="after") - def validate_fabric_prerequisites(self): - if self.connectivity_check.ifoe.fabric_checks and not self.node_check.enabled: - raise ValueError("connectivity_check.ifoe.fabric_checks requires node_check.enabled=true") - return self - - -def validate_config_file( - config_path: Union[str, Path], config_type: str = "auto" -) -> Union[ - AortaBenchmarkConfigFile, - ClusterConfigFile, - PytorchXditWanConfigFile, - PytorchXditFluxConfigFile, - PreflightConfigFile, -]: - """ - Load and validate a configuration file. - - Args: - config_path: Path to configuration file (YAML or JSON) - config_type: Type of config - "aorta", "cluster", "pytorch_xdit_wan", "pytorch_xdit_flux", "preflight", or "auto" (detect from content) - - Returns: - Validated Pydantic model - - Raises: - ValueError: If config is invalid with detailed error message - FileNotFoundError: If config file doesn't exist - """ - import json - import yaml - - config_path = Path(config_path) - - if not config_path.exists(): - raise FileNotFoundError(f"Configuration file not found: {config_path}") - - # Load file - with open(config_path) as f: - if config_path.suffix in ('.yaml', '.yml'): - raw_config = yaml.safe_load(f) - else: - raw_config = json.load(f) - - if raw_config is None: - raise ValueError(f"Configuration file is empty: {config_path}") - - # Determine config type - if config_type == "auto": - if "node_dict" in raw_config: - config_type = "cluster" - elif "preflight" in raw_config: - config_type = "preflight" - elif "aorta_path" in raw_config: - config_type = "aorta" - elif "config" in raw_config and "benchmark_params" in raw_config: - # Check if it's a pytorch_xdit config (WAN or Flux) - config_section = raw_config.get("config", {}) - benchmark_section = raw_config.get("benchmark_params", {}) - - # Detect Flux: check for flux1_dev_t2i in benchmark_params or FLUX in model_repo - if "flux1_dev_t2i" in benchmark_section or "FLUX" in config_section.get("model_repo", ""): - config_type = "pytorch_xdit_flux" - # Detect WAN: check for wan22_i2v_a14b in benchmark_params or Wan in model_repo - elif "wan22_i2v_a14b" in benchmark_section or "Wan" in config_section.get("model_repo", ""): - config_type = "pytorch_xdit_wan" - else: - # Generic pytorch_xdit - default to WAN for backward compatibility - config_type = "pytorch_xdit_wan" - else: - raise ValueError( - f"Cannot auto-detect config type for {config_path}. " - f"Specify config_type='aorta', config_type='cluster', config_type='pytorch_xdit_wan', config_type='pytorch_xdit_flux', or config_type='preflight'" - ) - - # Validate with appropriate schema - try: - if config_type == "cluster": - return ClusterConfigFile.model_validate(raw_config) - elif config_type == "preflight": - # Extract preflight section for validation - if "preflight" in raw_config: - return PreflightConfigFile.model_validate(raw_config["preflight"]) - else: - raise ValueError("Preflight config must contain 'preflight' section") - elif config_type == "aorta": - return AortaBenchmarkConfigFile.model_validate(raw_config) - elif config_type == "pytorch_xdit_wan": - return PytorchXditWanConfigFile.model_validate(raw_config) - elif config_type == "pytorch_xdit_flux": - return PytorchXditFluxConfigFile.model_validate(raw_config) - else: - raise ValueError(f"Unknown config_type: {config_type}") - except Exception as e: - # Re-raise with file context - raise ValueError(f"Invalid configuration in {config_path}:\n{e}") from e diff --git a/cvs/schema/__init__.py b/cvs/schema/__init__.py index e5a7bcffb..7402862c7 100644 --- a/cvs/schema/__init__.py +++ b/cvs/schema/__init__.py @@ -1,13 +1,49 @@ """CVS Pydantic schemas for data validation.""" -from .rccl import ( +from cvs.schema.cluster_file.cluster import ( + ClusterConfigFile, + ClusterNodeConfig, + HeadNodeConfig, + RackConfig, + RacksBlock, +) +from cvs.schema.config_file.aorta.benchmark import AortaBenchmarkConfigFile +from cvs.schema.config_file.inference.pytorch_xdit.config import ( + PytorchXditFluxConfigFile, + PytorchXditWanConfigFile, +) +from cvs.schema.config_file.preflight.config import PreflightConfigFile +from cvs.schema.config_file.training.jaxmaxtext.variant import TrainingVariantConfig +from cvs.schema.config_file.training.megatron.variant import MegatronVariantConfig +from cvs.schema.config_file.training.torchtitan.variant import TorchTitanVariantConfig +from cvs.schema.config_file.inference.atom.variant import AtomVariantConfig +from cvs.schema.config_file.inference.sglang.variant import SglangSingleVariantConfig +from cvs.schema.config_file.inference.vllm.variant import VariantConfig as VllmVariantConfig +from cvs.schema.rccl import ( RcclTests, - RcclTestsMultinodeRaw, RcclTestsAggregated, + RcclTestsMultinodeRaw, ) +from cvs.schema.validate import validate_config_file __all__ = [ 'RcclTests', 'RcclTestsMultinodeRaw', 'RcclTestsAggregated', + 'AortaBenchmarkConfigFile', + 'ClusterConfigFile', + 'ClusterNodeConfig', + 'HeadNodeConfig', + 'MegatronVariantConfig', + 'PreflightConfigFile', + 'PytorchXditFluxConfigFile', + 'TorchTitanVariantConfig', + 'TrainingVariantConfig', + 'AtomVariantConfig', + 'SglangSingleVariantConfig', + 'VllmVariantConfig', + 'PytorchXditWanConfigFile', + 'RackConfig', + 'RacksBlock', + 'validate_config_file', ] diff --git a/cvs/schema/base.py b/cvs/schema/base.py new file mode 100644 index 000000000..729e7f0d2 --- /dev/null +++ b/cvs/schema/base.py @@ -0,0 +1,11 @@ +"""Shared Pydantic base classes for CVS config schemas.""" + +from pydantic import BaseModel, ConfigDict + + +class _Forbid(BaseModel): + model_config = ConfigDict(extra="forbid") + + +class _Allow(BaseModel): + model_config = ConfigDict(extra="allow") diff --git a/cvs/schema/cluster_file/__init__.py b/cvs/schema/cluster_file/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/cvs/schema/cluster_file/cluster.py b/cvs/schema/cluster_file/cluster.py new file mode 100644 index 000000000..3f2694143 --- /dev/null +++ b/cvs/schema/cluster_file/cluster.py @@ -0,0 +1,171 @@ +""" +Cluster file configuration schema. + +Mirrors ``cvs/input/cluster_file/``. Used by ``cvs.schema.validate.validate_config_file``. +""" + +import warnings +from typing import Dict, List, Optional + +from pydantic import BaseModel, ConfigDict, Field, field_validator, model_validator + + +class ClusterNodeConfig(BaseModel): + """Schema for a single node entry in cluster.json node_dict.""" + + model_config = ConfigDict(extra="allow") # Allow extra fields like bmc_ip, rack_id + + vpc_ip: str = Field(description="VPC IP or hostname for inter-node communication") + bmc_ip: Optional[str] = Field(default=None, description="BMC IP for out-of-band management") + + +class HeadNodeConfig(BaseModel): + """Schema for head_node_dict in cluster.json.""" + + model_config = ConfigDict(extra="allow") + + mgmt_ip: str = Field(description="Management IP of head node") + + +class RackConfig(BaseModel): + """ + Schema for a single rack entry inside the 'racks' block of cluster.json. + + A rack groups compute trays (referenced via node_dict rack_id) and the + switch trays physically associated with that rack. + """ + + model_config = ConfigDict(extra="allow") + + platform: Optional[str] = Field(default=None, description="ARC platform name, e.g. 'HeliosP' or 'HeliosR'") + arc_controller: Optional[str] = Field( + default=None, + description="IP of the ARC controller node. Defaults to first sorted node_dict entry with matching rack_id.", + ) + switch_trays: List[str] = Field( + default_factory=list, + description="IPs of switch trays in this rack", + ) + rmc: Optional[str] = Field(default=None, description="IP of the Rack Management Controller") + + +class RacksBlock(BaseModel): + """ + Schema for the top-level 'racks' field in cluster.json. + + Holds optional global switch credentials and one RackConfig entry per rack + (keyed by rack ID, e.g. 'rack-01'). Extra keys (rack IDs) are accepted via + extra='allow' and retrieved via get_racks(). + + Switch credentials are fleet-wide (homogeneous across all racks). Per-rack + overrides are not supported in the current exec path; add them to RackConfig + when that need arises. + """ + + model_config = ConfigDict(extra="allow") + + switch_ssh_user: Optional[str] = Field( + default=None, + description="SSH username for all switch trays in every rack.", + ) + switch_ssh_password: Optional[str] = Field( + default=None, + description="SSH password for all switch trays. Ignored when switch_ssh_key_file is set.", + ) + switch_ssh_key_file: Optional[str] = Field( + default=None, + description="Path to SSH private key for all switch trays. Takes priority over switch_ssh_password when set.", + ) + + def get_racks(self): + """Return only the rack entries, excluding credential fields.""" + skip = {'switch_ssh_user', 'switch_ssh_password', 'switch_ssh_key_file'} + result = {} + for key, value in (self.__pydantic_extra__ or {}).items(): + if key not in skip and isinstance(value, dict): + result[key] = RackConfig(**value) + return result + + +class ClusterConfigFile(BaseModel): + """ + Schema for cluster.json configuration file. + + Validates the cluster configuration before running benchmarks. + Fails fast with clear error messages if required fields are missing. + """ + + model_config = ConfigDict(extra="allow") + + username: str = Field(description="SSH username for cluster nodes") + priv_key_file: Optional[str] = Field(default=None, description="Path to SSH private key") + password: Optional[str] = Field(default=None, description="SSH password (if not using key)") + + node_dict: Dict[str, ClusterNodeConfig] = Field( + description="Dictionary mapping node hostname/IP to node configuration" + ) + head_node_dict: Optional[HeadNodeConfig] = Field(default=None, description="Head node configuration") + + racks: Optional[RacksBlock] = Field( + default=None, + description=( + "Rack topology block. Contains optional global switch credentials and one entry per rack " + "(keyed by rack ID) listing switch_trays and platform." + ), + ) + rack_groups: Optional[RacksBlock] = Field( + default=None, + description="Deprecated alias for 'racks'. Use 'racks' instead.", + ) + + # Optional fields that may be present + home_mount_dir_name: Optional[str] = Field(default="home") + node_dir_name: Optional[str] = Field(default="root") + + @model_validator(mode='after') + def validate_auth_method(self): + """Ensure at least one authentication method is provided.""" + if not self.priv_key_file and not self.password: + raise ValueError("Authentication required: provide either 'priv_key_file' or 'password' in cluster config") + return self + + @model_validator(mode='after') + def validate_nodes_exist(self): + """Ensure at least one node is configured.""" + if not self.node_dict: + raise ValueError("No nodes configured in 'node_dict' - at least one node is required") + return self + + @model_validator(mode='after') + def warn_rack_groups_deprecated(self): + """Emit a deprecation warning when the old 'rack_groups' key is used.""" + if self.rack_groups is not None and self.racks is None: + warnings.warn( + "'rack_groups' in cluster.json is deprecated. Rename it to 'racks'.", + DeprecationWarning, + stacklevel=2, + ) + return self + + def get_racks_block(self): + """Return the active racks block, preferring 'racks' over the deprecated 'rack_groups'.""" + return self.racks if self.racks is not None else self.rack_groups + + @field_validator('username') + @classmethod + def validate_username_not_placeholder(cls, v): + """Check that username is not still a placeholder.""" + if '' in v.lower(): + raise ValueError( + "Username contains placeholder ''. Please set a valid username in cluster config." + ) + return v + + +__all__ = [ + "ClusterConfigFile", + "ClusterNodeConfig", + "HeadNodeConfig", + "RackConfig", + "RacksBlock", +] diff --git a/cvs/schema/cluster_file/unittests/__init__.py b/cvs/schema/cluster_file/unittests/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/cvs/schema/cluster_file/unittests/test_cluster.py b/cvs/schema/cluster_file/unittests/test_cluster.py new file mode 100644 index 000000000..d8b6de9ab --- /dev/null +++ b/cvs/schema/cluster_file/unittests/test_cluster.py @@ -0,0 +1,51 @@ +"""Unit tests for cluster.json schema (cvs/schema/cluster_file/cluster.py).""" + +import json +import unittest +from pathlib import Path + +from pydantic import ValidationError + +from cvs.schema.cluster_file.cluster import ClusterConfigFile + +_PACKAGE_ROOT = Path(__file__).resolve().parents[3] +_SAMPLE_CLUSTER = _PACKAGE_ROOT / "input" / "cluster_file" / "cluster.json" + + +class TestClusterConfigFile(unittest.TestCase): + def test_sample_cluster_json_validates(self): + raw = json.loads(_SAMPLE_CLUSTER.read_text()) + config = ClusterConfigFile.model_validate(raw) + self.assertGreater(len(config.node_dict), 0) + first_node = next(iter(config.node_dict.values())) + self.assertTrue(first_node.vpc_ip) + + def test_minimal_node_dict_required(self): + with self.assertRaises(ValidationError): + ClusterConfigFile.model_validate({"node_dict": {}}) + + def test_extra_top_level_keys_allowed(self): + config = ClusterConfigFile.model_validate( + { + "username": "testuser", + "priv_key_file": "/home/testuser/.ssh/id_rsa", + "node_dict": {"host1": {"vpc_ip": "10.0.0.1"}}, + "orchestrator": {"type": "baremetal"}, + "env_vars": {"PATH": "/bin"}, + } + ) + self.assertIn("host1", config.node_dict) + + def test_node_requires_vpc_ip(self): + with self.assertRaises(ValidationError): + ClusterConfigFile.model_validate( + { + "username": "testuser", + "priv_key_file": "/home/testuser/.ssh/id_rsa", + "node_dict": {"host1": {"bmc_ip": "1.2.3.4"}}, + } + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/cvs/schema/common/__init__.py b/cvs/schema/common/__init__.py new file mode 100644 index 000000000..c89e9af0c --- /dev/null +++ b/cvs/schema/common/__init__.py @@ -0,0 +1,17 @@ +"""Shared config schema primitives used across CVS suites.""" + +from cvs.schema.common.base import ( + BaseVariantConfig, + ContainerSpec, + ModelSpec, + Paths, + RuntimeSpec, +) + +__all__ = [ + "BaseVariantConfig", + "ContainerSpec", + "ModelSpec", + "Paths", + "RuntimeSpec", +] diff --git a/cvs/schema/common/base.py b/cvs/schema/common/base.py new file mode 100644 index 000000000..318991252 --- /dev/null +++ b/cvs/schema/common/base.py @@ -0,0 +1,78 @@ +"""Framework-agnostic variant config blocks shared by every CVS suite.""" + +from __future__ import annotations + +from typing import Any, Dict + +from pydantic import Field, model_validator +from typing_extensions import Literal + +from cvs.schema.base import _Allow, _Forbid + + +class Paths(_Forbid): + shared_fs: str + models_dir: str + log_dir: str + hf_token_file: str + # Host-user-namespaced scratch (jaxmaxtext launchers/yml). Optional so + # inference configs that omit it still load; jaxmaxtext configs set it to + # /tmp/{user-id}/jaxmaxtext so container-root /tmp/root is never used. + temp_dir: str = "" + + +class ModelSpec(_Forbid): + id: str + remote: Literal[0, 1] + precision: str = "" + + +class RuntimeSpec(_Allow): + name: str + args: Dict[str, Any] = Field(default_factory=dict) + + +class ContainerSpec(_Forbid): + lifetime: Literal["no_launch", "per_run", "persistent"] = "per_run" + name: str + image: str + runtime: RuntimeSpec + + +class BaseVariantConfig(_Forbid): + """The framework-agnostic skeleton of a variant config. + + Carries the fields every suite shares (schema/paths/model/image/container/ + thresholds + the enforce gate) and the remote-not-implemented guard. + Per-framework subclasses add their own ``framework``/``Params``/``Sweep`` and the + ``cell_key``/coverage-check pair that depend on them. + """ + + schema_version: Literal[1] + # When false, the threshold-coverage gate warns instead of raising and the + # test records metrics without asserting pass/fail (record-only). Use for + # un-calibrated shapes (e.g. a throughput characterization whose published + # numbers are curves, not tabulated values). Default true keeps the gate + # strict for calibrated configs -- no regression to the remediation work. + enforce_thresholds: bool = True + threshold_json: str = "" + paths: Paths + model: ModelSpec + # The container image is declared once, on container.image (ContainerSpec). + # There is no separate top-level image block. + container: ContainerSpec + thresholds: Dict[str, Dict[str, Any]] = Field(default_factory=dict) + + # pydantic runs @model_validator(mode="after") hooks in definition order, + # parent-class hooks before subclass hooks. This remote check is intentionally + # the first to run: an unimplemented remote config fails fast + # (NotImplementedError) before any subclass's threshold-coverage check runs, + # which is meaningless for a config we are going to reject anyway. + @model_validator(mode="after") + def _check_remote_not_implemented(self): + if self.model.remote == 1: + raise NotImplementedError( + "model.remote=1 (remote model download) is not implemented in the PoC. " + "Port from cvs-dtni-v1/resource_resolver.py before enabling." + ) + return self diff --git a/cvs/schema/config_file/__init__.py b/cvs/schema/config_file/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/cvs/schema/config_file/aorta/__init__.py b/cvs/schema/config_file/aorta/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/cvs/schema/config_file/aorta/benchmark.py b/cvs/schema/config_file/aorta/benchmark.py new file mode 100644 index 000000000..dc4165ad6 --- /dev/null +++ b/cvs/schema/config_file/aorta/benchmark.py @@ -0,0 +1,207 @@ +""" +Aorta benchmark configuration file schema. + +Mirrors ``cvs/input/config_file/aorta/`` (``aorta_benchmark.yaml``). +""" + +from pathlib import Path +from typing import Any, Dict, List, Optional + +from pydantic import BaseModel, ConfigDict, Field, field_validator + + +class AortaDockerConfigFile(BaseModel): + """Schema for docker section in aorta_benchmark.yaml.""" + + model_config = ConfigDict(extra="forbid") # Catch typos + + image: str = Field( + default="jeffdaily/pytorch:torchrec-dlrm-complete", description="Docker image for Aorta container" + ) + container_name: str = Field(default="aorta-benchmark", description="Name for the Docker container") + shm_size: str = Field(default="17G", description="Shared memory size") + network_mode: str = Field(default="host", description="Docker network mode") + privileged: bool = Field(default=True, description="Run container in privileged mode") + + +class AortaRcclConfigFile(BaseModel): + """Schema for rccl section in aorta_benchmark.yaml.""" + + model_config = ConfigDict(extra="forbid") + + clone_url: str = Field( + default="https://github.com/ROCmSoftwarePlatform/rccl.git", description="RCCL git repository URL" + ) + branch: str = Field(default="develop", description="RCCL branch to build") + build_path: str = Field(default="/mnt/rccl", description="Path inside container for RCCL build") + + +class AortaEnvironmentConfigFile(BaseModel): + """Schema for environment section in aorta_benchmark.yaml.""" + + model_config = ConfigDict(extra="allow") # Allow custom env vars + + NCCL_MAX_NCHANNELS: int = Field(default=112, ge=1, le=256, description="Maximum NCCL channels") + NCCL_MAX_P2P_NCHANNELS: int = Field(default=112, ge=1, le=256, description="Maximum NCCL P2P channels") + NCCL_DEBUG: str = Field(default="VERSION", description="NCCL debug level") + TORCH_NCCL_HIGH_PRIORITY: int = Field(default=1, ge=0, le=1, description="Enable high priority NCCL streams") + OMP_NUM_THREADS: int = Field(default=1, ge=1, description="OpenMP thread count") + RCCL_MSCCL_ENABLE: int = Field(default=0, ge=0, le=1, description="Enable MSCCL") + + +class AortaExpectedResultsConfigFile(BaseModel): + """Schema for expected_results section in aorta_benchmark.yaml.""" + + model_config = ConfigDict(extra="allow") # Allow custom thresholds + + max_avg_iteration_ms: Optional[float] = Field( + default=None, ge=0, description="Maximum acceptable average iteration time in ms" + ) + min_compute_ratio: Optional[float] = Field(default=None, ge=0, le=1, description="Minimum acceptable compute ratio") + min_overlap_ratio: Optional[float] = Field( + default=None, ge=0, le=1, description="Minimum acceptable compute-comm overlap ratio" + ) + max_time_variance_ratio: Optional[float] = Field( + default=None, ge=0, description="Maximum acceptable iteration time variance" + ) + + +class AortaAnalysisConfigFile(BaseModel): + """Schema for analysis section in aorta_benchmark.yaml.""" + + model_config = ConfigDict(extra="forbid") + + enable_tracelens: bool = Field(default=True, description="Run Aorta's TraceLens analysis after benchmark") + enable_gemm_analysis: bool = Field(default=False, description="Run Aorta's GEMM analysis (for sweep experiments)") + tracelens_script: str = Field( + default="scripts/tracelens_single_config/run_tracelens_single_config.sh", + description="Path to TraceLens analysis script relative to aorta_path", + ) + gemm_script: str = Field( + default="scripts/gemm_analysis/run_tracelens_analysis.sh", + description="Path to GEMM analysis script relative to aorta_path", + ) + skip_if_exists: bool = Field( + default=False, description="Skip analysis if tracelens_analysis directory already exists" + ) + + +class AortaBenchmarkConfigFile(BaseModel): + """ + Schema for the entire aorta_benchmark.yaml configuration file. + + Validates structure and provides sensible defaults. + Fails fast with clear error messages if configuration is invalid. + + For ``test_aorta``, load YAML, apply ``resolve_test_config_placeholders`` with the resolved + cluster dict (same as other CVS suites), then ``model_validate``. Standalone tools may validate + raw YAML without placeholder resolution if paths are already absolute. + """ + + model_config = ConfigDict(extra="forbid") # Catch typos in top-level keys + + # Path to Aorta repository on host (will be bind-mounted). If missing and aorta_auto_clone is true, it is cloned. + aorta_path: str = Field(description="Path to Aorta repository on host (will be bind-mounted)") + + # Optional: clone Aorta repo when aorta_path does not exist + aorta_auto_clone: bool = Field( + default=False, description="If true and aorta_path missing, clone from aorta_clone_url" + ) + aorta_clone_url: Optional[str] = Field(default=None, description="Git URL to clone when aorta_auto_clone is true") + + # Container settings + container_mount_path: str = Field(default="/mnt", description="Mount point inside container for aorta_path") + + # Aorta config + base_config: str = Field(default="config/distributed.yaml", description="Aorta config file relative to aorta_path") + + # Nested configuration sections + docker: AortaDockerConfigFile = Field( + default_factory=AortaDockerConfigFile, description="Docker container configuration" + ) + rccl: AortaRcclConfigFile = Field(default_factory=AortaRcclConfigFile, description="RCCL build configuration") + environment: AortaEnvironmentConfigFile = Field( + default_factory=AortaEnvironmentConfigFile, description="Environment variables for RCCL/NCCL" + ) + + # Training overrides + training_overrides: Dict[str, Any] = Field( + default_factory=dict, description="Overrides passed to Aorta via --override flag" + ) + + # Scripts + build_script: str = Field( + default="scripts/build_rccl.sh", description="RCCL build script relative to container mount" + ) + experiment_script: str = Field( + default="scripts/rccl_exp.sh", description="Experiment script relative to container mount" + ) + + # Hardware + gpus_per_node: int = Field(default=8, ge=1, description="Number of GPUs per node") + + # Execution settings + timeout_seconds: int = Field(default=10800, ge=60, description="Benchmark timeout in seconds") + skip_rccl_build: bool = Field(default=False, description="Skip RCCL build if already built") + + # Validation thresholds + expected_results: AortaExpectedResultsConfigFile = Field( + default_factory=AortaExpectedResultsConfigFile, description="Expected results for validation" + ) + + # Analysis configuration (use Aorta's built-in analysis scripts) + analysis: AortaAnalysisConfigFile = Field( + default_factory=AortaAnalysisConfigFile, description="Post-benchmark analysis configuration" + ) + + @field_validator('aorta_path') + @classmethod + def validate_aorta_path_not_placeholder(cls, v: str) -> str: + """Check that aorta_path is not a placeholder.""" + if '' in v.lower(): + raise ValueError( + "aorta_path contains placeholder ''. Please set the actual path to your Aorta installation." + ) + return v + + def validate_paths_exist(self) -> List[str]: + """ + Validate that referenced paths exist on the filesystem. + + Call this after loading config to check paths. + Returns list of error messages (empty if all valid). + """ + errors = [] + + aorta = Path(self.aorta_path) + if not aorta.exists(): + if self.aorta_auto_clone and self.aorta_clone_url: + # Runner will clone in setup(); skip path checks here + return errors + errors.append(f"aorta_path does not exist: {self.aorta_path}") + else: + # Check internal paths + base_cfg = aorta / self.base_config + if not base_cfg.exists(): + errors.append(f"base_config does not exist: {base_cfg}") + + build_script = aorta / self.build_script + if not build_script.exists(): + errors.append(f"build_script does not exist: {build_script}") + + exp_script = aorta / self.experiment_script + if not exp_script.exists(): + errors.append(f"experiment_script does not exist: {exp_script}") + + # Check analysis scripts if enabled + if self.analysis.enable_tracelens: + tracelens_script = aorta / self.analysis.tracelens_script + if not tracelens_script.exists(): + errors.append(f"tracelens_script does not exist: {tracelens_script}") + + if self.analysis.enable_gemm_analysis: + gemm_script = aorta / self.analysis.gemm_script + if not gemm_script.exists(): + errors.append(f"gemm_script does not exist: {gemm_script}") + + return errors diff --git a/cvs/schema/config_file/aorta/unittests/__init__.py b/cvs/schema/config_file/aorta/unittests/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/cvs/schema/config_file/aorta/unittests/test_benchmark.py b/cvs/schema/config_file/aorta/unittests/test_benchmark.py new file mode 100644 index 000000000..e88b9ad3f --- /dev/null +++ b/cvs/schema/config_file/aorta/unittests/test_benchmark.py @@ -0,0 +1,33 @@ +"""Unit tests for Aorta benchmark config schema (config_file/aorta/benchmark.py).""" + +import unittest +from pathlib import Path + +import yaml +from pydantic import ValidationError + +from cvs.schema.config_file.aorta.benchmark import AortaBenchmarkConfigFile + +_PACKAGE_ROOT = Path(__file__).resolve().parents[4] +_SAMPLE_CONFIG = _PACKAGE_ROOT / "input" / "config_file" / "aorta" / "aorta_benchmark.yaml" + + +class TestAortaBenchmarkConfigFile(unittest.TestCase): + def test_sample_yaml_validates(self): + raw = yaml.safe_load(_SAMPLE_CONFIG.read_text()) + config = AortaBenchmarkConfigFile.model_validate(raw) + self.assertIn("aorta", config.aorta_path) + self.assertFalse(config.analysis.enable_tracelens) + + def test_aorta_path_changeme_rejected(self): + with self.assertRaisesRegex(ValidationError, "placeholder ''"): + AortaBenchmarkConfigFile.model_validate({"aorta_path": "/path//aorta"}) + + def test_nested_defaults_applied(self): + config = AortaBenchmarkConfigFile.model_validate({"aorta_path": "/opt/aorta"}) + self.assertEqual(config.docker.container_name, "aorta-benchmark") + self.assertEqual(config.environment.NCCL_DEBUG, "VERSION") + + +if __name__ == "__main__": + unittest.main() diff --git a/cvs/schema/config_file/inference/__init__.py b/cvs/schema/config_file/inference/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/cvs/schema/config_file/inference/atom/__init__.py b/cvs/schema/config_file/inference/atom/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/cvs/schema/config_file/inference/atom/unittests/__init__.py b/cvs/schema/config_file/inference/atom/unittests/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/cvs/schema/config_file/inference/atom/unittests/test_variant.py b/cvs/schema/config_file/inference/atom/unittests/test_variant.py new file mode 100644 index 000000000..7fd388c05 --- /dev/null +++ b/cvs/schema/config_file/inference/atom/unittests/test_variant.py @@ -0,0 +1,582 @@ +"""Unit tests for ATOM inference variant schema (inference/atom/variant.py).""" + +import json +import unittest +from pathlib import Path + +from cvs.lib.inference.atom.atom_config_loader import ( + expand_sweep, + expand_sweep_parametrize, + load_variant, + orchestrator_container_from_variant, + placeholder_gated_threshold_cell, + reuse_server_flag, + server_session_key, +) +from cvs.schema.config_file.inference.common.sweep import Run, SeqCombo, Sweep +from cvs.schema.config_file.inference.atom.variant import AtomVariantConfig + +_PACKAGE_ROOT = Path(__file__).resolve().parents[5] + + +def _cluster_dict(): + return {"username": "testuser"} + + +class TestATOMAtomConfigLoader(unittest.TestCase): + def test_load_mi300x_sample_config(self): + root = _PACKAGE_ROOT + config = root / ("input/config_file/inference/atom/mi300x_atom_gpt-oss-120b_mxfp4_vllm_single.json") + variant = load_variant(config, _cluster_dict()) + self.assertEqual(variant.framework, "atom") + self.assertEqual(variant.params.driver, "vllm_atom") + self.assertEqual(variant.expected_cells(), ["ISL=8192,OSL=1024,TP=4,CONC=32", "ISL=8192,OSL=1024,TP=4,CONC=64"]) + self.assertIn("enforce-eager", variant.roles.server.serve_args) + + def test_load_w1_mi300x_atom_variant(self): + root = _PACKAGE_ROOT + config = root / ("input/config_file/inference/atom/mi300x_atom_deepseek-r1_fp8_single.json") + variant = load_variant(config, _cluster_dict()) + self.assertEqual(variant.threshold_json, "mi300x_atom_deepseek-r1_fp8_single_threshold.json") + self.assertEqual(variant.gpu_arch, "mi300x") + self.assertEqual(variant.params.driver, "atom") + self.assertEqual(variant.params.metric_percentiles, "95,99") + self.assertEqual( + variant.roles.server.atom_args[:4], + ["-tp", "8", "--kv_cache_dtype", "fp8"], + ) + self.assertEqual( + variant.expected_cells(), + ["ISL=1024,OSL=1024,TP=8,CONC=128", "ISL=1024,OSL=1024,TP=8,CONC=256"], + ) + cell = "ISL=1024,OSL=1024,TP=8,CONC=128" + for key in ( + "client.per_gpu_throughput", + "client.output_tput_per_gpu", + "client.p99_ttft_ms", + "client.p99_tpot_ms", + "client.p95_tpot_ms", + ): + self.assertIn(key, variant.thresholds[cell]) + + def test_load_w1_mi300x_multinode_variant(self): + root = _PACKAGE_ROOT + config = root / ("input/config_file/inference/atom/mi300x_atom_deepseek-r1_fp8_distributed.json") + variant = load_variant(config, _cluster_dict()) + self.assertEqual(variant.params.nnodes, "2") + self.assertEqual(variant.params.driver, "vllm_atom") + self.assertEqual(variant.params.pipeline_parallel_size, "2") + self.assertEqual(variant.roles.server.ib_netdev, "auto") + self.assertEqual(variant.roles.server.ib_hca_devices, "auto") + self.assertEqual(variant.params.scaling_baseline_output_throughput, "1500") + self.assertTrue(variant.enforce_thresholds) + self.assertEqual(len(variant.expected_cells()), 15) + cell = "ISL=512,OSL=512,TP=8,PP=2,NNODES=2,CONC=16" + self.assertIn(cell, variant.expected_cells()) + self.assertEqual( + variant.thresholds[cell]["scaling.efficiency_pct"], + {"kind": "min", "value": 11}, + ) + + def test_load_w1_mi300x_multinode_sglang_variant(self): + root = _PACKAGE_ROOT + config = root / ("input/config_file/inference/atom/mi300x_atom_deepseek-r1_fp8_sglang_distributed.json") + variant = load_variant(config, _cluster_dict()) + self.assertEqual(variant.params.driver, "sglang") + self.assertEqual(variant.params.pipeline_parallel_size, "2") + self.assertFalse(variant.enforce_thresholds) + cell = "ISL=512,OSL=512,TP=8,PP=2,NNODES=2,CONC=16" + self.assertIn(cell, variant.expected_cells()) + + def test_load_w1_mi355x_multinode_variant(self): + root = _PACKAGE_ROOT + config = root / ("input/config_file/inference/atom/mi355x_atom_deepseek-r1_fp8_distributed.json") + variant = load_variant(config, _cluster_dict()) + self.assertEqual(variant.gpu_arch, "mi355x") + self.assertEqual(variant.params.nnodes, "2") + self.assertEqual(variant.params.driver, "vllm_atom") + self.assertEqual(variant.params.pipeline_parallel_size, "2") + self.assertEqual(variant.params.scaling_baseline_output_throughput, "4000") + self.assertFalse(variant.enforce_thresholds) + self.assertEqual(len(variant.expected_cells()), 15) + cell = "ISL=512,OSL=512,TP=8,PP=2,NNODES=2,CONC=16" + self.assertIn(cell, variant.expected_cells()) + self.assertEqual( + variant.thresholds[cell]["scaling.efficiency_pct"], + {"kind": "min", "value": 50}, + ) + + def test_load_baseline_sweep_mi300x_variant(self): + root = _PACKAGE_ROOT + config = root / ("input/config_file/inference/atom/mi300x_atom_deepseek-r1_fp8_baseline_sweep.json") + variant = load_variant(config, _cluster_dict()) + self.assertEqual(variant.params.max_model_length, "10240") + self.assertTrue(variant.enforce_thresholds) + self.assertEqual(len(variant.expected_cells()), 14) + self.assertIn("ISL=1024,OSL=1024,TP=8,CONC=4", variant.expected_cells()) + self.assertIn("ISL=8192,OSL=1024,TP=8,CONC=256", variant.expected_cells()) + cell = "ISL=8192,OSL=1024,TP=8,CONC=128" + self.assertIn("client.output_throughput", variant.thresholds[cell]) + self.assertEqual(variant.thresholds[cell]["client.success_rate"]["value"], 1) + + def test_load_baseline_sweep_mi355x_variant(self): + root = _PACKAGE_ROOT + config = root / ("input/config_file/inference/atom/mi355x_atom_deepseek-r1_fp8_baseline_sweep.json") + variant = load_variant(config, _cluster_dict()) + self.assertEqual(variant.gpu_arch, "mi355x") + self.assertFalse(variant.enforce_thresholds) + self.assertEqual(len(variant.expected_cells()), 14) + + def test_load_w1_mi355x_atom_single_variant_and_thresholds(self): + root = _PACKAGE_ROOT + config = root / ("input/config_file/inference/atom/mi355x_atom_deepseek-r1_fp8_single.json") + variant = load_variant(config, _cluster_dict()) + self.assertEqual(variant.gpu_arch, "mi355x") + self.assertIn("--trust-remote-code", variant.roles.server.atom_args) + self.assertEqual( + variant.expected_cells(), + ["ISL=1024,OSL=1024,TP=8,CONC=128", "ISL=1024,OSL=1024,TP=8,CONC=256"], + ) + cell = "ISL=1024,OSL=1024,TP=8,CONC=128" + self.assertEqual( + variant.thresholds[cell]["client.output_throughput"]["value"], + 4004.66, + ) + self.assertEqual( + variant.thresholds[cell]["client.mean_ttft_ms"]["value"], + 362.18, + ) + + def test_load_w1_mi355x_atom_mtp3_inline_bench_args(self): + root = _PACKAGE_ROOT + config = root / ("input/config_file/inference/atom/mi355x_atom_deepseek-r1_fp8_mtp3.json") + variant = load_variant(config, _cluster_dict()) + self.assertIn("--method", variant.roles.server.atom_args) + self.assertEqual(variant.params.bench_extra_args, "--use-chat-template") + + def test_load_w1_mi355x_atom_mtp3_thresholds(self): + root = _PACKAGE_ROOT + config = root / ("input/config_file/inference/atom/mi355x_atom_deepseek-r1_fp8_mtp3.json") + variant = load_variant(config, _cluster_dict()) + cell = "ISL=1024,OSL=1024,TP=8,CONC=256" + self.assertEqual( + variant.thresholds[cell]["client.output_throughput"]["value"], + 6451.59, + ) + + def test_orchestrator_container_includes_server_env(self): + sweep = Sweep( + sequence_combinations=[SeqCombo(name="legacy_profile", isl="7168", osl="1024")], + runs=[Run(combo="legacy_profile", concurrency=64)], + ) + thresholds = { + "ISL=7168,OSL=1024,TP=8,CONC=64": placeholder_gated_threshold_cell(), + } + variant = AtomVariantConfig( + schema_version=1, + framework="atom", + gpu_arch="mi300x", + enforce_thresholds=False, + paths={ + "shared_fs": "/home/x", + "models_dir": "/home/x/models", + "log_dir": "/home/x/LOGS", + "hf_token_file": "/home/x/.hf", + }, + model={"id": "openai/gpt-oss-120b", "remote": 0, "precision": "bf16"}, + container={ + "name": "c", + "image": "img", + "runtime": {"name": "docker", "args": {"volumes": ["/home/x:/home/x"]}}, + }, + roles={"server": {"env": {"VLLM_ROCM_USE_AITER": "1"}}}, + params={"tensor_parallelism": "8"}, + sweep=sweep, + thresholds=thresholds, + ) + block = orchestrator_container_from_variant(variant) + self.assertEqual(block["env"]["VLLM_ROCM_USE_AITER"], "1") + + def test_expand_sweep_matches_w1_single(self): + root = _PACKAGE_ROOT + config = root / ("input/config_file/inference/atom/mi300x_atom_deepseek-r1_fp8_single.json") + import json + + raw = json.loads(config.read_text()) + cases, ids = expand_sweep(raw["sweep"]) + self.assertEqual(len(cases), 2) + self.assertEqual(ids[0], "w1_1k_1k-conc128") + self.assertEqual(ids[1], "w1_1k_1k-conc256") + self.assertEqual(cases[0][1], 128) + + def test_w1_single_threshold_health_gates_tight_when_enforcing(self): + root = _PACKAGE_ROOT + config = root / ("input/config_file/inference/atom/mi300x_atom_deepseek-r1_fp8_single.json") + variant = load_variant(config, _cluster_dict()) + self.assertTrue(variant.enforce_thresholds) + cell = "ISL=1024,OSL=1024,TP=8,CONC=128" + self.assertEqual(variant.thresholds[cell]["client.success_rate"]["value"], 1) + self.assertEqual(variant.thresholds[cell]["client.failed"]["value"], 0) + + def test_placeholder_threshold_cell_covers_gated_metrics(self): + cell = placeholder_gated_threshold_cell() + from cvs.lib.inference.atom.atom_parsing import GATED_METRICS + + for short in GATED_METRICS: + self.assertIn(f"client.{short}", cell, short) + + def test_atom_driver_requires_inline_atom_args(self): + sweep = Sweep( + sequence_combinations=[SeqCombo(name="w1", isl="1024", osl="1024")], + runs=[Run(combo="w1", concurrency=128)], + ) + thresholds = {"ISL=1024,OSL=1024,TP=8,CONC=128": placeholder_gated_threshold_cell()} + with self.assertRaises(ValueError): + AtomVariantConfig( + schema_version=1, + framework="atom", + gpu_arch="mi300x", + enforce_thresholds=False, + paths={ + "shared_fs": "/home/x", + "models_dir": "/home/x/models", + "log_dir": "/home/x/LOGS", + "hf_token_file": "/home/x/.hf", + }, + model={"id": "deepseek-ai/DeepSeek-R1-0528", "remote": 0, "precision": "fp8"}, + container={ + "name": "c", + "image": "img", + "runtime": {"name": "docker", "args": {"volumes": ["/home/x:/home/x"]}}, + }, + roles={"server": {"env": {}}}, + params={"driver": "atom", "tensor_parallelism": "8"}, + sweep=sweep, + thresholds=thresholds, + ) + + def test_reuse_server_flag_and_session_key_helpers(self): + from types import SimpleNamespace + + self.assertFalse(reuse_server_flag(SimpleNamespace())) + variant = SimpleNamespace( + model=SimpleNamespace(id="m"), + params=SimpleNamespace( + driver="atom", + tensor_parallelism="8", + nnodes="1", + pipeline_parallel_size="1", + master_addr="", + master_port="29501", + ), + roles=SimpleNamespace(server=SimpleNamespace(atom_args=("-tp", "8"), serve_args={}, sglang_args=[])), + ) + self.assertNotEqual(server_session_key(variant, "1", "2"), server_session_key(variant, "3", "4")) + + def test_expand_sweep_parametrize_tier_ids(self): + sweep = { + "sequence_combinations": [{"name": "w1", "isl": "1024", "osl": "1024"}], + "runs": [{"combo": "w1", "concurrency": 128}], + } + _, _, ids = expand_sweep_parametrize(sweep, ("metric_tier",)) + self.assertIn("w1-conc128-throughput", ids) + + def test_ib_netdev_coerces_mlx5_hca_name_to_auto(self): + sweep = Sweep( + sequence_combinations=[SeqCombo(name="w1", isl="512", osl="512")], + runs=[Run(combo="w1", concurrency=16)], + ) + variant = AtomVariantConfig( + schema_version=1, + framework="atom", + gpu_arch="mi300x", + enforce_thresholds=False, + paths={ + "shared_fs": "/home/x", + "models_dir": "/home/x/models", + "log_dir": "/home/x/LOGS", + "hf_token_file": "/home/x/.hf", + }, + model={"id": "deepseek-ai/DeepSeek-R1-0528", "remote": 0, "precision": "fp8"}, + container={ + "name": "c", + "image": "img", + "runtime": {"name": "docker", "args": {"volumes": ["/home/x:/home/x"]}}, + }, + roles={"server": {"ib_netdev": "mlx5_0"}}, + params={ + "driver": "vllm_atom", + "tensor_parallelism": "8", + "pipeline_parallel_size": "2", + "nnodes": "2", + "master_addr": "10.0.0.1", + }, + sweep=sweep, + thresholds={}, + ) + self.assertEqual(variant.roles.server.ib_netdev, "auto") + + def test_server_env_strips_orchestrator_network_keys(self): + sweep = Sweep( + sequence_combinations=[SeqCombo(name="w1", isl="512", osl="512")], + runs=[Run(combo="w1", concurrency=16)], + ) + variant = AtomVariantConfig( + schema_version=1, + framework="atom", + gpu_arch="mi300x", + enforce_thresholds=False, + paths={ + "shared_fs": "/home/x", + "models_dir": "/home/x/models", + "log_dir": "/home/x/LOGS", + "hf_token_file": "/home/x/.hf", + }, + model={"id": "deepseek-ai/DeepSeek-R1-0528", "remote": 0, "precision": "fp8"}, + container={ + "name": "c", + "image": "img", + "runtime": {"name": "docker", "args": {"volumes": ["/home/x:/home/x"]}}, + }, + roles={ + "server": { + "env": { + "GLOO_SOCKET_IFNAME": "mlx5_0", + "NCCL_IB_HCA": "mlx5_0", + "NCCL_IB_GID_INDEX": "1", + } + } + }, + params={ + "driver": "vllm_atom", + "tensor_parallelism": "8", + "pipeline_parallel_size": "2", + "nnodes": "2", + "master_addr": "10.0.0.1", + }, + sweep=sweep, + thresholds={}, + ) + self.assertEqual(variant.roles.server.env, {"NCCL_IB_GID_INDEX": "1"}) + + def test_load_w1_accuracy_variant(self): + root = _PACKAGE_ROOT + config = root / "input/config_file/inference/atom/mi300x_atom_deepseek-r1_fp8_accuracy.json" + variant = load_variant(config, _cluster_dict()) + self.assertEqual(len(variant.accuracy.tasks), 9) + self.assertEqual(variant.accuracy.tasks[0].id, "gsm8k_flex") + task_ids = {t.id for t in variant.accuracy.tasks} + self.assertIn("hellaswag", task_ids) + self.assertIn("mmlu_pro", task_ids) + self.assertIn("bbh", task_ids) + self.assertIn("arc_challenge", task_ids) + self.assertTrue(variant.quant_parity.enabled) + self.assertTrue(variant.enforce_thresholds) + self.assertIn( + "gsm8k.exact_match__flexible-extract", + variant.thresholds["accuracy"]["gsm8k_flex"], + ) + self.assertIn( + "hellaswag.acc_norm__none", + variant.thresholds["accuracy"]["hellaswag"], + ) + self.assertIn( + "mmlu_pro.exact_match__custom-extract", + variant.thresholds["accuracy"]["mmlu_pro"], + ) + self.assertIn("quant_parity", variant.thresholds) + + def test_load_mtp3_variant_mtp_quality_enabled(self): + root = _PACKAGE_ROOT + config = root / "input/config_file/inference/atom/mi300x_atom_deepseek-r1_fp8_mtp3.json" + variant = load_variant(config, _cluster_dict()) + self.assertTrue(variant.mtp_quality.enabled) + self.assertIn("mtp.acceptance_rate", variant.thresholds["mtp_quality"]) + + def test_mtp_quality_threshold_key_not_sweep_cell(self): + root = _PACKAGE_ROOT + config = root / "input/config_file/inference/atom/mi300x_atom_deepseek-r1_fp8_mtp3.json" + variant = load_variant(config, _cluster_dict()) + self.assertIn("mtp_quality", variant.thresholds) + self.assertEqual(len(variant.expected_cells()), 2) + + def test_load_w2_accuracy_long_context_cells(self): + root = _PACKAGE_ROOT + config = root / "input/config_file/inference/atom/mi300x_atom_gpt-oss-120b_mxfp4_accuracy.json" + variant = load_variant(config, _cluster_dict()) + self.assertTrue(variant.functional.api_smoke) + self.assertEqual(len(variant.long_context_accuracy.cells), 1) + self.assertEqual(variant.long_context_accuracy.cells[0].id, "niah_8k") + self.assertIn("long_context_accuracy", variant.thresholds) + + def test_load_mi355x_accuracy_variant(self): + root = _PACKAGE_ROOT + config = root / "input/config_file/inference/atom/mi355x_atom_deepseek-r1_fp8_accuracy.json" + variant = load_variant(config, _cluster_dict()) + self.assertEqual(variant.gpu_arch, "mi355x") + self.assertFalse(variant.enforce_thresholds) + + def test_load_phase_c_w2_mxfp4_perf(self): + root = _PACKAGE_ROOT + config = root / "input/config_file/inference/atom/mi300x_atom_gpt-oss-120b_mxfp4_single.json" + variant = load_variant(config, _cluster_dict()) + self.assertEqual(variant.params.driver, "atom") + self.assertEqual(variant.params.tensor_parallelism, "4") + self.assertEqual(variant.model.precision, "mxfp4") + self.assertEqual(variant.roles.server.env.get("ATOM_USE_TRITON_MOE"), "1") + self.assertEqual(variant.roles.server.env.get("ATOM_USE_TRITON_GEMM"), "1") + self.assertEqual( + variant.expected_cells(), + ["ISL=8192,OSL=1024,TP=4,CONC=32", "ISL=8192,OSL=1024,TP=4,CONC=64"], + ) + + def test_load_kimi_k27_mxfp4_triton_env(self): + root = _PACKAGE_ROOT + config = root / "input/config_file/inference/atom/mi355x_atom_kimi-k2.7-code_single.json" + variant = load_variant(config, _cluster_dict()) + self.assertEqual(variant.model.precision, "mxfp4") + self.assertEqual(variant.roles.server.env.get("ATOM_USE_TRITON_MOE"), "1") + self.assertEqual(variant.roles.server.env.get("ATOM_USE_TRITON_GEMM"), "1") + + def test_load_phase_c_w3_glm_perf(self): + root = _PACKAGE_ROOT + config = root / "input/config_file/inference/atom/mi300x_atom_glm-5.1_single.json" + variant = load_variant(config, _cluster_dict()) + self.assertEqual(variant.model.id, "zai-org/GLM-5.1") + self.assertEqual( + variant.expected_cells(), + ["ISL=1024,OSL=8192,TP=8,CONC=32", "ISL=1024,OSL=8192,TP=8,CONC=64"], + ) + + def test_load_phase_c_w13_code_perf(self): + root = _PACKAGE_ROOT + config = root / "input/config_file/inference/atom/mi355x_atom_kimi-k2.7-code_single.json" + variant = load_variant(config, _cluster_dict()) + self.assertEqual(variant.model.id, "moonshotai/Kimi-K2.7-Code") + self.assertFalse(variant.enforce_thresholds) + + def test_load_phase_c_w17_mxfp4_perf(self): + root = _PACKAGE_ROOT + config = root / "input/config_file/inference/atom/mi300x_atom_deepseek-r1_mxfp4_single.json" + variant = load_variant(config, _cluster_dict()) + self.assertEqual(variant.model.id, "amd/DeepSeek-R1-0528-MXFP4") + self.assertEqual(variant.params.tensor_parallelism, "8") + + def test_load_m4_vllm_single_parity(self): + root = _PACKAGE_ROOT + config = root / "input/config_file/inference/atom/mi300x_atom_deepseek-r1_fp8_vllm_single.json" + variant = load_variant(config, _cluster_dict()) + self.assertEqual(variant.params.driver, "vllm_atom") + self.assertEqual(variant.params.nnodes, "1") + self.assertIn("kv-cache-dtype", variant.roles.server.serve_args) + + def test_load_m4_sglang_single_parity(self): + root = _PACKAGE_ROOT + config = root / "input/config_file/inference/atom/mi300x_atom_deepseek-r1_fp8_sglang_single.json" + variant = load_variant(config, _cluster_dict()) + self.assertEqual(variant.params.driver, "sglang") + self.assertIn("--kv-cache-dtype", variant.roles.server.sglang_args) + + def test_load_qwen397b_fp8_single(self): + root = _PACKAGE_ROOT + config = root / "input/config_file/inference/atom/mi300x_atom_qwen3.5-397b-a17b_fp8_single.json" + variant = load_variant(config, _cluster_dict()) + self.assertEqual(variant.model.id, "amd/Qwen3.5-397B-A17B-FP8") + self.assertEqual(variant.expected_cells()[0], "ISL=1024,OSL=8192,TP=8,CONC=32") + + def test_load_kimi_k26_thinking_single_tp4(self): + root = _PACKAGE_ROOT + config = root / "input/config_file/inference/atom/mi355x_atom_kimi-k2.6-thinking_single.json" + variant = load_variant(config, _cluster_dict()) + self.assertEqual(variant.params.tensor_parallelism, "4") + self.assertIn("TP=4", variant.expected_cells()[0]) + + def test_load_kimi_k27_longctx_single(self): + root = _PACKAGE_ROOT + config = root / "input/config_file/inference/atom/mi355x_atom_kimi-k2.7-code_longctx_single.json" + variant = load_variant(config, _cluster_dict()) + self.assertEqual(variant.expected_cells()[0], "ISL=8192,OSL=1024,TP=8,CONC=32") + + def test_load_w1_single_gpu_metrics_poll(self): + root = _PACKAGE_ROOT + config = root / "input/config_file/inference/atom/mi300x_atom_deepseek-r1_fp8_single.json" + variant = load_variant(config, _cluster_dict()) + self.assertTrue(variant.platform.gpu_metrics_poll) + + def test_w1_gsm8k_threshold_fails_below_floor(self): + from cvs.lib.utils.verdict import ThresholdViolation, evaluate_all + + root = _PACKAGE_ROOT + config = root / "input/config_file/inference/atom/mi300x_atom_deepseek-r1_fp8_accuracy.json" + variant = load_variant(config, _cluster_dict()) + specs = variant.thresholds["accuracy"]["gsm8k_flex"] + with self.assertRaises(ThresholdViolation): + evaluate_all({"gsm8k.exact_match__flexible-extract": 0.90}, specs) + + def test_load_distributed_accuracy_scaffold(self): + root = _PACKAGE_ROOT + config = root / "input/config_file/inference/atom/mi300x_atom_deepseek-r1_fp8_distributed_accuracy.json" + variant = load_variant(config, _cluster_dict()) + self.assertEqual(variant.params.driver, "vllm_atom") + self.assertEqual(variant.params.nnodes, "2") + self.assertIn("PP=2", variant.expected_cells()[0]) + self.assertIn("accuracy", variant.thresholds) + + def test_load_v4_pro_longctx_stem(self): + root = _PACKAGE_ROOT + config = root / "input/config_file/inference/atom/mi300x_atom_deepseek-v4-pro_longctx_single.json" + variant = load_variant(config, _cluster_dict()) + self.assertEqual(variant.model.id, "deepseek-ai/DeepSeek-V4-Pro") + self.assertEqual(variant.expected_cells()[0], "ISL=5000,OSL=1024,TP=8,CONC=16") + self.assertTrue(variant.platform.gpu_metrics_poll) + + def test_load_v4_pro_vllm_single_stem(self): + root = _PACKAGE_ROOT + config = root / "input/config_file/inference/atom/mi300x_atom_deepseek-v4-pro_vllm_single.json" + variant = load_variant(config, _cluster_dict()) + self.assertEqual(variant.params.driver, "vllm_atom") + self.assertEqual(variant.params.tokenizer_mode, "deepseek_v4") + self.assertFalse(variant.roles.server.serve_args.get("enforce-eager")) + self.assertEqual(variant.roles.server.serve_args.get("moe-backend"), "triton_unfused") + + def test_load_v4_pro_sglang_single_stem(self): + root = _PACKAGE_ROOT + config = root / "input/config_file/inference/atom/mi300x_atom_deepseek-v4-pro_sglang_single.json" + variant = load_variant(config, _cluster_dict()) + self.assertEqual(variant.params.driver, "sglang") + self.assertEqual(variant.model.id, "deepseek-ai/DeepSeek-V4-Pro") + + def test_load_v4_pro_distributed_stem(self): + root = _PACKAGE_ROOT + config = root / "input/config_file/inference/atom/mi300x_atom_deepseek-v4-pro_distributed.json" + variant = load_variant(config, _cluster_dict()) + self.assertEqual(variant.params.driver, "vllm_atom") + self.assertEqual(variant.params.pipeline_parallel_size, "2") + self.assertIn("PP=2", variant.expected_cells()[0]) + + def test_load_w2_m4_vllm_parity(self): + root = _PACKAGE_ROOT + config = root / "input/config_file/inference/atom/mi300x_atom_gpt-oss-120b_mxfp4_vllm_single.json" + variant = load_variant(config, _cluster_dict()) + self.assertEqual(variant.params.driver, "vllm_atom") + self.assertEqual(variant.model.id, "openai/gpt-oss-120b") + + +class TestAtomVariantSamples(unittest.TestCase): + def test_all_committed_variant_samples_validate(self): + config_dir = _PACKAGE_ROOT / "input" / "config_file" / "inference" / "atom" + for path in sorted(config_dir.glob("*.json")): + if path.name.endswith("_threshold.json"): + continue + raw = json.loads(path.read_text()) + if raw.get("framework") != "atom": + continue + with self.subTest(sample=path.name): + known = {k: v for k, v in raw.items() if k in AtomVariantConfig.model_fields} + known["enforce_thresholds"] = False + known["thresholds"] = {} + AtomVariantConfig.model_validate(known) + + +if __name__ == "__main__": + unittest.main() diff --git a/cvs/schema/config_file/inference/atom/variant.py b/cvs/schema/config_file/inference/atom/variant.py new file mode 100644 index 000000000..2f792d2c5 --- /dev/null +++ b/cvs/schema/config_file/inference/atom/variant.py @@ -0,0 +1,244 @@ +""" +ATOM inference variant config schema. + +Mirrors ``cvs/input/config_file/inference/atom/``. +""" + +from __future__ import annotations + +import re +import warnings + +from pydantic import Field, field_validator, model_validator +from typing_extensions import Literal + +from cvs.lib import globals +from cvs.lib.inference.atom.atom_parsing import GATED_METRICS +from cvs.schema.common.base import BaseVariantConfig +from cvs.schema.base import _Forbid +from cvs.schema.config_file.inference.common.accuracy import AccuracyConfig +from cvs.schema.config_file.inference.common.functional import FunctionalConfig +from cvs.schema.config_file.inference.common.long_context_accuracy import LongContextAccuracyConfig +from cvs.schema.config_file.inference.common.platform import PlatformConfig +from cvs.schema.config_file.inference.common.sweep import ( + RoleServer, + Sweep, + validate_thresholds_cover_sweep, +) + +ATOM_DRIVERS = ("atom", "vllm", "vllm_atom", "sglang") +ATOM_PP_DRIVERS = ("vllm", "vllm_atom", "sglang") +_MXFP4_TRITON_ENV = { + "ATOM_USE_TRITON_MOE": "1", + "ATOM_USE_TRITON_GEMM": "1", +} + +log = globals.log +_ORCH_MANAGED_NETWORK_ENV = frozenset({"NCCL_SOCKET_IFNAME", "GLOO_SOCKET_IFNAME", "TP_SOCKET_IFNAME", "NCCL_IB_HCA"}) +_IB_HCA_NETDEV_RE = re.compile(r"^mlx5_\d+$", re.IGNORECASE) + + +def merge_mxfp4_triton_env(precision: str, env: dict[str, str]) -> dict[str, str]: + """Return server env with MXFP4 Triton defaults applied when unset.""" + merged = dict(env or {}) + if (precision or "").lower() == "mxfp4": + for key, value in _MXFP4_TRITON_ENV.items(): + merged.setdefault(key, value) + for key in _MXFP4_TRITON_ENV: + if str(merged.get(key, "")).lower() == "true": + merged[key] = "1" + return merged + + +class AtomRoleServer(RoleServer): + atom_args: list[str] = [] + sglang_args: list[str] = [] + ib_hca_devices: Literal["auto"] | list[str] | None = None + ib_netdev: Literal["auto"] | str | None = None + + @field_validator("ib_netdev", mode="after") + @classmethod + def _normalize_ib_netdev(cls, v): + raw = (v or "").strip() + if raw and raw.lower() != "auto" and _IB_HCA_NETDEV_RE.match(raw): + log.warning( + "roles.server.ib_netdev=%r looks like an IB HCA name; coercing to 'auto' " + "(socket netdev is discovered from cluster IPs at runtime)", + raw, + ) + return "auto" + return v + + @model_validator(mode="after") + def _strip_orchestrator_managed_network_env(self): + if not self.env: + return self + dropped = sorted(k for k in self.env if k in _ORCH_MANAGED_NETWORK_ENV) + if not dropped: + return self + log.warning( + "roles.server.env drops orchestrator-managed keys %s " + "(set by test_discover_topology / build_server_cmd instead)", + dropped, + ) + self.env = {k: v for k, v in self.env.items() if k not in _ORCH_MANAGED_NETWORK_ENV} + return self + + +class AtomRoles(_Forbid): + server: AtomRoleServer = Field(default_factory=AtomRoleServer) + + +class AtomParams(_Forbid): + driver: Literal["atom", "vllm", "vllm_atom", "sglang"] = "vllm" + backend: str = "vllm" + base_url: str = "http://0.0.0.0" + port_no: str = "8000" + dataset_name: str = "random" + burstiness: str = "1.0" + seed: str = "0" + request_rate: str = "inf" + random_range_ratio: str = "0.8" + random_prefix_len: str = "0" + tensor_parallelism: str = "8" + tokenizer_mode: str = "auto" + percentile_metrics: str = "ttft,tpot,itl,e2el" + metric_percentiles: str = "95,99" + num_prompts: str = "1000" + max_model_length: str = "8192" + client_poll_count: str = "50" + client_poll_wait_time: str = "60" + client_initial_wait_s: str = "120" + server_precheck_wait_s: str = "30" + server_warmup_wait_s: str = "330" + server_poll_count: str = "60" + server_poll_wait_time: str = "60" + reuse_server_across_sweep: str = "false" + bench_max_failed_requests: str = "0" + bench_extra_args: str = "" + result_filename: str = "results" + nnodes: str = "1" + pipeline_parallel_size: str = "1" + master_addr: str = "" + master_port: str = "29501" + scaling_baseline_output_throughput: str = "" + + +class AtomRunCard(_Forbid): + upstream_run_url: str = "" + atom_image_pin: str = "" + notes: str = "" + + +class MtpQualityConfig(_Forbid): + enabled: bool = False + chat_template_prompt: str = "Say hello in one short sentence." + chat_template_expected_sha256: str = "" + + +class QuantParityConfig(_Forbid): + enabled: bool = False + probe_prompt: str = "The capital of France is" + reference_config_stem: str = "" + + +class AtomVariantConfig(BaseVariantConfig): + framework: Literal["atom"] + gpu_arch: str + run_card: AtomRunCard = Field(default_factory=AtomRunCard) + roles: AtomRoles = Field(default_factory=AtomRoles) + params: AtomParams + sweep: Sweep + accuracy: AccuracyConfig = Field(default_factory=AccuracyConfig) + mtp_quality: MtpQualityConfig = Field(default_factory=MtpQualityConfig) + quant_parity: QuantParityConfig = Field(default_factory=QuantParityConfig) + functional: FunctionalConfig = Field(default_factory=FunctionalConfig) + long_context_accuracy: LongContextAccuracyConfig = Field(default_factory=LongContextAccuracyConfig) + platform: PlatformConfig = Field(default_factory=PlatformConfig) + + def cell_key(self, isl, osl, concurrency): + p = self.params + key = f"ISL={isl},OSL={osl},TP={p.tensor_parallelism}" + nnodes = int(p.nnodes) + pp = int(p.pipeline_parallel_size) + if p.driver == "atom": + if nnodes > 1: + key += f",DP={nnodes},NNODES={nnodes}" + elif p.driver in ATOM_PP_DRIVERS: + if pp > 1 or nnodes > 1: + key += f",PP={p.pipeline_parallel_size}" + if nnodes > 1: + key += f",NNODES={p.nnodes}" + return f"{key},CONC={concurrency}" + + def expected_cells(self) -> list[str]: + by_name = {c.name: c for c in self.sweep.sequence_combinations} + return [self.cell_key(by_name[r.combo].isl, by_name[r.combo].osl, r.concurrency) for r in self.sweep.runs] + + @model_validator(mode="after") + def _apply_mxfp4_triton_env_defaults(self): + self.roles.server.env = merge_mxfp4_triton_env(self.model.precision, self.roles.server.env) + return self + + @model_validator(mode="after") + def _check_thresholds_cover_sweep(self): + validate_thresholds_cover_sweep( + expected_cells=self.expected_cells(), + thresholds=self.thresholds, + enforce_thresholds=self.enforce_thresholds, + gated_metrics=GATED_METRICS, + ) + if int(self.params.nnodes) > 1 and (self.params.scaling_baseline_output_throughput or "").strip(): + missing = [] + for cell in self.expected_cells(): + specs = self.thresholds.get(cell) or {} + if "scaling.efficiency_pct" not in specs: + missing.append(cell) + if missing: + msg = ( + "multinode variant with scaling_baseline_output_throughput requires " + f"scaling.efficiency_pct in every cell; missing: {missing}" + ) + if self.enforce_thresholds: + raise ValueError(msg) + warnings.warn(f"{msg} (enforce_thresholds=false -> record-only)", stacklevel=2) + return self + + @model_validator(mode="after") + def _atom_multinode_uses_dp_not_pp(self): + if self.params.driver == "atom" and int(self.params.nnodes) > 1: + if int(self.params.pipeline_parallel_size) > 1: + raise ValueError( + "params.driver='atom' with nnodes>1 uses ATOM SPMD data parallel (-dp); " + "standalone ATOM cannot execute pipeline parallel. For true PP>1 use " + "params.driver='vllm_atom' or 'sglang'." + ) + return self + + @model_validator(mode="after") + def _pp_driver_distributed_consistency(self): + driver = self.params.driver + if driver not in ATOM_PP_DRIVERS: + return self + nn = int(self.params.nnodes) + pp = int(self.params.pipeline_parallel_size) + is_ray = self.roles.server.serve_args.get("distributed-executor-backend") == "ray" + if nn > 1 and pp == 1 and not is_ray: + raise ValueError( + f"params.driver={driver!r} with nnodes={nn} requires pipeline_parallel_size>1 " + f"(got pp={pp}) for multinode pipeline parallel" + ) + if pp > 1 and nn == 1: + raise ValueError( + f"pipeline_parallel_size={pp} > 1 requires nnodes > 1 (got nnodes={nn}) for params.driver={driver!r}" + ) + return self + + @model_validator(mode="after") + def _atom_driver_requires_inline_server_args(self): + if self.params.driver == "atom" and not self.roles.server.atom_args: + raise ValueError( + "params.driver='atom' requires roles.server.atom_args " + "(inline ATOM openai_server CLI tokens, vLLM-style)" + ) + return self diff --git a/cvs/schema/config_file/inference/common/__init__.py b/cvs/schema/config_file/inference/common/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/cvs/schema/config_file/inference/common/accuracy.py b/cvs/schema/config_file/inference/common/accuracy.py new file mode 100644 index 000000000..1914d7cf2 --- /dev/null +++ b/cvs/schema/config_file/inference/common/accuracy.py @@ -0,0 +1,34 @@ +"""Accuracy-evaluation config schema shared across inference suites.""" + +from __future__ import annotations + +from collections import Counter +from typing import Any, Dict, List + +from pydantic import model_validator + +from cvs.schema.base import _Forbid + + +class AccuracyTask(_Forbid): + id: str + task: str + num_fewshot: int = 0 + metadata: Dict[str, Any] = {} + include_path: str = "" + num_concurrent: int = 8 + apply_chat_template: bool = False + gen_kwargs: Dict[str, Any] = {} + + +class AccuracyConfig(_Forbid): + tasks: List[AccuracyTask] = [] + + @model_validator(mode="after") + def _check_unique_task_ids(self): + counts = Counter(t.id for t in self.tasks) + dupes = sorted(i for i, n in counts.items() if n > 1) + if dupes: + rendered = ", ".join(repr(d) for d in dupes) + raise ValueError(f"duplicate task id(s): {rendered}") + return self diff --git a/cvs/schema/config_file/inference/common/functional.py b/cvs/schema/config_file/inference/common/functional.py new file mode 100644 index 000000000..adf2454c1 --- /dev/null +++ b/cvs/schema/config_file/inference/common/functional.py @@ -0,0 +1,10 @@ +"""Optional functional smoke gates for inference suites.""" + +from __future__ import annotations + +from cvs.schema.base import _Forbid + + +class FunctionalConfig(_Forbid): + api_smoke: bool = False + health_check: bool = False diff --git a/cvs/schema/config_file/inference/common/long_context_accuracy.py b/cvs/schema/config_file/inference/common/long_context_accuracy.py new file mode 100644 index 000000000..a94416069 --- /dev/null +++ b/cvs/schema/config_file/inference/common/long_context_accuracy.py @@ -0,0 +1,31 @@ +"""Long-context accuracy cell selection for inference suites.""" + +from __future__ import annotations + +from collections import Counter +from typing import List + +from pydantic import model_validator + +from cvs.schema.base import _Forbid + + +class LongContextAccCell(_Forbid): + id: str + isl: int + osl: int = 32 + num_prompts: int = 8 + seed: int = 42 + + +class LongContextAccuracyConfig(_Forbid): + cells: List[LongContextAccCell] = [] + + @model_validator(mode="after") + def _check_unique_cell_ids(self): + counts = Counter(c.id for c in self.cells) + dupes = sorted(i for i, n in counts.items() if n > 1) + if dupes: + rendered = ", ".join(repr(d) for d in dupes) + raise ValueError(f"duplicate long-context cell id(s): {rendered}") + return self diff --git a/cvs/schema/config_file/inference/common/platform.py b/cvs/schema/config_file/inference/common/platform.py new file mode 100644 index 000000000..8e621e5b1 --- /dev/null +++ b/cvs/schema/config_file/inference/common/platform.py @@ -0,0 +1,10 @@ +"""Optional platform / post-run checks for inference suites.""" + +from __future__ import annotations + +from cvs.schema.base import _Forbid + + +class PlatformConfig(_Forbid): + dmesg_scan: bool = False + gpu_metrics_poll: bool = False diff --git a/cvs/schema/config_file/inference/common/sweep.py b/cvs/schema/config_file/inference/common/sweep.py new file mode 100644 index 000000000..8906b1c52 --- /dev/null +++ b/cvs/schema/config_file/inference/common/sweep.py @@ -0,0 +1,104 @@ +""" +Shared inference sweep selector types and threshold coverage helpers. +""" + +from __future__ import annotations + +import warnings +from collections import Counter +from typing import Any, Dict, List, Optional + +from pydantic import model_validator + +from cvs.schema.base import _Forbid + +NON_SWEEP_THRESHOLD_KEYS = {"accuracy", "mtp_quality", "long_context_accuracy", "quant_parity"} + + +class RoleServer(_Forbid): + serve_args: Dict[str, Any] = {} + env: Dict[str, str] = {} + + +class Roles(_Forbid): + server: RoleServer = RoleServer() + + +class GoodputSlo(_Forbid): + ttft_ms: float + tpot_ms: float + e2el_ms: float + + +class SeqCombo(_Forbid): + name: str + isl: str + osl: str + goodput_slo: Optional[GoodputSlo] = None + + +class Run(_Forbid): + combo: str + concurrency: int + + +def validate_thresholds_cover_sweep( + *, + expected_cells, + thresholds, + enforce_thresholds: bool, + gated_metrics=None, + gated_gpu_metrics=None, +) -> None: + """Shared sweep/threshold coverage check for inference variant configs.""" + expected = set(expected_cells) + present = set(thresholds.keys()) - NON_SWEEP_THRESHOLD_KEYS + missing = sorted(expected - present) + extra = sorted(present - expected) + problems = [] + if missing: + problems.append(f"sweep cells with no threshold entry: {missing}") + if extra: + problems.append(f"threshold keys matching no sweep cell (typo?): {extra}") + gated = gated_metrics if gated_metrics is not None else set() + gated_keys = [f"client.{m}" for m in sorted(gated)] + if gated_gpu_metrics: + gated_keys += [f"gpu.{m}" for m in sorted(gated_gpu_metrics)] + gated_gaps = {} + for cell in sorted(expected & present): + specs = thresholds.get(cell) or {} + absent = [k for k in gated_keys if k not in specs] + if absent: + gated_gaps[cell] = absent + if gated_gaps: + problems.append(f"cells missing gated-metric specs: {gated_gaps}") + if problems: + msg = "threshold.json does not match the sweep matrix; " + "; ".join(problems) + if enforce_thresholds: + raise ValueError(msg) + warnings.warn(f"{msg} (enforce_thresholds=false -> record-only)", stacklevel=3) + + +def validate_sweep_selector(combo_names, run_combo_refs): + """Sweep-selector rule: combo names unique, every run.combo names one.""" + counts = Counter(combo_names) + dupes = sorted(name for name, count in counts.items() if count > 1) + if dupes: + raise ValueError(f"duplicate sequence_combination names: {dupes}") + known = set(counts) + unknown = sorted({r for r in run_combo_refs if r not in known}) + if unknown: + raise ValueError(f"run.combo names no sequence_combination: {unknown} (known: {sorted(known)})") + + +class Sweep(_Forbid): + sequence_combinations: List[SeqCombo] + runs: List[Run] + + @model_validator(mode="after") + def _check_runs_reference_known_combos(self): + validate_sweep_selector( + [c.name for c in self.sequence_combinations], + [r.combo for r in self.runs], + ) + return self diff --git a/cvs/schema/config_file/inference/common/unittests/__init__.py b/cvs/schema/config_file/inference/common/unittests/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/cvs/schema/config_file/inference/common/unittests/test_accuracy.py b/cvs/schema/config_file/inference/common/unittests/test_accuracy.py new file mode 100644 index 000000000..e329ff0db --- /dev/null +++ b/cvs/schema/config_file/inference/common/unittests/test_accuracy.py @@ -0,0 +1,496 @@ +"""Unit tests for AccuracyConfig / AccuracyTask (inference/common/accuracy.py).""" + +import unittest + +from pydantic import ValidationError + +from cvs.schema.config_file.inference.common.accuracy import AccuracyConfig, AccuracyTask + + +def _task(**over): + """A valid AccuracyTask with fields overridable per case.""" + base = {"id": "a", "task": "gsm8k"} + base.update(over) + return AccuracyTask(**base) + + +def _dupes_message(exc): + """Isolate the duplicate-id validator's own message from pydantic's wrapper. + + str(ValidationError) appends '[type=..., input_value=..., input_type=...]', + and input_value repeats each task (and thus each id). To assert on the + validator's rendered sorted-dupes list (count/order/quoting) without the + wrapper's echoes, slice from the documented prefix up to pydantic's + '[type=' metadata marker. + """ + msg = str(exc) + prefix = "duplicate task id(s):" + i = msg.find(prefix) + if i == -1: + return None + tail = msg[i:] + j = tail.find("[type=") + if j != -1: + tail = tail[:j] + return tail + + +class TestAccuracyTaskDefaults(unittest.TestCase): + """AC12, AC23, AC24: defaults, empty-string id, include_path (no I/O).""" + + def test_defaults_on_minimal_task(self): + t = AccuracyTask(id="a", task="gsm8k") + self.assertEqual(t.id, "a") + self.assertEqual(t.task, "gsm8k") + self.assertEqual(t.num_fewshot, 0) + self.assertEqual(t.metadata, {}) + self.assertEqual(t.include_path, "") + self.assertEqual(t.num_concurrent, 8) + self.assertIs(t.apply_chat_template, False) + self.assertEqual(t.gen_kwargs, {}) + + def test_empty_string_id_is_valid(self): + # AC23: "" is a valid id, not treated as missing. + t = AccuracyTask(id="", task="gsm8k") + self.assertEqual(t.id, "") + + def test_include_path_no_filesystem_check(self): + # AC24: nonexistent path accepted as plain string, no I/O. + t = AccuracyTask(id="a", task="gsm8k", include_path="/some/nonexistent/path") + self.assertEqual(t.include_path, "/some/nonexistent/path") + + def test_default_dicts_are_isolated_per_instance(self): + # Mutable-default isolation for metadata/gen_kwargs on AccuracyTask. + t1 = AccuracyTask(id="a", task="gsm8k") + t2 = AccuracyTask(id="b", task="gsm8k") + self.assertIsNot(t1.metadata, t2.metadata) + self.assertIsNot(t1.gen_kwargs, t2.gen_kwargs) + + +class TestAccuracyTaskRequiredFields(unittest.TestCase): + """AC13: id and task are required.""" + + def test_missing_required_field_raises(self): + for missing in ("id", "task"): + with self.subTest(missing=missing): + kwargs = {"id": "a", "task": "gsm8k"} + del kwargs[missing] + with self.assertRaises(ValidationError): + AccuracyTask(**kwargs) + + +class TestAccuracyTaskExplicitNone(unittest.TestCase): + """Explicit None is a distinct equivalence class from omission: no field is + Optional, so None is rejected for every field (required str fields AND the + non-None-defaulted int/dict/bool/str fields). Mirror of AC28 for the config's + tasks field. Guards against a mutated schema (e.g. id: Optional[str], or + metadata: Optional[Dict] = {}) silently accepting None while still passing + the omission-only required-field test.""" + + def test_explicit_none_per_field_raises(self): + # (field, None value passed via a valid base task) + for field in ( + "id", + "task", + "num_fewshot", + "metadata", + "include_path", + "num_concurrent", + "apply_chat_template", + "gen_kwargs", + ): + with self.subTest(field=field): + with self.assertRaises(ValidationError): + _task(**{field: None}) + + +class TestAccuracyTaskIntCoercion(unittest.TestCase): + """AC14, AC15, AC16: int fields coerce numeric strings; no range constraint.""" + + def test_int_coercion_success(self): + # (field, input, expected int) + cases = [ + ("num_fewshot", "5", 5), + ("num_fewshot", 5, 5), + ("num_fewshot", 5.0, 5), # whole-number float coerces (vs 1.9 which rejects) + ("num_fewshot", -1, -1), # AC16: negative allowed, no ge + ("num_fewshot", 0, 0), + ("num_fewshot", True, 1), # pydantic lax int: bool coerces to 1/0 + ("num_concurrent", "3", 3), + ("num_concurrent", 3, 3), + ("num_concurrent", 4.0, 4), # whole-number float coerces (vs 2.5 which rejects) + ("num_concurrent", 0, 0), # AC16: zero allowed, no gt + ("num_concurrent", -1, -1), + ("num_concurrent", False, 0), # pydantic lax int: bool coerces to 1/0 + ] + for field, value, expected in cases: + with self.subTest(field=field, value=value): + t = _task(**{field: value}) + got = getattr(t, field) + self.assertEqual(got, expected) + self.assertIsInstance(got, int) + + def test_int_coercion_failure_raises(self): + cases = [ + ("num_fewshot", "not-an-int"), + ("num_fewshot", 1.9), # float with fractional part + ("num_concurrent", "bad"), + ("num_concurrent", 2.5), + ] + for field, value in cases: + with self.subTest(field=field, value=value): + with self.assertRaises(ValidationError): + _task(**{field: value}) + + +class TestAccuracyTaskDictCoercion(unittest.TestCase): + """AC17, AC18, AC19: dict fields accept mappings only.""" + + def test_dict_success(self): + cases = [ + ("metadata", {"k": "v"}), + ("metadata", {}), + ("gen_kwargs", {"a": 1}), + ("gen_kwargs", {}), + ] + for field, value in cases: + with self.subTest(field=field, value=value): + t = _task(**{field: value}) + self.assertEqual(getattr(t, field), value) + + def test_non_mapping_raises(self): + cases = [ + ("metadata", "not-a-dict"), + ("metadata", [1, 2]), + ("metadata", 123), + ("gen_kwargs", 123), + ("gen_kwargs", "not-a-dict"), + ("gen_kwargs", [1, 2]), + ] + for field, value in cases: + with self.subTest(field=field, value=value): + with self.assertRaises(ValidationError): + _task(**{field: value}) + + +class TestAccuracyTaskBoolCoercion(unittest.TestCase): + """AC20, AC21: bool field; 'maybe' is the guaranteed-fail non-bool word.""" + + def test_bool_true_accepted(self): + t = _task(apply_chat_template=True) + self.assertIs(t.apply_chat_template, True) + + def test_bool_false_accepted(self): + t = _task(apply_chat_template=False) + self.assertIs(t.apply_chat_template, False) + + def test_non_bool_word_raises(self): + with self.assertRaises(ValidationError): + _task(apply_chat_template="maybe") + + +class TestAccuracyTaskStringTyping(unittest.TestCase): + """AC22: id/task accept only str; non-str scalars are not auto-coerced.""" + + def test_non_str_id_or_task_raises(self): + cases = [ + {"id": 123, "task": "gsm8k"}, + {"id": "a", "task": 123}, + {"id": 1.5, "task": "gsm8k"}, + {"id": True, "task": "gsm8k"}, # bool is a non-str scalar; not coerced to str + {"id": "a", "task": False}, # bool is a non-str scalar; not coerced to str + ] + for kwargs in cases: + with self.subTest(kwargs=kwargs): + with self.assertRaises(ValidationError): + AccuracyTask(**kwargs) + + +class TestAccuracyTaskExtraForbid(unittest.TestCase): + """AC10: unknown fields rejected (extra='forbid' from _Forbid).""" + + def test_unknown_field_raises(self): + with self.assertRaises(ValidationError): + AccuracyTask(id="a", task="gsm8k", extra_field=1) + + +class TestAccuracyConfigConstruction(unittest.TestCase): + """AC1, AC2, AC3, AC25, AC26, AC29: happy-path construction + element typing.""" + + def test_empty_config_has_empty_tasks(self): + # AC1 + edge case: zero tasks constructs, tasks == []. + cfg = AccuracyConfig() + self.assertEqual(cfg.tasks, []) + + def test_single_task_constructs(self): + # AC2 + edge case: exactly one task is trivially unique. + cfg = AccuracyConfig(tasks=[AccuracyTask(id="a", task="gsm8k")]) + self.assertEqual(len(cfg.tasks), 1) + self.assertEqual(cfg.tasks[0].id, "a") + + def test_three_distinct_ids_construct(self): + # AC3. + cfg = AccuracyConfig( + tasks=[ + AccuracyTask(id="a", task="gsm8k"), + AccuracyTask(id="b", task="gsm8k"), + AccuracyTask(id="c", task="gsm8k"), + ] + ) + self.assertEqual([t.id for t in cfg.tasks], ["a", "b", "c"]) + + def test_list_of_dicts_becomes_tasks(self): + # AC25. + cfg = AccuracyConfig(tasks=[{"id": "a", "task": "gsm8k"}]) + self.assertIsInstance(cfg.tasks[0], AccuracyTask) + self.assertEqual(cfg.tasks[0].id, "a") + + def test_mixed_dicts_and_instances(self): + # AC26: every element ends up an AccuracyTask. + cfg = AccuracyConfig(tasks=[AccuracyTask(id="a", task="gsm8k"), {"id": "b", "task": "gsm8k"}]) + self.assertTrue(all(isinstance(t, AccuracyTask) for t in cfg.tasks)) + self.assertEqual([t.id for t in cfg.tasks], ["a", "b"]) + + def test_order_and_length_preserved(self): + # AC29. + cfg = AccuracyConfig( + tasks=[ + AccuracyTask(id="a", task="m"), + AccuracyTask(id="b", task="m"), + AccuracyTask(id="c", task="m"), + ] + ) + self.assertEqual([t.id for t in cfg.tasks], ["a", "b", "c"]) + + +class TestAccuracyConfigTasksField(unittest.TestCase): + """AC11, AC27, AC28: extra forbid + tasks element/None handling.""" + + def test_unknown_field_raises(self): + # AC11. + with self.assertRaises(ValidationError): + AccuracyConfig(tasks=[], extra_field=1) + + def test_non_dict_non_instance_element_raises(self): + # AC27. + with self.assertRaises(ValidationError): + AccuracyConfig(tasks=["not-a-task"]) + + def test_tasks_none_raises(self): + # AC28: field is not Optional; only omission yields the [] default. + with self.assertRaises(ValidationError): + AccuracyConfig(tasks=None) + + +class TestAccuracyConfigDuplicateIds(unittest.TestCase): + """AC4-AC9: the model_validator(mode='after') duplicate-id contract.""" + + def test_single_duplicate_group(self): + # AC4: prefix + the offending id present. + with self.assertRaises(ValidationError) as ctx: + AccuracyConfig( + tasks=[ + AccuracyTask(id="dup-mmlu", task="mmlu"), + AccuracyTask(id="dup-mmlu", task="mmlu"), + ] + ) + dupes = _dupes_message(ctx.exception) + self.assertIsNotNone(dupes, "expected 'duplicate task id(s):' prefix") + self.assertIn("dup-mmlu", dupes) + + def test_two_groups_sorted_ascending(self): + # AC5: both ids present; 'dup-gsm8k' before 'dup-mmlu' (sorted, not + # encounter order -- input intentionally lists mmlu first). + with self.assertRaises(ValidationError) as ctx: + AccuracyConfig( + tasks=[ + AccuracyTask(id="dup-mmlu", task="mmlu"), + AccuracyTask(id="dup-gsm8k", task="gsm8k"), + AccuracyTask(id="dup-mmlu", task="mmlu"), + AccuracyTask(id="dup-gsm8k", task="gsm8k"), + ] + ) + dupes = _dupes_message(ctx.exception) + self.assertIsNotNone(dupes) + i_gsm = dupes.find("dup-gsm8k") + i_mmlu = dupes.find("dup-mmlu") + self.assertNotEqual(i_gsm, -1) + self.assertNotEqual(i_mmlu, -1) + self.assertLess(i_gsm, i_mmlu, "ids must be sorted ascending in the message") + + def test_mixed_case_dupes_sorted_case_sensitively(self): + # AC5 (sort discriminator): the sort must be case-SENSITIVE lexicographic, + # distinct from AC8's case-sensitive equality. All-lowercase fixtures + # (dup-gsm8k/dup-mmlu) cannot tell a correct sorted(dupes) from a + # spec-violating sorted(dupes, key=str.lower). Use ids that differ in + # leading case: case-sensitive sort orders uppercase before lowercase + # ('Dup-Zebra' < 'dup-apple'), whereas a case-insensitive key flips them + # ('dup-apple' < 'Dup-Zebra' since 'a' < 'z'). + with self.assertRaises(ValidationError) as ctx: + AccuracyConfig( + tasks=[ + AccuracyTask(id="dup-apple", task="mmlu"), + AccuracyTask(id="Dup-Zebra", task="gsm8k"), + AccuracyTask(id="dup-apple", task="mmlu"), + AccuracyTask(id="Dup-Zebra", task="gsm8k"), + ] + ) + dupes = _dupes_message(ctx.exception) + self.assertIsNotNone(dupes) + i_zebra = dupes.find("Dup-Zebra") + i_apple = dupes.find("dup-apple") + self.assertNotEqual(i_zebra, -1) + self.assertNotEqual(i_apple, -1) + self.assertLess( + i_zebra, + i_apple, + "dupes must be sorted case-sensitively: uppercase-leading 'Dup-Zebra' " + "precedes 'dup-apple' (a case-insensitive sort key would reverse this)", + ) + + def test_triple_duplicate_id_listed_once(self): + # AC6: id repeated 3x appears exactly once in the sorted-dupes list. + with self.assertRaises(ValidationError) as ctx: + AccuracyConfig( + tasks=[ + AccuracyTask(id="dup-mmlu", task="mmlu"), + AccuracyTask(id="dup-mmlu", task="mmlu"), + AccuracyTask(id="dup-mmlu", task="mmlu"), + ] + ) + dupes = _dupes_message(ctx.exception) + self.assertIsNotNone(dupes) + self.assertEqual(dupes.count("dup-mmlu"), 1) + + def test_duplicate_by_id_only_ignores_other_fields(self): + # AC7: same id, different task -> still duplicates. The raise itself is + # the discriminator: full-object comparison would construct successfully. + # Use a distinctive id ("dup-x") that cannot be a substring of the fixed + # "duplicate task id(s):" prefix, so assertIn actually probes the + # validator's rendered dupes list rather than the constant prefix text. + with self.assertRaises(ValidationError) as ctx: + AccuracyConfig( + tasks=[ + AccuracyTask(id="dup-x", task="mmlu"), + AccuracyTask(id="dup-x", task="gsm8k"), + ] + ) + dupes = _dupes_message(ctx.exception) + self.assertIsNotNone(dupes, "must be the duplicate-id validator, not another error") + self.assertIn("dup-x", dupes) + + def test_only_repeated_ids_listed_not_unique_ones(self): + # AC4 (message contents): the dupes list must contain ONLY ids that + # actually repeat, not every distinct id in the config. Mix a duplicated + # id with an id that appears exactly once and assert the unique one is + # absent -- this fails a validator that reports sorted(set(all_ids)). + with self.assertRaises(ValidationError) as ctx: + AccuracyConfig( + tasks=[ + AccuracyTask(id="dup-repeated", task="mmlu"), + AccuracyTask(id="dup-repeated", task="gsm8k"), + AccuracyTask(id="only-once", task="mmlu"), + ] + ) + dupes = _dupes_message(ctx.exception) + self.assertIsNotNone(dupes) + self.assertIn("dup-repeated", dupes) + self.assertNotIn("only-once", dupes) + + def test_case_sensitive_ids_not_duplicates(self): + # AC8: 'MMLU' vs 'mmlu' differ only by case -> NOT duplicates. + cfg = AccuracyConfig( + tasks=[ + AccuracyTask(id="MMLU", task="mmlu"), + AccuracyTask(id="mmlu", task="mmlu"), + ] + ) + self.assertEqual([t.id for t in cfg.tasks], ["MMLU", "mmlu"]) + + def test_empty_string_duplicates_render_as_quotes(self): + # AC9: two id="" -> duplicate; renders as '' in the sorted list. + with self.assertRaises(ValidationError) as ctx: + AccuracyConfig( + tasks=[ + AccuracyTask(id="", task="a"), + AccuracyTask(id="", task="b"), + ] + ) + dupes = _dupes_message(ctx.exception) + self.assertIsNotNone(dupes) + self.assertIn("''", dupes) + + +class TestAccuracyConfigNonMutation(unittest.TestCase): + """AC30, AC31: input list not mutated; default tasks list not shared.""" + + def test_caller_dict_list_not_mutated(self): + # AC30. + lst = [{"id": "a", "task": "m"}] + AccuracyConfig(tasks=lst) + self.assertEqual(len(lst), 1) + self.assertIsInstance(lst[0], dict) + self.assertEqual(lst[0], {"id": "a", "task": "m"}) + + def test_default_tasks_not_shared_between_instances(self): + # AC31. + a = AccuracyConfig() + b = AccuracyConfig() + self.assertIsNot(a.tasks, b.tasks) + a.tasks.append(AccuracyTask(id="x", task="m")) + self.assertEqual(b.tasks, []) + + +class TestValidationPrecedence(unittest.TestCase): + """AC32: per-element field error surfaces before the duplicate-id validator.""" + + def test_field_error_preempts_duplicate_validator(self): + with self.assertRaises(ValidationError) as ctx: + AccuracyConfig(tasks=[{"id": "a", "task": "gsm8k"}, {"id": "a"}]) + msg = str(ctx.exception) + # The missing required 'task' field on element index 1 is what surfaces. + # Assert the fully-qualified error location "tasks.1.task" rather than a + # bare "task": the parent field name "tasks" means a plain "task" + # substring would also match "tasks.1.id" (i.e. the *other* field being + # the one missing), so it cannot tell which required field failed. + self.assertIn("tasks.1.task", msg) + self.assertTrue( + ("Field required" in msg) or ("missing" in msg), + f"expected a missing-required-field marker, got: {msg}", + ) + # ...and the duplicate-id validator must NOT have run. + self.assertNotIn("duplicate task id(s):", msg) + + +class TestModelFieldMembership(unittest.TestCase): + """AC33, AC34 + regression constraints: closed-set field membership.""" + + def test_accuracy_task_fields_exact(self): + # AC33. + self.assertEqual( + set(AccuracyTask.model_fields), + { + "id", + "task", + "num_fewshot", + "metadata", + "include_path", + "num_concurrent", + "apply_chat_template", + "gen_kwargs", + }, + ) + + def test_accuracy_config_fields_exact(self): + # AC34. + self.assertEqual(set(AccuracyConfig.model_fields), {"tasks"}) + + def test_no_threshold_or_gate_fields(self): + # Regression constraint: no gating/threshold wiring exists on these models. + forbidden = {"threshold", "gate", "min_score", "accuracy_gate", "accuracy"} + self.assertEqual(set(AccuracyTask.model_fields) & forbidden, set()) + self.assertEqual(set(AccuracyConfig.model_fields) & forbidden, set()) + + +if __name__ == "__main__": + unittest.main() diff --git a/cvs/schema/config_file/inference/common/unittests/test_helpers.py b/cvs/schema/config_file/inference/common/unittests/test_helpers.py new file mode 100644 index 000000000..69eac39c9 --- /dev/null +++ b/cvs/schema/config_file/inference/common/unittests/test_helpers.py @@ -0,0 +1,41 @@ +"""Unit tests for shared inference helper schemas (inference/common/).""" + +import unittest + +from cvs.schema.config_file.inference.common.functional import FunctionalConfig +from cvs.schema.config_file.inference.common.long_context_accuracy import ( + LongContextAccCell, + LongContextAccuracyConfig, +) +from cvs.schema.config_file.inference.common.platform import PlatformConfig + + +class TestFunctionalConfig(unittest.TestCase): + def test_defaults_api_smoke_false(self): + cfg = FunctionalConfig() + self.assertFalse(cfg.api_smoke) + + +class TestPlatformConfig(unittest.TestCase): + def test_defaults_dmesg_scan_false(self): + cfg = PlatformConfig() + self.assertFalse(cfg.dmesg_scan) + + def test_defaults_gpu_metrics_poll_false(self): + cfg = PlatformConfig() + self.assertFalse(cfg.gpu_metrics_poll) + + +class TestLongContextAccuracyConfig(unittest.TestCase): + def test_rejects_duplicate_cell_ids(self): + with self.assertRaises(ValueError): + LongContextAccuracyConfig( + cells=[ + LongContextAccCell(id="a", isl=1024), + LongContextAccCell(id="a", isl=2048), + ] + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/cvs/schema/config_file/inference/common/unittests/test_sweep.py b/cvs/schema/config_file/inference/common/unittests/test_sweep.py new file mode 100644 index 000000000..9c59c255e --- /dev/null +++ b/cvs/schema/config_file/inference/common/unittests/test_sweep.py @@ -0,0 +1,116 @@ +"""Unit tests for shared inference sweep schema (inference/common/sweep.py).""" + +import unittest +import warnings + +from pydantic import ValidationError + +from cvs.schema.config_file.inference.common.sweep import ( + GoodputSlo, + Run, + SeqCombo, + Sweep, + validate_sweep_selector, + validate_thresholds_cover_sweep, +) + + +def _combo(name, isl="128", osl="2048"): + return SeqCombo(name=name, isl=isl, osl=osl) + + +class TestSweepValidator(unittest.TestCase): + def test_valid_runs_selector_constructs(self): + sw = Sweep( + sequence_combinations=[_combo("a"), _combo("b", osl="4096")], + runs=[Run(combo="a", concurrency=16), Run(combo="b", concurrency=32)], + ) + self.assertEqual([r.combo for r in sw.runs], ["a", "b"]) + + def test_unknown_run_combo_raises(self): + with self.assertRaises(ValidationError) as ctx: + Sweep( + sequence_combinations=[_combo("a")], + runs=[Run(combo="typo", concurrency=16)], + ) + self.assertIn("names no sequence_combination", str(ctx.exception)) + + def test_duplicate_combo_names_raise(self): + with self.assertRaises(ValidationError) as ctx: + Sweep( + sequence_combinations=[_combo("a"), _combo("a", osl="4096")], + runs=[Run(combo="a", concurrency=16)], + ) + self.assertIn("duplicate sequence_combination names", str(ctx.exception)) + + def test_concurrency_levels_is_rejected(self): + with self.assertRaises(ValidationError): + Sweep( + sequence_combinations=[_combo("a")], + runs=[Run(combo="a", concurrency=16)], + concurrency_levels=[16], + ) + + +class TestValidateSweepSelector(unittest.TestCase): + def test_unknown_combo_raises(self): + with self.assertRaisesRegex(ValueError, "names no sequence_combination"): + validate_sweep_selector(["a"], ["b"]) + + def test_duplicate_names_raises(self): + with self.assertRaisesRegex(ValueError, "duplicate sequence_combination names"): + validate_sweep_selector(["a", "a"], ["a"]) + + +class TestValidateThresholdsCoverSweep(unittest.TestCase): + _CELL = "ISL=128,OSL=2048,TP=8,CONC=16" + _GATED = {"p99_ttft_ms", "failed"} + + def test_missing_cell_raises_when_enforced(self): + with self.assertRaises(ValueError): + validate_thresholds_cover_sweep( + expected_cells=[self._CELL], + thresholds={}, + enforce_thresholds=True, + gated_metrics=self._GATED, + ) + + def test_missing_gated_metric_raises_when_enforced(self): + with self.assertRaises(ValueError): + validate_thresholds_cover_sweep( + expected_cells=[self._CELL], + thresholds={self._CELL: {}}, + enforce_thresholds=True, + gated_metrics=self._GATED, + ) + + def test_missing_gated_metric_warns_when_record_only(self): + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always") + validate_thresholds_cover_sweep( + expected_cells=[self._CELL], + thresholds={self._CELL: {}}, + enforce_thresholds=False, + gated_metrics=self._GATED, + ) + self.assertTrue(any("missing gated-metric specs" in str(x.message) for x in caught)) + + +class TestGoodputSlo(unittest.TestCase): + def test_valid_goodput_slo_constructs(self): + slo = GoodputSlo(ttft_ms=100.0, tpot_ms=50.0, e2el_ms=5000.0) + self.assertEqual(slo.e2el_ms, 5000.0) + + def test_extra_key_raises(self): + with self.assertRaises(ValidationError): + GoodputSlo(ttft_ms=1.0, tpot_ms=1.0, e2el_ms=1.0, ttft_msec=1.0) + + +class TestSeqComboForbid(unittest.TestCase): + def test_extra_key_raises(self): + with self.assertRaises(ValidationError): + SeqCombo(name="a", isl="128", osl="2048", unknown_field="x") + + +if __name__ == "__main__": + unittest.main() diff --git a/cvs/schema/config_file/inference/pytorch_xdit/__init__.py b/cvs/schema/config_file/inference/pytorch_xdit/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/cvs/schema/config_file/inference/pytorch_xdit/config.py b/cvs/schema/config_file/inference/pytorch_xdit/config.py new file mode 100644 index 000000000..2ea006fab --- /dev/null +++ b/cvs/schema/config_file/inference/pytorch_xdit/config.py @@ -0,0 +1,266 @@ +""" +PyTorch xDiT (WAN / Flux) inference configuration file schemas. + +Mirrors ``cvs/input/config_file/inference/pytorch_xdit/``. +""" + +from typing import Dict, List, Optional + +from pydantic import BaseModel, ConfigDict, Field, field_validator, model_validator + + +# ============================================================================= +# PyTorch XDit (WAN/Flux) Schemas +# ============================================================================= + + +class PytorchXditContainerConfig(BaseModel): + """Schema for container_config section in pytorch-xdit configs.""" + + model_config = ConfigDict(extra="allow") + + device_list: List[str] = Field( + default=["/dev/dri", "/dev/kfd"], description="List of device paths to mount in container" + ) + volume_dict: Dict[str, str] = Field(default_factory=dict, description="Host:container volume mount mappings") + env_dict: Dict[str, str] = Field(default_factory=dict, description="Environment variables for container") + + +class PytorchXditExpectedResults(BaseModel): + """Schema for expected_results in pytorch-xdit benchmark params.""" + + model_config = ConfigDict(extra="forbid") + + max_avg_total_time_s: float = Field(gt=0, description="Maximum acceptable average total_time in seconds") + + +class PytorchXditWan22Benchmarks(BaseModel): + """Schema for wan22_i2v_a14b benchmark parameters.""" + + model_config = ConfigDict(extra="allow") # Allow comment fields + + prompt: str = Field(description="Text prompt for image-to-video generation") + size: str = Field(default="720*1280", pattern=r"^\d+\*\d+$", description="Video resolution (format: height*width)") + frame_num: int = Field(default=81, ge=1, description="Number of frames to generate") + num_benchmark_steps: int = Field(default=5, ge=1, description="Number of benchmark iterations to run") + compile: bool = Field(default=True, description="Whether to use torch.compile for optimization") + torchrun_nproc: int = Field(default=8, ge=1, description="Number of processes for torchrun (usually num GPUs)") + expected_results: Dict[str, PytorchXditExpectedResults] = Field( + description="Expected results by GPU type (auto, mi300x, mi355, etc.)" + ) + + @field_validator('expected_results') + @classmethod + def validate_has_auto_or_specific( + cls, v: Dict[str, PytorchXditExpectedResults] + ) -> Dict[str, PytorchXditExpectedResults]: + """Ensure either 'auto' or a specific GPU type is present.""" + if not v: + raise ValueError("expected_results must contain at least one GPU type threshold") + if 'auto' not in v and not any(k in v for k in ['mi300x', 'mi325', 'mi350', 'mi355']): + raise ValueError("expected_results must contain either 'auto' or a specific GPU type (mi300x, mi325, etc.)") + return v + + +class PytorchXditFluxExpectedResults(BaseModel): + """Schema for expected_results in Flux benchmark params.""" + + model_config = ConfigDict(extra="forbid") + + max_avg_pipe_time_s: float = Field(gt=0, description="Maximum acceptable average pipe_time in seconds") + + +class PytorchXditFlux1DevBenchmarks(BaseModel): + """Schema for flux1_dev_t2i benchmark parameters.""" + + model_config = ConfigDict(extra="allow") # Allow comment fields + + prompt: str = Field(description="Text prompt for text-to-image generation") + seed: int = Field(default=42, description="Random seed for reproducibility") + num_inference_steps: int = Field(default=25, ge=1, description="Number of denoising steps") + max_sequence_length: int = Field(default=256, ge=1, description="Maximum sequence length for text encoder") + no_use_resolution_binning: bool = Field(default=True, description="Disable resolution binning") + warmup_steps: int = Field(default=1, ge=0, description="Number of warmup steps before benchmarking") + warmup_calls: int = Field(default=5, ge=0, description="Number of warmup calls") + num_repetitions: int = Field(default=25, ge=1, description="Number of benchmark repetitions") + height: int = Field(default=1024, ge=1, description="Output image height in pixels") + width: int = Field(default=1024, ge=1, description="Output image width in pixels") + ulysses_degree: int = Field(default=8, ge=1, description="Ulysses parallelism degree") + ring_degree: int = Field(default=1, ge=1, description="Ring parallelism degree") + use_torch_compile: bool = Field(default=True, description="Whether to use torch.compile for optimization") + torchrun_nproc: int = Field(default=8, ge=1, description="Number of processes for torchrun (usually num GPUs)") + expected_results: Dict[str, PytorchXditFluxExpectedResults] = Field( + description="Expected results by GPU type (auto, mi300x, mi355, etc.)" + ) + + @field_validator('expected_results') + @classmethod + def validate_has_auto_or_specific( + cls, v: Dict[str, PytorchXditFluxExpectedResults] + ) -> Dict[str, PytorchXditFluxExpectedResults]: + """Ensure either 'auto' or a specific GPU type is present.""" + if not v: + raise ValueError("expected_results must contain at least one GPU type threshold") + if 'auto' not in v and not any(k in v for k in ['mi300x', 'mi325', 'mi350', 'mi355']): + raise ValueError("expected_results must contain either 'auto' or a specific GPU type (mi300x, mi325, etc.)") + return v + + +class PytorchXditBenchmarkParams(BaseModel): + """Schema for benchmark_params section in pytorch-xdit configs.""" + + model_config = ConfigDict(extra="forbid") + + wan22_i2v_a14b: Optional[PytorchXditWan22Benchmarks] = Field( + default=None, description="WAN 2.2 image-to-video A14B benchmark parameters" + ) + flux1_dev_t2i: Optional[PytorchXditFlux1DevBenchmarks] = Field( + default=None, description="FLUX.1-dev text-to-image benchmark parameters" + ) + + +class PytorchXditWanConfigFile(BaseModel): + """ + Schema for PyTorch XDit WAN microbenchmark configuration file. + + Validates WAN inference config structure and provides fail-fast validation. + + Usage: + with open("mi300x_wan22_i2v_a14b.json") as f: + raw = json.load(f) + config = PytorchXditWanConfigFile.model_validate(raw) + """ + + model_config = ConfigDict(extra="forbid") + + config: 'PytorchXditWanConfig' = Field(description="Main configuration section") + benchmark_params: PytorchXditBenchmarkParams = Field(description="Benchmark parameters section") + + @model_validator(mode='after') + def validate_benchmark_present(self): + """Ensure at least one benchmark is configured.""" + if not self.benchmark_params.wan22_i2v_a14b: + raise ValueError("No benchmarks configured in 'benchmark_params' - at least wan22_i2v_a14b is required") + return self + + +class PytorchXditWanConfig(BaseModel): + """Schema for config section in pytorch-xdit WAN configs.""" + + model_config = ConfigDict(extra="forbid") + + container_image: str = Field( + default="amdsiloai/pytorch-xdit:v25.11.2", description="Docker image for pytorch-xdit container" + ) + container_name: str = Field(default="wan22-benchmark", description="Name for the Docker container") + hf_token_file: str = Field( + default="", + description=( + "Optional path to Hugging Face token file. " + "Not required when using a pre-staged local model path (recommended) or pre-cached HF snapshots (offline)." + ), + ) + hf_home: str = Field(description="Host directory for Hugging Face cache (mounted to /hf_home)") + output_base_dir: str = Field(description="Host base directory for benchmark outputs") + model_repo: str = Field( + default="Wan-AI/Wan2.2-I2V-A14B", + description=( + "Model identifier. Prefer an explicit local filesystem path (e.g., /models/Wan-AI/Wan2.2-I2V-A14B) " + "to avoid any runtime downloads. For backward compatibility, a Hugging Face repo id may be used only if " + "the snapshot is already cached under hf_home." + ), + ) + model_rev: str = Field( + default="206a9ee1b7bfaaf8f7e4d81335650533490646a3", + description="Model revision (commit hash). Ignored if model_repo is an explicit local filesystem path.", + ) + container_config: PytorchXditContainerConfig = Field( + default_factory=PytorchXditContainerConfig, description="Container device/volume/env configuration" + ) + + @field_validator('hf_token_file', 'hf_home', 'output_base_dir') + @classmethod + def validate_path_not_placeholder(cls, v: str, info) -> str: + """Check that paths are not still placeholders.""" + if not v: + return v + if '' in v.lower(): + raise ValueError(f"{info.field_name} contains placeholder ''. Please set a valid path in config.") + return v + + +class PytorchXditFluxConfigFile(BaseModel): + """ + Schema for PyTorch XDit Flux microbenchmark configuration file. + + Validates Flux inference config structure and provides fail-fast validation. + + Usage: + with open("mi300x_flux1_dev_t2i.json") as f: + raw = json.load(f) + config = PytorchXditFluxConfigFile.model_validate(raw) + """ + + model_config = ConfigDict(extra="forbid") + + config: 'PytorchXditFluxConfig' = Field(description="Main configuration section") + benchmark_params: PytorchXditBenchmarkParams = Field(description="Benchmark parameters section") + + @model_validator(mode='after') + def validate_benchmark_present(self): + """Ensure at least one benchmark is configured.""" + if not self.benchmark_params.flux1_dev_t2i: + raise ValueError("No benchmarks configured in 'benchmark_params' - at least flux1_dev_t2i is required") + return self + + +class PytorchXditFluxConfig(BaseModel): + """Schema for config section in pytorch-xdit Flux configs.""" + + model_config = ConfigDict(extra="forbid") + + container_image: str = Field( + default="amdsiloai/pytorch-xdit:v25.11.2", description="Docker image for pytorch-xdit container" + ) + container_name: str = Field(default="flux-benchmark", description="Name for the Docker container") + hf_token_file: str = Field( + default="", + description=( + "Optional path to Hugging Face token file. " + "Not required when using a pre-staged local model path (recommended) or pre-cached HF snapshots (offline)." + ), + ) + hf_home: str = Field(description="Host directory for Hugging Face cache (mounted to /hf_home)") + output_base_dir: str = Field(description="Host base directory for benchmark outputs") + model_repo: str = Field( + default="black-forest-labs/FLUX.1-dev", + description=( + "Model identifier. Prefer an explicit local filesystem path (e.g., /models/black-forest-labs/FLUX.1-dev) " + "to avoid any runtime downloads. For backward compatibility, a Hugging Face repo id may be used only if " + "the snapshot is already cached under hf_home." + ), + ) + model_rev: str = Field( + default="", + description=( + "Model revision (commit hash). Empty means use any available cached snapshot under hf_home. " + "Ignored if model_repo is an explicit local filesystem path." + ), + ) + container_config: PytorchXditContainerConfig = Field( + default_factory=PytorchXditContainerConfig, description="Container device/volume/env configuration" + ) + + @field_validator('hf_token_file', 'hf_home', 'output_base_dir') + @classmethod + def validate_path_not_placeholder(cls, v: str, info) -> str: + """Check that paths are not still placeholders.""" + if not v: + return v + if '' in v.lower(): + raise ValueError(f"{info.field_name} contains placeholder ''. Please set a valid path in config.") + return v + + +PytorchXditWanConfigFile.model_rebuild() +PytorchXditFluxConfigFile.model_rebuild() diff --git a/cvs/schema/config_file/inference/pytorch_xdit/unittests/__init__.py b/cvs/schema/config_file/inference/pytorch_xdit/unittests/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/cvs/schema/config_file/inference/pytorch_xdit/unittests/test_config.py b/cvs/schema/config_file/inference/pytorch_xdit/unittests/test_config.py new file mode 100644 index 000000000..08e60edbc --- /dev/null +++ b/cvs/schema/config_file/inference/pytorch_xdit/unittests/test_config.py @@ -0,0 +1,79 @@ +"""Unit tests for PyTorch xDiT config schemas (inference/pytorch_xdit/config.py).""" + +import json +import unittest +from pathlib import Path + +from pydantic import ValidationError + +from cvs.schema.config_file.inference.pytorch_xdit.config import ( + PytorchXditFluxConfigFile, + PytorchXditWanConfigFile, +) + +_PACKAGE_ROOT = Path(__file__).resolve().parents[5] +_WAN_SAMPLE = ( + _PACKAGE_ROOT / "input" / "config_file" / "inference" / "xdit" / "mi3xx_pytorch_xdit_wan22_14b_single.json" +) +_FLUX_SAMPLE = ( + _PACKAGE_ROOT / "input" / "config_file" / "inference" / "xdit" / "mi3xx_pytorch_xdit_flux1_dev_single.json" +) + + +class TestPytorchXditWanConfigFile(unittest.TestCase): + def test_sample_wan_json_validates(self): + raw = json.loads(_WAN_SAMPLE.read_text()) + config = PytorchXditWanConfigFile.model_validate(raw) + self.assertIsNotNone(config.benchmark_params.wan22_i2v_a14b) + + def test_wan_requires_wan_benchmark_block(self): + with self.assertRaisesRegex(ValidationError, "wan22_i2v_a14b"): + PytorchXditWanConfigFile.model_validate( + { + "config": { + "hf_home": "/hf", + "output_base_dir": "/out", + }, + "benchmark_params": {}, + } + ) + + def test_hf_home_changeme_rejected(self): + with self.assertRaisesRegex(ValidationError, "placeholder ''"): + PytorchXditWanConfigFile.model_validate( + { + "config": { + "hf_home": "/home//hf", + "output_base_dir": "/out", + }, + "benchmark_params": { + "wan22_i2v_a14b": { + "prompt": "test", + "expected_results": {"auto": {"max_avg_total_time_s": 100.0}}, + } + }, + } + ) + + +class TestPytorchXditFluxConfigFile(unittest.TestCase): + def test_sample_flux_json_validates(self): + raw = json.loads(_FLUX_SAMPLE.read_text()) + config = PytorchXditFluxConfigFile.model_validate(raw) + self.assertIsNotNone(config.benchmark_params.flux1_dev_t2i) + + def test_flux_requires_flux_benchmark_block(self): + with self.assertRaisesRegex(ValidationError, "flux1_dev_t2i"): + PytorchXditFluxConfigFile.model_validate( + { + "config": { + "hf_home": "/hf", + "output_base_dir": "/out", + }, + "benchmark_params": {}, + } + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/cvs/schema/config_file/inference/sglang/__init__.py b/cvs/schema/config_file/inference/sglang/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/cvs/schema/config_file/inference/sglang/unittests/__init__.py b/cvs/schema/config_file/inference/sglang/unittests/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/cvs/schema/config_file/inference/sglang/unittests/test_variant.py b/cvs/schema/config_file/inference/sglang/unittests/test_variant.py new file mode 100644 index 000000000..c1fbf5fa1 --- /dev/null +++ b/cvs/schema/config_file/inference/sglang/unittests/test_variant.py @@ -0,0 +1,80 @@ +"""Unit tests for SGLang inference variant schema (inference/sglang/variant.py).""" + +import unittest + +from pydantic import ValidationError + +from cvs.schema.config_file.inference.sglang.variant import ( + SglangSingleVariantConfig, + perf_cell_key, +) + + +def _minimal_variant(**overrides): + payload = { + "schema_version": 1, + "framework": "sglang_single", + "gpu_arch": "mi30x", + "enforce_thresholds": False, + "threshold_json": "test_threshold.json", + "paths": { + "shared_fs": "/home/test", + "models_dir": "/home/test/models", + "log_dir": "/home/test/LOGS", + "hf_token_file": "/home/test/.hf_token", + }, + "model": {"id": "test/model", "remote": 0}, + "container": { + "name": "sglang_test", + "image": "rocm/sglang:latest", + "runtime": {"name": "docker", "args": {}}, + }, + "benchmark_params": { + "tensor_parallelism": "8", + "pipeline_parallelism": "1", + "max_concurrency": "32", + "inference_tests": { + "bench_serv_random": {"input_length": 128, "output_length": 2048}, + }, + }, + } + payload.update(overrides) + return payload + + +class TestPerfCellKey(unittest.TestCase): + def test_builds_from_benchmark_params(self): + key = perf_cell_key( + { + "tensor_parallelism": "8", + "pipeline_parallelism": "2", + "max_concurrency": "16", + "inference_tests": {"bench_serv_random": {"input_length": 128, "output_length": 2048}}, + } + ) + self.assertEqual(key, "ISL=128,OSL=2048,TP=8,PP=2,CONC=16") + + +class TestSglangSingleVariantConfig(unittest.TestCase): + def test_minimal_payload_validates(self): + config = SglangSingleVariantConfig.model_validate(_minimal_variant()) + self.assertEqual(config.framework, "sglang_single") + self.assertEqual(config.perf_cell_key(), "ISL=128,OSL=2048,TP=8,PP=1,CONC=32") + + def test_syncs_legacy_inference_container_name(self): + config = SglangSingleVariantConfig.model_validate( + _minimal_variant( + inference={"container_name": "old"}, + ) + ) + self.assertEqual(config.inference["container_name"], config.container.name) + + def test_unknown_top_level_field_rejected(self): + payload = _minimal_variant() + payload["bogus"] = True + with self.assertRaises(ValidationError): + SglangSingleVariantConfig.model_validate(payload) + + +if __name__ == "__main__": + unittest.main() diff --git a/cvs/schema/config_file/inference/sglang/variant.py b/cvs/schema/config_file/inference/sglang/variant.py new file mode 100644 index 000000000..0e8c0b832 --- /dev/null +++ b/cvs/schema/config_file/inference/sglang/variant.py @@ -0,0 +1,66 @@ +""" +SGLang single-node inference variant config schema. + +Mirrors ``cvs/input/config_file/inference/sglang/``. +""" + +from __future__ import annotations + +from typing import Any, Dict, Mapping + +from pydantic import Field, model_validator +from typing_extensions import Literal + +from cvs.schema.common.base import BaseVariantConfig +from cvs.schema.base import _Forbid + + +class SglangRoleServer(_Forbid): + env: Dict[str, str] = Field(default_factory=dict) + serve_port: str = "8000" + + +class SglangRoles(_Forbid): + server: SglangRoleServer = Field(default_factory=SglangRoleServer) + + +def perf_cell_key(bp_dict: Mapping[str, Any]) -> str: + bench = (bp_dict.get("inference_tests") or {}).get("bench_serv_random") or {} + return ( + f"ISL={bench.get('input_length', '-')}," + f"OSL={bench.get('output_length', '-')}," + f"TP={bp_dict.get('tensor_parallelism', '8')}," + f"PP={bp_dict.get('pipeline_parallelism', '1')}," + f"CONC={bp_dict.get('max_concurrency', '-')}" + ) + + +class SglangSingleVariantConfig(BaseVariantConfig): + """Typed config for ``sglang_single`` + ContainerOrchestrator.""" + + framework: Literal["sglang_single"] + gpu_arch: str + variant_key: str = "" + config_path: str = "" + inference: Dict[str, Any] = Field(default_factory=dict) + benchmark_params: Dict[str, Any] = Field(default_factory=dict) + roles: SglangRoles = Field(default_factory=SglangRoles) + + def cell_key(self, isl, osl, concurrency) -> str: + tp = self.benchmark_params.get("tensor_parallelism", "-") + pp = self.benchmark_params.get("pipeline_parallelism", "-") + return f"ISL={isl},OSL={osl},TP={tp},PP={pp},CONC={concurrency}" + + def perf_cell_key(self) -> str: + return perf_cell_key(self.benchmark_params) + + @property + def hf_token_file(self) -> str: + return self.paths.hf_token_file + + @model_validator(mode="after") + def _sync_legacy_inference_container_name(self): + if self.inference and self.container.name: + self.inference["container_name"] = self.container.name + self.inference["container_image"] = self.container.image + return self diff --git a/cvs/schema/config_file/inference/vllm/__init__.py b/cvs/schema/config_file/inference/vllm/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/cvs/schema/config_file/inference/vllm/unittests/__init__.py b/cvs/schema/config_file/inference/vllm/unittests/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/cvs/schema/config_file/inference/vllm/unittests/test_variant.py b/cvs/schema/config_file/inference/vllm/unittests/test_variant.py new file mode 100644 index 000000000..136cf54c3 --- /dev/null +++ b/cvs/schema/config_file/inference/vllm/unittests/test_variant.py @@ -0,0 +1,168 @@ +"""Unit tests for vLLM inference variant schema (inference/vllm/variant.py).""" + +import json +import unittest +from pathlib import Path + +from cvs.lib.inference.utils.vllm_config_loader import GATED_GPU_METRICS, GATED_PROM_METRICS +from cvs.lib.inference.utils.vllm_parsing import GATED_METRICS +from cvs.schema.config_file.inference.common.sweep import Run, SeqCombo, Sweep +from cvs.schema.config_file.inference.vllm.variant import VariantConfig + +_PACKAGE_ROOT = Path(__file__).resolve().parents[5] + + +def _combo(name, isl="128", osl="2048"): + return SeqCombo(name=name, isl=isl, osl=osl) + + +def _full_gated_specs(): + """A spec for every gated client.*, gpu.*, and prom.* metric -- the + minimum that satisfies coverage. Values are inert so the set passes + without asserting anything; these tests pin the coverage gate, not the + numbers.""" + out = {} + for m in GATED_METRICS: + kind = "max_ms" if m.endswith("_ms") else "max" if m == "failed" else "min" + out[f"client.{m}"] = {"kind": kind, "value": 0 if kind == "min" else 1e12} + for m in GATED_GPU_METRICS: + kind = "max" if m in ("peak_gpu_memory_mb", "model_load_memory_mb", "model_load_s") else "min" + out[f"gpu.{m}"] = {"kind": kind, "value": 0 if kind == "min" else 1e12} + for m in GATED_PROM_METRICS: + out[f"prom.{m}"] = {"kind": "max_ms", "value": 1e12} + return out + + +class TestGpuGatedMetricCoverage(unittest.TestCase): + """The gpu.* axis of vllm_config_loader's _check_thresholds_cover_sweep.""" + + _CELL = "ISL=128,OSL=2048,TP=8,CONC=16" + + def _variant_with(self, thresholds, enforce): + sw = Sweep( + sequence_combinations=[_combo("a")], + runs=[Run(combo="a", concurrency=16)], + ) + return VariantConfig( + schema_version=1, + framework="vllm", + gpu_arch="mi300x", + enforce_thresholds=enforce, + paths={ + "shared_fs": "/home/x", + "models_dir": "/home/x/models", + "log_dir": "/home/x/LOGS", + "hf_token_file": "/home/x/.hf", + }, + model={"id": "amd/Llama-3.1-70B-Instruct-FP8-KV", "remote": 0}, + params={"tensor_parallelism": "8"}, + sweep=sw, + thresholds=thresholds, + ) + + def test_full_gated_set_constructs(self): + vc = self._variant_with({self._CELL: _full_gated_specs()}, enforce=True) + self.assertEqual(vc.enforce_thresholds, True) + + def test_missing_gpu_metric_does_not_raise_when_enforced(self): + # Operators may gate only a subset of gpu.* metrics; an absent one is + # simply not gated, not an authoring error. + specs = _full_gated_specs() + del specs["gpu.peak_gpu_memory_mb"] + vc = self._variant_with({self._CELL: specs}, enforce=True) + self.assertNotIn("gpu.peak_gpu_memory_mb", vc.thresholds[self._CELL]) + + def test_no_gpu_specs_at_all_does_not_raise_when_enforced(self): + vc = self._variant_with({self._CELL: {}}, enforce=True) + self.assertEqual(vc.thresholds[self._CELL], {}) + + def test_all_five_gpu_metrics_are_gated(self): + self.assertEqual( + GATED_GPU_METRICS, + { + "peak_gpu_memory_mb", + "model_load_memory_mb", + "model_load_s", + "gpu_bandwidth_util_pct", + "gpu_compute_util_pct", + }, + ) + + +class TestPromGatedMetricCoverage(unittest.TestCase): + """The prom.* axis of vllm_config_loader's _check_thresholds_cover_sweep. + + Mirrors TestGpuGatedMetricCoverage: prom.* is a fully separate, parallel + gated family, not part of client.*'s tiering machinery, so its coverage + is proven independently here rather than in test_vllm_deck_profile.py. + """ + + _CELL = "ISL=128,OSL=2048,TP=8,CONC=16" + + def _variant_with(self, thresholds, enforce): + sw = Sweep( + sequence_combinations=[_combo("a")], + runs=[Run(combo="a", concurrency=16)], + ) + return VariantConfig( + schema_version=1, + framework="vllm", + gpu_arch="mi300x", + enforce_thresholds=enforce, + paths={ + "shared_fs": "/home/x", + "models_dir": "/home/x/models", + "log_dir": "/home/x/LOGS", + "hf_token_file": "/home/x/.hf", + }, + model={"id": "amd/Llama-3.1-70B-Instruct-FP8-KV", "remote": 0}, + params={"tensor_parallelism": "8"}, + sweep=sw, + thresholds=thresholds, + ) + + def test_full_gated_set_constructs(self): + vc = self._variant_with({self._CELL: _full_gated_specs()}, enforce=True) + self.assertEqual(vc.enforce_thresholds, True) + + def test_missing_prom_metric_does_not_raise_when_enforced(self): + # Operators may gate only a subset of prom.* metrics; an absent one is + # simply not gated, not an authoring error. + specs = _full_gated_specs() + del specs["prom.queue_time_p50_ms"] + vc = self._variant_with({self._CELL: specs}, enforce=True) + self.assertNotIn("prom.queue_time_p50_ms", vc.thresholds[self._CELL]) + + def test_only_one_prom_metric_gated_does_not_raise_when_enforced(self): + specs = {"prom.queue_time_p50_ms": {"kind": "max_ms", "value": 200}} + vc = self._variant_with({self._CELL: specs}, enforce=True) + self.assertEqual(vc.thresholds[self._CELL], specs) + + def test_all_four_prom_metrics_are_gated(self): + self.assertEqual( + GATED_PROM_METRICS, + { + "queue_time_p50_ms", + "queue_time_p95_ms", + "prefill_time_p50_ms", + "prefill_time_p95_ms", + }, + ) + + +class TestVllmVariantSamples(unittest.TestCase): + def test_all_committed_variant_samples_validate(self): + config_dir = _PACKAGE_ROOT / "input" / "config_file" / "inference" / "vllm" + for path in sorted(config_dir.glob("*.json")): + if path.name.endswith("_threshold.json"): + continue + with self.subTest(sample=path.name): + raw = json.loads(path.read_text()) + known = {k: v for k, v in raw.items() if k in VariantConfig.model_fields} + known["enforce_thresholds"] = False + known["thresholds"] = {} + VariantConfig.model_validate(known) + + +if __name__ == "__main__": + unittest.main() diff --git a/cvs/schema/config_file/inference/vllm/variant.py b/cvs/schema/config_file/inference/vllm/variant.py new file mode 100644 index 000000000..2f79e8a30 --- /dev/null +++ b/cvs/schema/config_file/inference/vllm/variant.py @@ -0,0 +1,139 @@ +""" +Unified vLLM inference variant config schema. + +Mirrors ``cvs/input/config_file/inference/vllm/``. +""" + +from __future__ import annotations + +from typing import Any, Dict, List, Optional, Union + +from pydantic import Field, field_validator, model_validator +from typing_extensions import Literal + +from cvs.schema.base import _Allow, _Forbid +from cvs.schema.config_file.inference.common.accuracy import AccuracyConfig +from cvs.schema.config_file.inference.common.sweep import ( + RoleServer, + Sweep, + validate_thresholds_cover_sweep, +) + +_VLLM_LOG_LEVELS = {"debug", "info", "warning", "error", "critical"} + + +class ContainerConfig(_Allow): + lifetime: str = "per_run" + name: str = "" + image: str = "" + + +class Paths(_Forbid): + shared_fs: str + models_dir: str + log_dir: str + hf_token_file: str + + +class ModelSpec(_Forbid): + id: str + remote: Literal[0, 1] + + +class VllmRoleServer(RoleServer): + ib_hca_devices: Union[Literal["auto"], List[str], None] = None + ib_netdev: Optional[str] = None + + @field_validator("serve_args", mode="after") + @classmethod + def _check_log_level(cls, v): + level = v.get("log-level") + if level is not None and level not in _VLLM_LOG_LEVELS: + raise ValueError(f"serve_args.log-level must be one of {sorted(_VLLM_LOG_LEVELS)}, got: {level!r}") + return v + + +class VllmRoles(_Forbid): + server: VllmRoleServer = Field(default_factory=VllmRoleServer) + + +class Params(_Forbid): + backend: str = "vllm" + base_url: str = "http://0.0.0.0" + port_no: str = "8888" + dataset_name: str = "random" + burstiness: str = "1.0" + seed: str = "0" + request_rate: str = "inf" + random_range_ratio: str = "0.8" + random_prefix_len: str = "0" + tensor_parallelism: str = "8" + pipeline_parallel_size: str = "1" + master_addr: str = "localhost" + master_port: str = "29501" + nnodes: str = "1" + tokenizer_mode: str = "auto" + percentile_metrics: str = "ttft,tpot,itl,e2el" + metric_percentiles: str = "50,90,95,99" + num_prompts: str = "3200" + client_poll_count: str = "20" + + +class VariantConfig(_Forbid): + """Unified typed config for both single-node and distributed vllm runs.""" + + schema_version: Literal[1] + framework: Literal["vllm"] + gpu_arch: str + enforce_thresholds: bool = True + container: ContainerConfig = Field(default_factory=ContainerConfig) + paths: Paths + model: ModelSpec + roles: VllmRoles = Field(default_factory=VllmRoles) + params: Params = Field(default_factory=Params) + sweep: Sweep + thresholds: Dict[str, Dict[str, Any]] = Field(default_factory=dict) + accuracy: AccuracyConfig = Field(default_factory=AccuracyConfig) + + @model_validator(mode="after") + def _check_distributed_consistency(self): + nn = int(self.params.nnodes) + pp = int(self.params.pipeline_parallel_size) + is_ray = self.roles.server.serve_args.get("distributed-executor-backend") == "ray" + if nn > 1 and pp == 1 and not is_ray: + raise ValueError(f"nnodes={nn} > 1 requires pipeline_parallel_size > 1 (got pp={pp})") + if pp > 1 and nn == 1: + raise ValueError(f"pipeline_parallel_size={pp} > 1 requires nnodes > 1 (got nnodes={nn})") + if nn > 1 and not self.roles.server.ib_netdev: + raise ValueError( + "ib_netdev is required in roles.server when nnodes > 1. " + "Set it to the Linux network interface name for NCCL_SOCKET_IFNAME " + '(e.g. "ens51f1np1"). Cannot be auto-derived from HCA names.' + ) + return self + + @model_validator(mode="after") + def _check_remote_not_implemented(self): + if self.model.remote == 1: + raise NotImplementedError("model.remote=1 (remote model download) is not implemented.") + return self + + def cell_key(self, isl, osl, concurrency): + base = f"ISL={isl},OSL={osl},TP={self.params.tensor_parallelism}," + if int(self.params.pipeline_parallel_size) > 1: + base += f"PP={self.params.pipeline_parallel_size}," + return base + f"CONC={concurrency}" + + def expected_cells(self): + by_name = {c.name: c for c in self.sweep.sequence_combinations} + return [self.cell_key(by_name[r.combo].isl, by_name[r.combo].osl, r.concurrency) for r in self.sweep.runs] + + @model_validator(mode="after") + def _check_thresholds_cover_sweep(self): + validate_thresholds_cover_sweep( + expected_cells=self.expected_cells(), + thresholds=self.thresholds, + enforce_thresholds=self.enforce_thresholds, + gated_metrics=set(), + ) + return self diff --git a/cvs/schema/config_file/preflight/__init__.py b/cvs/schema/config_file/preflight/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/cvs/schema/config_file/preflight/config.py b/cvs/schema/config_file/preflight/config.py new file mode 100644 index 000000000..c8a453083 --- /dev/null +++ b/cvs/schema/config_file/preflight/config.py @@ -0,0 +1,637 @@ +""" +Preflight check configuration file schema. + +Mirrors ``cvs/input/config_file/preflight/``. +""" + +import warnings +from copy import deepcopy +from typing import List, Optional + +from pydantic import BaseModel, ConfigDict, Field, field_validator, model_validator + + +# ============================================================================= +# Preflight Check Configuration Schema +# ============================================================================= + + +LEGACY_PREFLIGHT_RDMA_PATHS = { + "gid_index": "gid_index", + "rdma_interfaces": "interfaces", +} + +PREFLIGHT_METADATA_PREFIXES = ("_comment", "_example") + + +def strip_preflight_metadata(value): + """Remove documentation-only pseudo-fields before schema validation. + + Preflight JSON files conventionally carry ``_comment*`` and ``_example*`` + keys so that the files remain self-documenting. They are not runtime + options. Strip only those reserved prefixes recursively, preserving strict + rejection of every other unknown customer-facing option. + """ + if isinstance(value, dict): + return { + key: strip_preflight_metadata(item) + for key, item in value.items() + if not (isinstance(key, str) and key.startswith(PREFLIGHT_METADATA_PREFIXES)) + } + if isinstance(value, list): + return [strip_preflight_metadata(item) for item in value] + return value + + +def normalize_legacy_preflight_rdma_config(value): + """Move the two deprecated node-check RDMA keys to their canonical block. + + Returns a deep-copied configuration and one consolidated warning message, + or the original value and ``None`` when no legacy keys are present. + Conflicting legacy and canonical values fail rather than silently choosing + which RDMA inventory should be tested. + """ + if not isinstance(value, dict): + return value, None + + node_check = value.get("node_check") + if not isinstance(node_check, dict): + return value, None + + legacy_keys = [key for key in LEGACY_PREFLIGHT_RDMA_PATHS if key in node_check] + if not legacy_keys: + return value, None + + normalized = deepcopy(value) + normalized_node_check = normalized["node_check"] + connectivity_check = normalized.setdefault("connectivity_check", {}) + if not isinstance(connectivity_check, dict): + raise ValueError( + "preflight.connectivity_check must be an object when deprecated node_check RDMA options are used" + ) + rdma = connectivity_check.setdefault("rdma", {}) + if not isinstance(rdma, dict): + raise ValueError( + "preflight.connectivity_check.rdma must be an object when deprecated node_check RDMA options are used" + ) + + migrations = [] + for legacy_key in legacy_keys: + canonical_key = LEGACY_PREFLIGHT_RDMA_PATHS[legacy_key] + legacy_value = normalized_node_check.pop(legacy_key) + if canonical_key in rdma and rdma[canonical_key] != legacy_value: + raise ValueError( + f"Conflicting preflight RDMA options: preflight.node_check.{legacy_key} and " + f"preflight.connectivity_check.rdma.{canonical_key} must have the same value when both are provided" + ) + rdma.setdefault(canonical_key, legacy_value) + migrations.append(f"preflight.node_check.{legacy_key} -> preflight.connectivity_check.rdma.{canonical_key}") + + warning_message = ( + "Deprecated preflight RDMA configuration detected: " + + ", ".join(migrations) + + ". Use the preflight.connectivity_check.rdma paths; legacy paths will be removed in a future release." + ) + return normalized, warning_message + + +LEGACY_PREFLIGHT_NODE_SMOKE_SECTIONS = { + "node_smoke": "node_smoke_tier1", + "tier3_info": "node_smoke_tier3", +} + + +def _preflight_section_has_values(section: dict) -> bool: + if not isinstance(section, dict): + return False + return any(value not in (None, "") for value in section.values()) + + +def normalize_legacy_preflight_node_smoke_sections(value): + """Copy legacy Node Smoke section names to their canonical tier keys. + + Returns a deep-copied configuration and one consolidated warning message, + or the original value and ``None`` when no legacy keys need migration. + Canonical sections win when both legacy and canonical blocks are populated. + """ + if not isinstance(value, dict): + return value, None + + legacy_keys = [key for key in LEGACY_PREFLIGHT_NODE_SMOKE_SECTIONS if key in value] + if not legacy_keys: + return value, None + + normalized = deepcopy(value) + migrations = [] + for legacy_key in legacy_keys: + canonical_key = LEGACY_PREFLIGHT_NODE_SMOKE_SECTIONS[legacy_key] + legacy_block = normalized.get(legacy_key) + if not isinstance(legacy_block, dict): + continue + canonical_block = normalized.get(canonical_key) + if isinstance(canonical_block, dict) and _preflight_section_has_values(canonical_block): + continue + normalized[canonical_key] = deepcopy(legacy_block) + migrations.append(f"preflight.{legacy_key} -> preflight.{canonical_key}") + + if not migrations: + return value, None + + warning_message = ( + "Deprecated preflight Node Smoke section name(s) detected: " + + ", ".join(migrations) + + ". Prefer node_smoke_tier1 and node_smoke_tier3 in new configs." + ) + return normalized, warning_message + + +class PreflightParallelismConfig(BaseModel): + """Legacy parallelism settings for preflight checks.""" + + model_config = ConfigDict(extra="allow") + + parallel_group_size: int = Field( + default=128, + ge=2, + le=512, + description=("Legacy alias for RDMA grouping. Prefer connectivity_check.rdma.nodes_per_full_mesh_group."), + ) + + +class PreflightDebugConfig(BaseModel): + """Debug and troubleshooting settings for preflight checks.""" + + model_config = ConfigDict(extra="allow") + + scriptlet: bool = Field( + default=False, + description=( + "Enable ScriptLet debug: preserve generated scripts/logs on remote nodes. " + "For RDMA connectivity, also wraps each ibv_rc_pingpong server in strace with " + "per-test traces under /tmp/preflight/strace_server__.log (expensive at scale)." + ), + ) + + +class PreflightNodeCheckConfig(BaseModel): + """Individual node validation settings.""" + + model_config = ConfigDict(extra="forbid") + + enabled: bool = Field(default=True, description="Enable generic GPU node health and ROCm validation") + gpus_per_node: int = Field(default=4, ge=1, description="Expected AMD GPU count on each node") + expected_rocm_version: str = Field(default="6.2.0", description="Expected ROCm version across all cluster nodes") + + +class PreflightRdmaConfig(BaseModel): + """RDMA connectivity testing settings.""" + + model_config = ConfigDict(extra="allow") + + connectivity_mode: str = Field(default="basic", description="RDMA connectivity testing: basic, full_mesh, or skip") + gid_index: str = Field(default="3", description="GID index to check on all RDMA interfaces (typically 3 for RoCE)") + interfaces: List[str] = Field( + default_factory=lambda: ["rocep28s0", "rocep62s0", "rocep79s0", "rocep96s0"], + min_length=1, + description="RDMA device names checked for presence, GID consistency, and connectivity", + ) + nodes_per_full_mesh_group: int = Field( + default=128, + ge=2, + le=512, + description=( + "Number of nodes in each full-mesh partition group (2-512). " + "Smaller groups use fewer resources per node but require more rounds." + ), + ) + parallel_group_size: int = Field( + default=128, + ge=2, + le=512, + description="Legacy alias for nodes_per_full_mesh_group.", + ) + ibv_test_timeout: int = Field( + default=90, + ge=1, + description="Timeout in seconds for RDMA connectivity tests using ibv_rc_pingpong", + ) + ibv_test_port_range: str = Field( + default="10000-50000", description="Port range for RDMA connectivity tests (format: start-end)" + ) + inter_full_mesh_group_pairs_per_wave: str = Field( + default="auto", description="Max ordered group-pairs per wave during inter-group testing ('auto' or integer)" + ) + inter_group_wave_pairs: str = Field( + default="auto", + description="Legacy alias for inter_full_mesh_group_pairs_per_wave.", + ) + prune_failure_threshold: float = Field( + default=0.5, + gt=0.0, + le=1.0, + description=( + "Round 1 (intra) prune before inter-group: prune nodes whose fraction of peers with ≥1 FAIL " + "intra test is ≥ this value (default 0.5). Peers counted per distinct other node in the same partition group." + ), + ) + port_retry_max: int = Field( + default=3, + ge=0, + le=10, + description=( + "After each ScriptLet wave (intra/inter), rerun only pairs whose logs show PORT_LISTEN_FAILED, " + "up to this many extra batches with new TCP ports (default 3)." + ), + ) + port_retry_gap: int = Field( + default=1000, + ge=1, + le=65535, + description=( + "When remapping ports for PORT_LISTEN_FAILED retries, start at (max port in batch) + this gap " + "to reduce overlap with ephemeral ports." + ), + ) + exclude_failed_interface_nodes: str = Field( + default="true", + description=( + "Legacy hint for reporting: preflight now prunes interface- and GID-failed nodes from the SSH " + "host list before RDMA; interface failures are not run in the mesh regardless of this flag." + ), + ) + + @field_validator('connectivity_mode') + @classmethod + def validate_connectivity_check(cls, v: str) -> str: + """Validate RDMA connectivity check setting.""" + valid_modes = ['basic', 'full_mesh', 'skip'] + if v not in valid_modes: + raise ValueError(f"connectivity_mode must be one of: {', '.join(valid_modes)}") + return v + + @field_validator('ibv_test_port_range') + @classmethod + def validate_port_range(cls, v: str) -> str: + """Validate port range format.""" + try: + start, end = map(int, v.split('-')) + if start >= end or start < 1024 or end > 65535: + raise ValueError("Invalid port range") + except (ValueError, AttributeError): + raise ValueError("ibv_test_port_range must be in format 'start-end' with valid port numbers") + return v + + +class PreflightL2PingConfig(BaseModel): + """Small customer-facing IFoE L2 ping policy.""" + + model_config = ConfigDict(extra="forbid") + + enabled: bool = Field(default=False, description="Enable the mandatory IFoE L2 connectivity gate") + pings_per_port: int = Field(default=3, ge=1, description="Ping samples per selected IFoE port pair") + + +class PreflightTransferBenchConfig(BaseModel): + """Small customer-facing TransferBench preflight policy.""" + + model_config = ConfigDict(extra="forbid") + + enabled: bool = Field(default=False, description="Enable the mandatory TransferBench preflight gate") + scope: str = Field(default="node", description="node for independent runs or cluster for one multi-rank run") + profile: str = Field(default="smoketest", description="CVS-supported TransferBench validation profile") + message_sizes: List[str] = Field( + default_factory=lambda: ["1K", "16M"], + min_length=1, + description="Message sizes exercised by the selected profile", + ) + iterations: int = Field(default=2, ge=1, description="Validated iterations per test and message size") + warmup_iterations: int = Field(default=0, ge=0, description="Warmup iterations before validation") + + @field_validator('scope') + @classmethod + def validate_transferbench_scope(cls, value: str) -> str: + normalized = value.strip().lower() + if normalized not in ('node', 'cluster'): + raise ValueError("TransferBench scope must be one of: node, cluster") + return normalized + + @field_validator('profile') + @classmethod + def validate_transferbench_profile(cls, value: str) -> str: + normalized = value.strip().lower() + if normalized != 'smoketest': + raise ValueError("TransferBench profile must be a CVS-supported profile: smoketest") + return normalized + + @field_validator('message_sizes') + @classmethod + def validate_transferbench_message_sizes(cls, value: List[str]) -> List[str]: + normalized = [str(size).strip() for size in value] + if any(not size for size in normalized): + raise ValueError("TransferBench message_sizes entries must not be empty") + return normalized + + +class PreflightIfoeConfig(BaseModel): + """MI4XX IFoE admission and data-path checks.""" + + model_config = ConfigDict(extra="forbid") + + fabric_checks: bool = Field( + default=False, + description="Enable MI4XX AIFM, AFM, vPOD, station-mask, and IFoE port admission", + ) + l2ping: PreflightL2PingConfig = Field( + default_factory=PreflightL2PingConfig, + description="Strict IFoE L2 connectivity admission", + ) + transferbench: PreflightTransferBenchConfig = Field( + default_factory=PreflightTransferBenchConfig, + description="TransferBench IFoE data-path validation", + ) + + +class PreflightConnectivityCheckConfig(BaseModel): + """Connectivity check settings by protocol.""" + + model_config = ConfigDict(extra="allow") + + rdma: PreflightRdmaConfig = Field(default_factory=PreflightRdmaConfig, description="RDMA connectivity settings") + ifoe: PreflightIfoeConfig = Field(default_factory=PreflightIfoeConfig, description="IFoE connectivity settings") + + +class PreflightNodeSmokeConfig(BaseModel): + """Primus node_smoke settings (primus-cli direct -- node_smoke).""" + + model_config = ConfigDict(extra="allow") + + connectivity_mode: str = Field( + default="skip", + description="Primus node_smoke mode: 'run' (host/GPU/RDMA roll-call) or 'skip' (default)", + ) + auto_setup: bool = Field( + default=True, + description="Clone/update Primus and prepare venv on each node before node_smoke", + ) + setup_timeout: int = Field(default=600, ge=60, description="SSH timeout in seconds for Primus auto_setup") + force_reclone: bool = Field( + default=False, + description="Remove primus_dir and clone fresh on every run (destructive)", + ) + shared_install: bool = Field( + default=True, + description=( + "When true (default), clone and venv setup run only on the first reachable node; " + "other nodes wait for the shared NFS home install. Set false only if each node has " + "a local primus_dir/venv_activate path." + ), + ) + pip_install_mode: str = Field( + default="minimal", + description="Venv deps: minimal (torch only), requirements, or skip", + ) + torch_pip_index_url: str = Field( + default="https://download.pytorch.org/whl/rocm6.2", + description="PyTorch ROCm wheel index URL for minimal pip_install_mode", + ) + primus_git_url: str = Field( + default="https://github.com/AMD-AIG-AIMA/Primus.git", + description="Primus repository URL for auto_setup clone", + ) + primus_git_branch: str = Field( + default="dev/preflight-direct-test", + description="Git branch to checkout during auto_setup", + ) + primus_git_recurse_submodules: bool = Field( + default=False, + description="Clone git submodules during auto_setup (not required for node_smoke)", + ) + primus_dir: str = Field( + default="/home/{user-id}/INSTALL/Primus", + description="Path to cloned Primus repo under the user's home directory (required when connectivity_mode is 'run')", + ) + venv_activate: str = Field( + default="/home/{user-id}/envs/preflight/.venv/bin/activate", + description="Path to Python venv activate script on each node (required when connectivity_mode is 'run')", + ) + gpus_per_node: int = Field(default=8, ge=1, description="GPUs per node for node_smoke") + master_port: int = Field(default=1234, ge=1024, le=65535, description="Distributed master port for node_smoke") + dump_path: str = Field( + default="", + description="Per-node dump directory for smoke JSON (default: /node_smoke)", + ) + expected_rdma_nics: Optional[int] = Field( + default=None, + ge=1, + description="Hard-fail when training RDMA NIC count differs (default: len(connectivity_check.rdma.interfaces))", + ) + ulimit_l_min_gb: float = Field(default=32.0, ge=0, description="Minimum RLIMIT_MEMLOCK in GiB (0 disables)") + shm_min_gb: float = Field(default=8.0, ge=0, description="Minimum /dev/shm size in GiB (0 disables)") + skip_dmesg: bool = Field(default=False, description="Skip dmesg error scan (e.g. unprivileged containers)") + allow_foreign_procs: bool = Field( + default=False, + description="Do not FAIL nodes with foreign GPU processes (still reported)", + ) + allowed_procs: str = Field( + default="gpuagent,rocm-smi-daemon,amd-smi,dcgm-exporter", + description="Comma-separated process names allowed to hold GPUs", + ) + require_tools: str = Field( + default="", + description="Comma-separated CLI tools that must exist in PATH (empty = warn only)", + ) + nccl_socket_ifname: str = Field(default="", description="NCCL_SOCKET_IFNAME override for node_smoke") + gloo_socket_ifname: str = Field( + default="", description="GLOO_SOCKET_IFNAME override (defaults to nccl_socket_ifname)" + ) + nccl_ib_hca: str = Field( + default="", + description="NCCL_IB_HCA override (defaults to comma-joined connectivity_check.rdma.interfaces)", + ) + nccl_ib_gid_index: Optional[int] = Field( + default=None, + description="NCCL_IB_GID_INDEX override (defaults to connectivity_check.rdma.gid_index)", + ) + rdma_nic_allowlist: str = Field( + default="", + description="Training NIC allowlist for node_smoke (defaults to connectivity_check.rdma.interfaces)", + ) + ssh_timeout: int = Field(default=300, ge=30, description="SSH timeout in seconds for each node_smoke run") + tier2_perf: bool = Field( + default=False, + description=( + "Enable Primus node_smoke Tier 2 perf sanity (--tier2-perf): " + "8192³ GEMM TFLOPS floor, HBM D2D bandwidth, local multi-GPU RCCL all-reduce" + ), + ) + gemm_tflops_min: float = Field( + default=600.0, + ge=0, + description="Tier 2 large GEMM TFLOPS floor (--gemm-tflops-min); used when tier2_perf is true", + ) + hbm_gbs_min: float = Field( + default=2000.0, + ge=0, + description="Tier 2 HBM device-to-device bandwidth floor in GB/s (--hbm-gbs-min)", + ) + rccl_gbs_min: float = Field( + default=100.0, + ge=0, + description="Tier 2 local multi-GPU RCCL all-reduce bandwidth floor in GB/s (--rccl-gbs-min)", + ) + rccl_size_mb: int = Field( + default=64, + ge=1, + description="Tier 2 local RCCL all-reduce message size in MB (--rccl-size-mb)", + ) + rccl_timeout_sec: int = Field( + default=120, + ge=30, + description="Tier 2 local RCCL all-reduce hard timeout in seconds (--rccl-timeout-sec)", + ) + extra_args: List[str] = Field( + default_factory=list, + description="Additional node_smoke CLI flags forwarded to primus-cli", + ) + + @field_validator("connectivity_mode") + @classmethod + def validate_node_smoke_mode(cls, v: str) -> str: + valid_modes = ["run", "skip"] + if v not in valid_modes: + raise ValueError(f"node_smoke.connectivity_mode must be one of: {', '.join(valid_modes)}") + return v + + +class PreflightTier3InfoConfig(BaseModel): + """Primus preflight Tier 3 Host/GPU/Network info (primus-cli direct -- preflight).""" + + model_config = ConfigDict(extra="allow") + + connectivity_mode: str = Field( + default="skip", + description="Tier 3 info mode: 'run' (preflight --host --gpu --network) or 'skip' (default)", + ) + auto_setup: bool = Field( + default=True, + description="Clone/update Primus and prepare venv before Tier 3 (uses node_smoke git/pip settings via PrimusSetup fallback)", + ) + primus_dir: str = Field( + default="", + description="Primus checkout path; empty uses tier3_info then node_smoke.primus_dir", + ) + venv_activate: str = Field( + default="", + description="Venv activate script; empty uses tier3_info then node_smoke.venv_activate", + ) + gpus_per_node: int = Field(default=8, ge=1, description="GPUs per node for torchrun") + master_port: int = Field(default=1234, ge=1024, le=65535, description="Distributed master port") + dump_path: str = Field( + default="", + description="Tier 3 report directory (default: /tier3_info)", + ) + report_file_name: str = Field(default="tier3_info", description="Base name for Primus markdown/PDF reports") + dist_timeout_sec: int = Field( + default=120, ge=30, description="Timeout for torch.distributed init during aggregated report" + ) + save_pdf: bool = Field(default=False, description="Generate PDF report via Primus") + nccl_socket_ifname: str = Field(default="", description="NCCL_SOCKET_IFNAME override") + gloo_socket_ifname: str = Field(default="", description="GLOO_SOCKET_IFNAME override") + nccl_ib_hca: str = Field( + default="", + description="NCCL_IB_HCA override (defaults to comma-joined connectivity_check.rdma.interfaces when empty)", + ) + nccl_ib_gid_index: Optional[int] = Field( + default=None, + description="NCCL_IB_GID_INDEX override (defaults to connectivity_check.rdma.gid_index when null)", + ) + ssh_timeout: int = Field(default=600, ge=30, description="SSH timeout in seconds for the Tier 3 cluster run") + extra_args: List[str] = Field(default_factory=list, description="Additional preflight CLI flags") + + @field_validator("connectivity_mode") + @classmethod + def validate_tier3_info_mode(cls, v: str) -> str: + valid_modes = ["run", "skip"] + if v not in valid_modes: + raise ValueError(f"tier3_info.connectivity_mode must be one of: {', '.join(valid_modes)}") + return v + + +class PreflightReportingConfig(BaseModel): + """Report generation and output settings.""" + + model_config = ConfigDict(extra="allow") + + generate_html_report: bool = Field(default=True, description="Whether to generate HTML report") + artifacts_root_dir: str = Field( + default="/tmp/preflight", + description=( + "Root directory for preflight artifacts. HTML report output and RDMA full_mesh ScriptLet logs use " + "/rdma_connectivity_workspace/// on each node (NFS-friendly). " + "Sample configs use /home/{user-id}/preflight; that placeholder is resolved only when present in JSON." + ), + ) + generate_rdma_pairs_csv: bool = Field( + default=True, + description="If true, write preflight_report_*_rdma_pairs.csv beside the HTML report (failed pairs only)", + ) + + +class PreflightConfigFile(BaseModel): + """ + Schema for preflight check configuration file. + + Uses nested structure organized by execution phase for better organization. + """ + + model_config = ConfigDict(extra="allow") # Allow comment fields + + parallelism: PreflightParallelismConfig = Field( + default_factory=PreflightParallelismConfig, description="Parallel execution settings" + ) + debug: PreflightDebugConfig = Field( + default_factory=PreflightDebugConfig, description="Debug and troubleshooting options" + ) + node_check: PreflightNodeCheckConfig = Field( + default_factory=PreflightNodeCheckConfig, description="Individual node validation settings" + ) + connectivity_check: PreflightConnectivityCheckConfig = Field( + default_factory=PreflightConnectivityCheckConfig, description="Inter-node connectivity tests" + ) + node_smoke: PreflightNodeSmokeConfig = Field( + default_factory=PreflightNodeSmokeConfig, description="Primus node_smoke checks" + ) + tier3_info: PreflightTier3InfoConfig = Field( + default_factory=PreflightTier3InfoConfig, + description="Primus Tier 3 preflight Host/GPU/Network info checks", + ) + reporting: PreflightReportingConfig = Field( + default_factory=PreflightReportingConfig, description="Report generation and output settings" + ) + + @model_validator(mode="before") + @classmethod + def reject_flat_preflight_checks(cls, value): + if not isinstance(value, dict): + return value + cleaned = strip_preflight_metadata(value) + removed = sorted(set(cleaned) & {"node_health", "l2ping", "transferbench"}) + if removed: + raise ValueError( + "Unsupported flat preflight block(s): " + + ", ".join(removed) + + "; use node_check and connectivity_check.ifoe" + ) + normalized, rdma_warning = normalize_legacy_preflight_rdma_config(cleaned) + if rdma_warning: + warnings.warn(rdma_warning, FutureWarning, stacklevel=2) + normalized, smoke_warning = normalize_legacy_preflight_node_smoke_sections(normalized) + if smoke_warning: + warnings.warn(smoke_warning, FutureWarning, stacklevel=2) + return normalized + + @model_validator(mode="after") + def validate_fabric_prerequisites(self): + if self.connectivity_check.ifoe.fabric_checks and not self.node_check.enabled: + raise ValueError("connectivity_check.ifoe.fabric_checks requires node_check.enabled=true") + return self diff --git a/cvs/schema/config_file/preflight/unittests/__init__.py b/cvs/schema/config_file/preflight/unittests/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/cvs/schema/config_file/preflight/unittests/test_config.py b/cvs/schema/config_file/preflight/unittests/test_config.py new file mode 100644 index 000000000..4465f6feb --- /dev/null +++ b/cvs/schema/config_file/preflight/unittests/test_config.py @@ -0,0 +1,317 @@ +"""Unit tests for preflight configuration schema (config_file/preflight/config.py).""" + +import unittest +import warnings + +from pydantic import ValidationError + +from cvs.schema.config_file.preflight.config import ( + PreflightConfigFile, + normalize_legacy_preflight_node_smoke_sections, + normalize_legacy_preflight_rdma_config, +) + + +class TestPreflightRdmaConfigSchema(unittest.TestCase): + def test_legacy_rdma_inventory_is_normalized_with_one_warning(self): + legacy = { + "node_check": { + "expected_rocm_version": "7.15.0", + "gid_index": "7", + "rdma_interfaces": ["enp4s0np0"], + }, + "connectivity_check": {"rdma": {"connectivity_mode": "basic"}}, + } + + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always") + config = PreflightConfigFile.model_validate(legacy) + + self.assertEqual(len(caught), 1) + self.assertIs(caught[0].category, FutureWarning) + self.assertIn("legacy paths will be removed in a future release", str(caught[0].message)) + self.assertEqual(config.connectivity_check.rdma.gid_index, "7") + self.assertEqual(config.connectivity_check.rdma.interfaces, ["enp4s0np0"]) + self.assertNotIn("gid_index", config.node_check.model_dump()) + self.assertNotIn("rdma_interfaces", config.node_check.model_dump()) + + def test_matching_legacy_and_canonical_rdma_values_are_accepted(self): + with self.assertWarns(FutureWarning): + config = PreflightConfigFile.model_validate( + { + "node_check": { + "gid_index": "7", + "rdma_interfaces": ["enp4s0np0"], + }, + "connectivity_check": { + "rdma": { + "gid_index": "7", + "interfaces": ["enp4s0np0"], + } + }, + } + ) + + self.assertEqual(config.connectivity_check.rdma.gid_index, "7") + self.assertEqual(config.connectivity_check.rdma.interfaces, ["enp4s0np0"]) + + def test_conflicting_legacy_and_canonical_rdma_values_are_rejected(self): + with self.assertRaisesRegex(ValidationError, "Conflicting preflight RDMA options"): + PreflightConfigFile.model_validate( + { + "node_check": {"gid_index": "3"}, + "connectivity_check": {"rdma": {"gid_index": "7"}}, + } + ) + + with self.assertRaisesRegex(ValidationError, "Conflicting preflight RDMA options"): + PreflightConfigFile.model_validate( + { + "node_check": {"rdma_interfaces": ["enp4s0np0"]}, + "connectivity_check": {"rdma": {"interfaces": ["mlx5_0"]}}, + } + ) + + def test_invalid_legacy_interface_value_uses_canonical_validation(self): + with self.assertWarns(FutureWarning): + with self.assertRaises(ValidationError): + PreflightConfigFile.model_validate( + {"node_check": {"rdma_interfaces": "enp4s0np0"}}, + ) + + def test_canonical_rdma_inventory_under_connectivity_check(self): + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always") + config = PreflightConfigFile.model_validate( + { + "connectivity_check": { + "rdma": { + "connectivity_mode": "skip", + "gid_index": "7", + "interfaces": ["enp4s0np0"], + } + }, + "reporting": { + "generate_html_report": True, + "generate_rdma_pairs_csv": False, + }, + } + ) + + self.assertEqual(caught, []) + self.assertEqual(config.connectivity_check.rdma.gid_index, "7") + self.assertEqual(config.connectivity_check.rdma.interfaces, ["enp4s0np0"]) + + +class TestNodeHealthConfigSchema(unittest.TestCase): + def test_documentation_pseudo_fields_stripped_typos_rejected(self): + config = PreflightConfigFile.model_validate( + { + "_comment": "Preflight settings", + "node_check": { + "_comment": "Node checks", + "_example_gpus_per_node": 8, + "enabled": True, + "gpus_per_node": 4, + "expected_rocm_version": "7.15.0", + }, + "connectivity_check": { + "ifoe": { + "_comment": "IFoE checks", + "l2ping": { + "_comment_enabled": "Enable strict L2 validation", + "enabled": True, + }, + } + }, + } + ) + + self.assertEqual(config.node_check.gpus_per_node, 4) + self.assertTrue(config.connectivity_check.ifoe.l2ping.enabled) + self.assertNotIn("_comment", config.node_check.model_extra or {}) + + with self.assertRaises(ValidationError): + PreflightConfigFile.model_validate( + {"node_check": {"enabled": True, "gpus_per_nod": 4}}, + ) + + def test_node_check_accepts_only_documented_customer_fields(self): + config = PreflightConfigFile.model_validate( + { + "node_check": { + "enabled": True, + "gpus_per_node": 4, + "expected_rocm_version": "7.15.0", + }, + "connectivity_check": {"ifoe": {"fabric_checks": True}}, + } + ) + + self.assertTrue(config.node_check.enabled) + self.assertEqual(config.node_check.gpus_per_node, 4) + self.assertTrue(config.connectivity_check.ifoe.fabric_checks) + + with self.assertRaises(ValidationError): + PreflightConfigFile.model_validate( + { + "node_check": { + "enabled": True, + "gpus_per_node": 4, + "failure_mode": "report", + } + } + ) + + with self.assertRaises(ValidationError): + PreflightConfigFile.model_validate( + { + "node_health": { + "enabled": True, + "gpus_per_node": 4, + "fabric_checks": True, + } + } + ) + + +class TestTransferBenchConfigSchema(unittest.TestCase): + def test_accepts_only_six_customer_facing_options(self): + config = PreflightConfigFile.model_validate( + { + "connectivity_check": { + "ifoe": { + "transferbench": { + "enabled": True, + "scope": "cluster", + "profile": "smoketest", + "message_sizes": ["1K", "16M"], + "iterations": 3, + "warmup_iterations": 1, + } + } + } + } + ) + + transferbench = config.connectivity_check.ifoe.transferbench + self.assertTrue(transferbench.enabled) + self.assertEqual(transferbench.scope, "cluster") + + with self.assertRaises(ValidationError): + PreflightConfigFile.model_validate( + { + "connectivity_check": { + "ifoe": { + "transferbench": { + "enabled": True, + "scope": "node", + "profile": "bandwidth", + } + } + } + } + ) + + with self.assertRaises(ValidationError): + PreflightConfigFile.model_validate( + { + "connectivity_check": { + "ifoe": { + "transferbench": { + "enabled": True, + "tb_binary": "/custom/TransferBench", + } + } + } + } + ) + + with self.assertRaises(ValidationError): + PreflightConfigFile.model_validate({"transferbench": {"enabled": True}}) + + +class TestL2PingConfigSchema(unittest.TestCase): + def test_accepts_only_two_customer_facing_options(self): + config = PreflightConfigFile.model_validate( + { + "connectivity_check": { + "ifoe": { + "l2ping": { + "enabled": True, + "pings_per_port": 5, + } + } + } + } + ) + + self.assertTrue(config.connectivity_check.ifoe.l2ping.enabled) + self.assertEqual(config.connectivity_check.ifoe.l2ping.pings_per_port, 5) + + with self.assertRaises(ValidationError): + PreflightConfigFile.model_validate( + { + "connectivity_check": { + "ifoe": { + "l2ping": { + "enabled": True, + "pings_per_port": 3, + "loss_threshold_pct": 1.0, + } + } + } + } + ) + + with self.assertRaises(ValidationError): + PreflightConfigFile.model_validate( + {"l2ping": {"enabled": True, "pings_per_port": 3}}, + ) + + +class TestLegacyNodeSmokeNormalization(unittest.TestCase): + def test_legacy_node_smoke_copied_to_tier1(self): + cfg = {"node_smoke": {"connectivity_mode": "run", "primus_dir": "/home/user/Primus"}} + normalized, warning = normalize_legacy_preflight_node_smoke_sections(cfg) + self.assertIsNotNone(warning) + self.assertEqual(normalized["node_smoke_tier1"]["primus_dir"], "/home/user/Primus") + + def test_canonical_tier1_not_overwritten_by_legacy(self): + cfg = { + "node_smoke_tier1": {"primus_dir": "/tier1/Primus"}, + "node_smoke": {"primus_dir": "/legacy/Primus"}, + } + normalized, warning = normalize_legacy_preflight_node_smoke_sections(cfg) + self.assertIsNone(warning) + self.assertEqual(normalized["node_smoke_tier1"]["primus_dir"], "/tier1/Primus") + + +class TestLegacyRdmaNormalizer(unittest.TestCase): + def test_normalize_legacy_rdma_returns_warning_message(self): + normalized, warning = normalize_legacy_preflight_rdma_config( + { + "node_check": { + "gid_index": "7", + "rdma_interfaces": ["enp4s0np0"], + }, + } + ) + self.assertIsNotNone(warning) + self.assertEqual(normalized["connectivity_check"]["rdma"]["gid_index"], "7") + self.assertEqual(normalized["connectivity_check"]["rdma"]["interfaces"], ["enp4s0np0"]) + + +class TestFabricPrerequisites(unittest.TestCase): + def test_fabric_checks_requires_node_check_enabled(self): + with self.assertRaisesRegex(ValidationError, "fabric_checks requires node_check.enabled"): + PreflightConfigFile.model_validate( + { + "node_check": {"enabled": False}, + "connectivity_check": {"ifoe": {"fabric_checks": True}}, + } + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/cvs/schema/config_file/training/__init__.py b/cvs/schema/config_file/training/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/cvs/schema/config_file/training/jaxmaxtext/__init__.py b/cvs/schema/config_file/training/jaxmaxtext/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/cvs/schema/config_file/training/jaxmaxtext/unittests/__init__.py b/cvs/schema/config_file/training/jaxmaxtext/unittests/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/cvs/schema/config_file/training/jaxmaxtext/unittests/test_variant.py b/cvs/schema/config_file/training/jaxmaxtext/unittests/test_variant.py new file mode 100644 index 000000000..837b6c5b8 --- /dev/null +++ b/cvs/schema/config_file/training/jaxmaxtext/unittests/test_variant.py @@ -0,0 +1,103 @@ +"""Unit tests for JAX MaxText training variant schema.""" + +import json +import unittest +import warnings +from pathlib import Path + +from pydantic import ValidationError + +from cvs.schema.config_file.training.jaxmaxtext.variant import ( + CheckpointResume, + Convergence, + LossCurve, + NcclConfig, + ScalingBaseline, + SmokeTest, + TrainingVariantConfig, + validate_thresholds_cover_training, +) + +_PACKAGE_ROOT = Path(__file__).resolve().parents[5] + + +class TestJaxMaxTextSchemaDefaults(unittest.TestCase): + def test_scaling_baseline_defaults(self): + sb = ScalingBaseline() + self.assertEqual(sb.tokens_per_sec_total, 0.0) + self.assertEqual(sb.num_nodes, 1) + + def test_convergence_defaults(self): + c = Convergence() + self.assertEqual(c.target_metric, "auto") + self.assertEqual(c.target_value, 0.0) + + def test_loss_curve_defaults(self): + lc = LossCurve() + self.assertEqual(lc.sample_every, 10) + self.assertEqual(lc.milestone_steps, [100, 500, 1000, 5000]) + self.assertTrue(lc.enforce) + + def test_smoke_defaults(self): + s = SmokeTest() + self.assertTrue(s.enabled) + self.assertEqual(s.steps, 5) + + def test_checkpoint_resume_defaults(self): + cr = CheckpointResume() + self.assertFalse(cr.enabled) + + +class TestNcclConfig(unittest.TestCase): + def test_ib_gid_index_changeme_rejected(self): + with self.assertRaises(ValidationError): + NcclConfig(ib_gid_index="") + + +class TestValidateThresholdsCoverTraining(unittest.TestCase): + _GATED = { + "training.tflops_per_sec_per_gpu": {"kind": "min", "value": 1}, + "training.tokens_per_sec_per_gpu": {"kind": "min", "value": 1}, + "training.final_loss": {"kind": "max", "value": 15}, + "training.loss_decreased": {"kind": "min", "value": 1}, + } + + def test_missing_cell_raises_when_enforced(self): + with self.assertRaises(ValueError): + validate_thresholds_cover_training( + expected_cells=["CELL_A"], + thresholds={}, + enforce_thresholds=True, + gated_metrics={"final_loss"}, + ) + + def test_full_coverage_passes(self): + with warnings.catch_warnings(): + warnings.simplefilter("error") + validate_thresholds_cover_training( + expected_cells=["CELL_A"], + thresholds={"CELL_A": dict(self._GATED)}, + enforce_thresholds=True, + gated_metrics={"final_loss", "loss_decreased", "tflops_per_sec_per_gpu", "tokens_per_sec_per_gpu"}, + ) + + +class TestJaxMaxTextVariantSamples(unittest.TestCase): + def test_all_committed_variant_samples_validate(self): + config_dir = _PACKAGE_ROOT / "input" / "config_file" / "training" / "jaxmaxtext" + for path in sorted(config_dir.glob("*.json")): + if path.name.endswith("_threshold.json"): + continue + with self.subTest(sample=path.name): + raw = json.loads(path.read_text()) + nccl = (raw.get("training") or {}).get("nccl") or {} + if any(v == "" for v in nccl.values() if isinstance(v, str)): + self.skipTest("distributed template with cluster-specific nccl placeholders") + known = {k: v for k, v in raw.items() if k in TrainingVariantConfig.model_fields} + known["enforce_thresholds"] = False + known["thresholds"] = {} + TrainingVariantConfig.model_validate(known) + + +if __name__ == "__main__": + unittest.main() diff --git a/cvs/schema/config_file/training/jaxmaxtext/variant.py b/cvs/schema/config_file/training/jaxmaxtext/variant.py new file mode 100644 index 000000000..ed2171b52 --- /dev/null +++ b/cvs/schema/config_file/training/jaxmaxtext/variant.py @@ -0,0 +1,180 @@ +""" +JAX MaxText training variant config schema. + +Mirrors ``cvs/input/config_file/training/jaxmaxtext/``. +""" + +from __future__ import annotations + +import warnings +from typing import Any, Dict, List, Literal + +from pydantic import field_validator + +from cvs.schema.common.base import BaseVariantConfig +from cvs.schema.base import _Allow, _Forbid + + +class Tokenizer(_Forbid): + hf_model_id: str + tokenizer_path: str + + +class NcclConfig(_Allow): + ib_hca_list: str = "" + ib_hca: str = "" + socket_ifname: str = "" + gloo_socket_ifname: str = "" + ib_gid_index: str = "3" + + @field_validator("ib_hca_list", "ib_hca", "socket_ifname", "gloo_socket_ifname", "ib_gid_index") + @classmethod + def _reject_changeme(cls, v, info): + if isinstance(v, str) and "" in v.lower(): + raise ValueError( + f"nccl.{info.field_name} is still ''. Set your cluster's RDMA/NIC " + "device/interface (see the sibling _example_* value) before running distributed training." + ) + return v + + +class JaxDistributed(_Forbid): + coordinator_ip: str = "auto" + coordinator_port: str = "12346" + initialization_timeout_seconds: str = "1800" + heartbeat_timeout_seconds: str = "900" + + +class RdmaLib(_Allow): + host_source_file: str = "" + container_mount_file: str = "" + container_dest_file: str = "" + + +class ScalingBaseline(_Allow): + tokens_per_sec_total: float = 0.0 + num_nodes: int = 1 + + +class Convergence(_Allow): + target_metric: Literal["auto", "train_loss", "eval_loss"] = "auto" + target_value: float = 0.0 + + +class LossCurve(_Allow): + sample_every: int = 10 + milestone_steps: List[int] = [100, 500, 1000, 5000] + max_slope: float = 0.0 + enforce: bool = True + + +class SmokeTest(_Allow): + enabled: bool = True + steps: int = 5 + per_device_batch_size: int = 1 + max_target_length: int = 2048 + + +class CheckpointResume(_Allow): + enabled: bool = False + sweep: str = "" + steps_before_ckpt: int = 6 + steps_after_resume: int = 6 + checkpoint_period: int = 5 + loss_tolerance: float = 0.1 + max_save_seconds: float = 0.0 + max_load_seconds: float = 0.0 + delete_ckpt_dir: bool = True + smoke_model_overrides: Dict[str, Any] = {} + + +class Sweep(_Allow): + name: str + maxtext_overrides: Dict[str, Any] = {} + + +class TrainingConfig(_Allow): + distributed: bool = True + gpus_per_node: int = 8 + verify_dmesg: bool = True + steps: int = 30 + enable_checkpointing: bool = False + train_script_paths: List[str] = [ + "/workspace/maxtext/src/maxtext/trainers/pre_train/train.py", + "/workspace/maxtext/src/MaxText/train.py", + ] + train_script: str = "/workspace/maxtext/src/MaxText/train.py" + maxtext_config: Dict[str, Any] = {} + tokenizer: Tokenizer + nic_type: str = "thor2" + rdma_lib: RdmaLib = RdmaLib() + env_vars: Dict[str, str] = {} + xla_flags: Dict[str, str] = {} + error_patterns: Dict[str, str] = {} + nccl: NcclConfig = NcclConfig() + jax_distributed: JaxDistributed = JaxDistributed() + scaling_baseline: ScalingBaseline = ScalingBaseline() + convergence: Convergence = Convergence() + loss_curve: LossCurve = LossCurve() + smoke: SmokeTest = SmokeTest() + checkpoint_resume: CheckpointResume = CheckpointResume() + sweeps: List[Sweep] = [] + enabled_sweep_list: List[str] = [] + + +def validate_thresholds_cover_training( + *, + expected_cells, + thresholds, + enforce_thresholds: bool, + gated_metrics=None, +) -> None: + """Shared training threshold/cell coverage check.""" + expected = set(expected_cells) + present = {k for k in thresholds.keys() if not str(k).startswith("_")} + missing = sorted(expected - present) + extra = sorted(present - expected) + problems = [] + if missing: + problems.append(f"training cells with no threshold entry: {missing}") + if extra: + problems.append(f"threshold keys matching no training cell (typo?): {extra}") + gated = gated_metrics if gated_metrics is not None else set() + gated_keys = [f"training.{m}" for m in sorted(gated)] + gated_gaps = {} + for cell in sorted(expected & present): + specs = thresholds.get(cell) or {} + absent = [k for k in gated_keys if k not in specs] + if absent: + gated_gaps[cell] = absent + if gated_gaps: + problems.append(f"cells missing gated-metric specs: {gated_gaps}") + if problems: + msg = "threshold.json does not match the training config; " + "; ".join(problems) + if enforce_thresholds: + raise ValueError(msg) + warnings.warn(f"{msg} (enforce_thresholds=false -> record-only)", stacklevel=3) + + +class TrainingVariantConfig(BaseVariantConfig): + framework: Literal["jaxmaxtext"] + gpu_arch: str + training: TrainingConfig + + def expected_cells(self): + names = [s.name for s in self.training.sweeps] + return names or ["default"] + + def enabled_sweeps(self): + sweeps = self.training.sweeps + if not sweeps: + return [Sweep(name="default")] + by_name = {s.name: s for s in sweeps} + names = self.training.enabled_sweep_list or [s.name for s in sweeps] + selected = [] + for n in names: + if n in by_name: + selected.append(by_name[n]) + else: + warnings.warn(f"enabled_sweep_list references unknown sweep '{n}'", stacklevel=2) + return selected diff --git a/cvs/schema/config_file/training/megatron/__init__.py b/cvs/schema/config_file/training/megatron/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/cvs/schema/config_file/training/megatron/unittests/__init__.py b/cvs/schema/config_file/training/megatron/unittests/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/cvs/schema/config_file/training/megatron/unittests/test_variant.py b/cvs/schema/config_file/training/megatron/unittests/test_variant.py new file mode 100644 index 000000000..9f992fbb5 --- /dev/null +++ b/cvs/schema/config_file/training/megatron/unittests/test_variant.py @@ -0,0 +1,95 @@ +"""Unit tests for Megatron training variant schema (training/megatron/variant.py).""" + +import json +import unittest +from pathlib import Path + +from pydantic import ValidationError + +from cvs.schema.config_file.training.megatron.variant import ( + MegatronVariantConfig, + validate_sweep_selector, +) + +_PACKAGE_ROOT = Path(__file__).resolve().parents[5] +_SAMPLE_CONFIG = ( + _PACKAGE_ROOT / "input" / "config_file" / "training" / "megatron" / "mi300x_megatron_llama-3.1-8b_single.json" +) + + +def _minimal_variant(**overrides): + payload = { + "schema_version": 1, + "framework": "megatron_single", + "gpu_arch": "MI300X", + "enforce_thresholds": False, + "config": {"training_iterations": "10"}, + "model_params": {"model_name": "test"}, + "container": { + "name": "megatron_test", + "image": "rocm/megatron:latest", + "runtime": {"name": "docker", "args": {}}, + }, + "sweep": { + "combinations": { + "cell_a": { + "name": "cell_a", + "micro_batch_size": "1", + "global_batch_size": "8", + "precision": "BF16", + } + }, + "runs": ["cell_a"], + }, + } + payload.update(overrides) + return payload + + +class TestMegatronSweepSelector(unittest.TestCase): + def test_duplicate_combination_keys_rejected(self): + with self.assertRaisesRegex(ValueError, "duplicate sweep.combinations keys"): + validate_sweep_selector(["a", "a"], ["a"]) + + def test_unknown_run_reference_rejected(self): + with self.assertRaisesRegex(ValueError, "unknown combinations"): + validate_sweep_selector(["a"], ["b"]) + + +class TestMegatronVariantConfig(unittest.TestCase): + def test_minimal_payload_validates(self): + config = MegatronVariantConfig.model_validate(_minimal_variant()) + self.assertEqual(config.framework, "megatron_single") + self.assertEqual(config.cell_key("cell_a"), "MBS=1,GBS=8,PRECISION=BF16") + + def test_sample_json_validates_without_thresholds_when_not_enforced(self): + raw = json.loads(_SAMPLE_CONFIG.read_text()) + known = {k: v for k, v in raw.items() if k in MegatronVariantConfig.model_fields} + known["enforce_thresholds"] = False + known["thresholds"] = {} + config = MegatronVariantConfig.model_validate(known) + self.assertEqual(config.gpu_arch, "MI300X") + + def test_enforce_thresholds_requires_matching_threshold_cells(self): + with self.assertRaisesRegex(ValidationError, "threshold.json does not match"): + MegatronVariantConfig.model_validate( + _minimal_variant(enforce_thresholds=True, thresholds={}), + ) + + def test_all_committed_variant_samples_validate(self): + config_dir = _PACKAGE_ROOT / "input" / "config_file" / "training" / "megatron" + for path in sorted(config_dir.glob("*.json")): + if path.name.endswith("_threshold.json"): + continue + with self.subTest(sample=path.name): + raw = json.loads(path.read_text()) + if raw.get("schema_version") != 1: + self.skipTest("legacy config without schema_version") + known = {k: v for k, v in raw.items() if k in MegatronVariantConfig.model_fields} + known["enforce_thresholds"] = False + known["thresholds"] = {} + MegatronVariantConfig.model_validate(known) + + +if __name__ == "__main__": + unittest.main() diff --git a/cvs/schema/config_file/training/megatron/variant.py b/cvs/schema/config_file/training/megatron/variant.py new file mode 100644 index 000000000..7c3f9d0c3 --- /dev/null +++ b/cvs/schema/config_file/training/megatron/variant.py @@ -0,0 +1,141 @@ +""" +Megatron training variant config schema. + +Mirrors ``cvs/input/config_file/training/megatron/``. +""" + +import warnings +from collections import Counter +from typing import Any, Dict, List + +from pydantic import Field, model_validator +from typing_extensions import Literal + +from cvs.schema.common.base import ContainerSpec +from cvs.schema.base import _Forbid + + +class MegatronSweepCombo(_Forbid): + name: str + micro_batch_size: str + global_batch_size: str + precision: str = "" + + +def validate_sweep_selector(combo_keys, run_refs): + """The sweep-selector rule: combination keys unique, every run references one.""" + counts = Counter(combo_keys) + dupes = sorted(k for k, count in counts.items() if count > 1) + if dupes: + raise ValueError(f"duplicate sweep.combinations keys: {dupes}") + known = set(counts) + unknown = sorted(r for r in run_refs if r not in known) + if unknown: + raise ValueError(f"sweep.runs references unknown combinations: {unknown} (known: {sorted(known)})") + + +def validate_thresholds_cover_sweep( + *, + expected_cells, + thresholds, + enforce_thresholds: bool, + gated_metrics=None, +) -> None: + """Shared sweep/threshold coverage check for training variant configs.""" + expected = set(expected_cells) + present = set(thresholds.keys()) + missing = sorted(expected - present) + extra = sorted(present - expected) + problems = [] + if missing: + problems.append(f"sweep cells with no threshold entry: {missing}") + if extra: + problems.append(f"threshold keys matching no sweep cell (typo?): {extra}") + gated = gated_metrics if gated_metrics is not None else set() + gated_keys = [f"training.{m}" for m in sorted(gated)] + gated_gaps = {} + for cell in sorted(expected & present): + specs = thresholds.get(cell) or {} + absent = [k for k in gated_keys if k not in specs] + if absent: + gated_gaps[cell] = absent + if gated_gaps: + problems.append(f"cells missing gated-metric specs: {gated_gaps}") + if problems: + msg = "threshold.json does not match the sweep matrix; " + "; ".join(problems) + if enforce_thresholds: + raise ValueError(msg) + warnings.warn(f"{msg} (enforce_thresholds=false -> record-only)", stacklevel=3) + + +class MegatronSweep(_Forbid): + combinations: Dict[str, MegatronSweepCombo] + runs: List[str] + + @model_validator(mode="after") + def _check_runs_reference_known_combos(self): + validate_sweep_selector( + list(self.combinations.keys()), + self.runs, + ) + return self + + +class ScalingBaseline(_Forbid): + tokens_per_sec_total: float = 0.0 + num_nodes: int = 1 + + +class LossCurveConfig(_Forbid): + sample_every: int = 10 + milestone_steps: List[int] = Field(default_factory=lambda: [100, 500, 1000, 5000]) + max_slope: float = 0.0 + enforce: bool = True + + +class ConvergenceConfig(_Forbid): + target_metric: Literal["auto", "train_loss", "eval_loss"] = "auto" + target_value: float = 0.0 + + +class CheckpointConfig(_Forbid): + enforce: bool = False + save_interval: int = 20 + save_iters: int = 21 + resume_iters: int = 25 + loss_rtol: float = 0.05 + checkpoint_dir: str = "" + + +class MegatronVariantConfig(_Forbid): + schema_version: Literal[1] + framework: Literal["megatron_single", "megatron_distributed"] + gpu_arch: str + enforce_thresholds: bool = True + threshold_json: str = "" + scaling_baseline: ScalingBaseline = Field(default_factory=ScalingBaseline) + loss_curve: LossCurveConfig = Field(default_factory=LossCurveConfig) + convergence: ConvergenceConfig = Field(default_factory=ConvergenceConfig) + checkpoint: CheckpointConfig = Field(default_factory=CheckpointConfig) + config: Dict[str, Any] + model_params: Dict[str, Any] + container: ContainerSpec + sweep: MegatronSweep + thresholds: Dict[str, Dict[str, Any]] = Field(default_factory=dict) + + def cell_key(self, combo_key: str) -> str: + combo = self.sweep.combinations[combo_key] + return f"MBS={combo.micro_batch_size},GBS={combo.global_batch_size},PRECISION={combo.precision}" + + def expected_cells(self) -> List[str]: + return [self.cell_key(k) for k in self.sweep.runs] + + @model_validator(mode="after") + def _check_thresholds_cover_sweep(self): + validate_thresholds_cover_sweep( + expected_cells=self.expected_cells(), + thresholds=self.thresholds, + enforce_thresholds=self.enforce_thresholds, + gated_metrics=set(), + ) + return self diff --git a/cvs/schema/config_file/training/torchtitan/__init__.py b/cvs/schema/config_file/training/torchtitan/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/cvs/schema/config_file/training/torchtitan/unittests/__init__.py b/cvs/schema/config_file/training/torchtitan/unittests/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/cvs/schema/config_file/training/torchtitan/unittests/test_variant.py b/cvs/schema/config_file/training/torchtitan/unittests/test_variant.py new file mode 100644 index 000000000..d9a05e86a --- /dev/null +++ b/cvs/schema/config_file/training/torchtitan/unittests/test_variant.py @@ -0,0 +1,80 @@ +"""Unit tests for TorchTitan training variant schema (training/torchtitan/variant.py).""" + +import json +import unittest +from pathlib import Path + +from pydantic import ValidationError + +from cvs.schema.config_file.training.torchtitan.variant import ( + TorchTitanVariantConfig, + validate_sweep_selector, +) + +_PACKAGE_ROOT = Path(__file__).resolve().parents[5] + + +def _minimal_variant(**overrides): + payload = { + "schema_version": 1, + "framework": "torchtitan_single", + "gpu_arch": "MI355X", + "enforce_thresholds": False, + "config": {"training_iterations": "10"}, + "model_params": {"model_name": "test"}, + "container": { + "name": "torchtitan_test", + "image": "rocm/torchtitan:latest", + "runtime": {"name": "docker", "args": {}}, + }, + "sweep": { + "combinations": { + "cell_a": { + "name": "cell_a", + "micro_batch_size": "2", + "global_batch_size": "16", + "precision": "BF16", + } + }, + "runs": ["cell_a"], + }, + } + payload.update(overrides) + return payload + + +class TestTorchTitanSweepSelector(unittest.TestCase): + def test_duplicate_combination_keys_rejected(self): + with self.assertRaisesRegex(ValueError, "duplicate sweep.combinations keys"): + validate_sweep_selector(["x", "x"], ["x"]) + + +class TestTorchTitanVariantConfig(unittest.TestCase): + def test_minimal_payload_validates(self): + config = TorchTitanVariantConfig.model_validate(_minimal_variant()) + self.assertEqual(config.framework, "torchtitan_single") + self.assertEqual(config.cell_key("cell_a"), "MBS=2,GBS=16,PRECISION=BF16") + + def test_enforce_thresholds_requires_matching_threshold_cells(self): + with self.assertRaisesRegex(ValidationError, "threshold.json does not match"): + TorchTitanVariantConfig.model_validate( + _minimal_variant(enforce_thresholds=True, thresholds={}), + ) + + def test_all_committed_variant_samples_validate(self): + config_dir = _PACKAGE_ROOT / "input" / "config_file" / "training" / "torchtitan" + for path in sorted(config_dir.glob("*.json")): + if path.name.endswith("_threshold.json"): + continue + with self.subTest(sample=path.name): + raw = json.loads(path.read_text()) + if raw.get("schema_version") != 1: + self.skipTest("legacy config without schema_version") + known = {k: v for k, v in raw.items() if k in TorchTitanVariantConfig.model_fields} + known["enforce_thresholds"] = False + known["thresholds"] = {} + TorchTitanVariantConfig.model_validate(known) + + +if __name__ == "__main__": + unittest.main() diff --git a/cvs/schema/config_file/training/torchtitan/variant.py b/cvs/schema/config_file/training/torchtitan/variant.py new file mode 100644 index 000000000..61f67412b --- /dev/null +++ b/cvs/schema/config_file/training/torchtitan/variant.py @@ -0,0 +1,141 @@ +""" +TorchTitan training variant config schema. + +Mirrors ``cvs/input/config_file/training/torchtitan/``. +""" + +import warnings +from collections import Counter +from typing import Any, Dict, List + +from pydantic import Field, model_validator +from typing_extensions import Literal + +from cvs.schema.common.base import ContainerSpec +from cvs.schema.base import _Forbid + + +class TorchTitanSweepCombo(_Forbid): + name: str + micro_batch_size: str + global_batch_size: str + precision: str = "" + + +def validate_sweep_selector(combo_keys, run_refs): + """The sweep-selector rule: combination keys unique, every run references one.""" + counts = Counter(combo_keys) + dupes = sorted(k for k, count in counts.items() if count > 1) + if dupes: + raise ValueError(f"duplicate sweep.combinations keys: {dupes}") + known = set(counts) + unknown = sorted(r for r in run_refs if r not in known) + if unknown: + raise ValueError(f"sweep.runs references unknown combinations: {unknown} (known: {sorted(known)})") + + +def validate_thresholds_cover_sweep( + *, + expected_cells, + thresholds, + enforce_thresholds: bool, + gated_metrics=None, +) -> None: + """Shared sweep/threshold coverage check for training variant configs.""" + expected = set(expected_cells) + present = set(thresholds.keys()) + missing = sorted(expected - present) + extra = sorted(present - expected) + problems = [] + if missing: + problems.append(f"sweep cells with no threshold entry: {missing}") + if extra: + problems.append(f"threshold keys matching no sweep cell (typo?): {extra}") + gated = gated_metrics if gated_metrics is not None else set() + gated_keys = [f"training.{m}" for m in sorted(gated)] + gated_gaps = {} + for cell in sorted(expected & present): + specs = thresholds.get(cell) or {} + absent = [k for k in gated_keys if k not in specs] + if absent: + gated_gaps[cell] = absent + if gated_gaps: + problems.append(f"cells missing gated-metric specs: {gated_gaps}") + if problems: + msg = "threshold.json does not match the sweep matrix; " + "; ".join(problems) + if enforce_thresholds: + raise ValueError(msg) + warnings.warn(f"{msg} (enforce_thresholds=false -> record-only)", stacklevel=3) + + +class TorchTitanSweep(_Forbid): + combinations: Dict[str, TorchTitanSweepCombo] + runs: List[str] + + @model_validator(mode="after") + def _check_runs_reference_known_combos(self): + validate_sweep_selector( + list(self.combinations.keys()), + self.runs, + ) + return self + + +class ScalingBaseline(_Forbid): + tokens_per_sec_total: float = 0.0 + num_nodes: int = 1 + + +class LossCurveConfig(_Forbid): + sample_every: int = 10 + milestone_steps: List[int] = Field(default_factory=lambda: [100, 500, 1000, 5000]) + max_slope: float = 0.0 + enforce: bool = True + + +class ConvergenceConfig(_Forbid): + target_metric: Literal["auto", "train_loss", "eval_loss"] = "auto" + target_value: float = 0.0 + + +class CheckpointConfig(_Forbid): + enforce: bool = False + save_interval: int = 20 + save_iters: int = 21 + resume_iters: int = 25 + loss_rtol: float = 0.05 + checkpoint_dir: str = "" + + +class TorchTitanVariantConfig(_Forbid): + schema_version: Literal[1] + framework: Literal["torchtitan_single", "torchtitan_distributed"] + gpu_arch: str + enforce_thresholds: bool = True + threshold_json: str = "" + scaling_baseline: ScalingBaseline = Field(default_factory=ScalingBaseline) + loss_curve: LossCurveConfig = Field(default_factory=LossCurveConfig) + convergence: ConvergenceConfig = Field(default_factory=ConvergenceConfig) + checkpoint: CheckpointConfig = Field(default_factory=CheckpointConfig) + config: Dict[str, Any] + model_params: Dict[str, Any] + container: ContainerSpec + sweep: TorchTitanSweep + thresholds: Dict[str, Dict[str, Any]] = Field(default_factory=dict) + + def cell_key(self, combo_key: str) -> str: + combo = self.sweep.combinations[combo_key] + return f"MBS={combo.micro_batch_size},GBS={combo.global_batch_size},PRECISION={combo.precision}" + + def expected_cells(self) -> List[str]: + return [self.cell_key(k) for k in self.sweep.runs] + + @model_validator(mode="after") + def _check_thresholds_cover_sweep(self): + validate_thresholds_cover_sweep( + expected_cells=self.expected_cells(), + thresholds=self.thresholds, + enforce_thresholds=self.enforce_thresholds, + gated_metrics=set(), + ) + return self diff --git a/cvs/schema/unittests/__init__.py b/cvs/schema/unittests/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/cvs/schema/unittests/test_validate.py b/cvs/schema/unittests/test_validate.py new file mode 100644 index 000000000..90f41e8cf --- /dev/null +++ b/cvs/schema/unittests/test_validate.py @@ -0,0 +1,92 @@ +"""Unit tests for cvs/schema/validate.py file-loading helper.""" + +import json +import tempfile +import unittest +from pathlib import Path + + +from cvs.schema.cluster_file.cluster import ClusterConfigFile +from cvs.schema.config_file.aorta.benchmark import AortaBenchmarkConfigFile +from cvs.schema.config_file.inference.pytorch_xdit.config import PytorchXditWanConfigFile +from cvs.schema.config_file.preflight.config import PreflightConfigFile +from cvs.schema.config_file.training.megatron.variant import MegatronVariantConfig +from cvs.schema.validate import validate_config_file + +_PACKAGE_ROOT = Path(__file__).resolve().parents[2] + + +class TestValidateConfigFile(unittest.TestCase): + def test_auto_detect_cluster_json(self): + sample = _PACKAGE_ROOT / "input" / "cluster_file" / "cluster.json" + config = validate_config_file(sample, config_type="auto") + self.assertIsInstance(config, ClusterConfigFile) + + def test_preflight_unwraps_top_level_key(self): + with tempfile.NamedTemporaryFile(mode="w", suffix=".json", delete=False) as handle: + json.dump( + {"preflight": {"node_check": {"enabled": True, "gpus_per_node": 4}}}, + handle, + ) + path = handle.name + + try: + config = validate_config_file(path, config_type="preflight") + self.assertIsInstance(config, PreflightConfigFile) + self.assertTrue(config.node_check.enabled) + finally: + Path(path).unlink(missing_ok=True) + + def test_aorta_yaml_load(self): + sample = _PACKAGE_ROOT / "input" / "config_file" / "aorta" / "aorta_benchmark.yaml" + config = validate_config_file(sample, config_type="aorta") + self.assertIsInstance(config, AortaBenchmarkConfigFile) + + def test_auto_detect_megatron_variant(self): + sample = ( + _PACKAGE_ROOT + / "input" + / "config_file" + / "training" + / "megatron" + / "mi300x_megatron_llama-3.1-8b_single.json" + ) + config = validate_config_file(sample, config_type="auto") + self.assertIsInstance(config, MegatronVariantConfig) + + def test_auto_detect_pytorch_xdit_wan(self): + sample = ( + _PACKAGE_ROOT / "input" / "config_file" / "inference" / "xdit" / "mi3xx_pytorch_xdit_wan22_14b_single.json" + ) + config = validate_config_file(sample, config_type="auto") + self.assertIsInstance(config, PytorchXditWanConfigFile) + + def test_missing_file_raises(self): + with self.assertRaises(FileNotFoundError): + validate_config_file("/nonexistent/config.json") + + def test_unknown_auto_detect_raises(self): + with tempfile.NamedTemporaryFile(mode="w", suffix=".json", delete=False) as handle: + json.dump({"unknown_top_level": True}, handle) + path = handle.name + + try: + with self.assertRaisesRegex(ValueError, "Cannot auto-detect config type"): + validate_config_file(path, config_type="auto") + finally: + Path(path).unlink(missing_ok=True) + + def test_empty_yaml_raises(self): + with tempfile.NamedTemporaryFile(mode="w", suffix=".yaml", delete=False) as handle: + handle.write("") + path = handle.name + + try: + with self.assertRaisesRegex(ValueError, "empty"): + validate_config_file(path, config_type="aorta") + finally: + Path(path).unlink(missing_ok=True) + + +if __name__ == "__main__": + unittest.main() diff --git a/cvs/schema/validate.py b/cvs/schema/validate.py new file mode 100644 index 000000000..a76c2b2ea --- /dev/null +++ b/cvs/schema/validate.py @@ -0,0 +1,147 @@ +""" +Load and validate CVS configuration files against Pydantic schemas. +""" + +from pathlib import Path +from typing import Any, Type, Union + +from cvs.schema.cluster_file.cluster import ClusterConfigFile +from cvs.schema.config_file.aorta.benchmark import AortaBenchmarkConfigFile +from cvs.schema.config_file.inference.pytorch_xdit.config import ( + PytorchXditFluxConfigFile, + PytorchXditWanConfigFile, +) +from cvs.schema.config_file.preflight.config import PreflightConfigFile +from cvs.schema.config_file.training.jaxmaxtext.variant import TrainingVariantConfig +from cvs.schema.config_file.training.megatron.variant import MegatronVariantConfig +from cvs.schema.config_file.training.torchtitan.variant import TorchTitanVariantConfig +from cvs.schema.config_file.inference.atom.variant import AtomVariantConfig +from cvs.schema.config_file.inference.sglang.variant import SglangSingleVariantConfig +from cvs.schema.config_file.inference.vllm.variant import VariantConfig as VllmVariantConfig + +# Populated as additional variant schemas land. +_VARIANT_FRAMEWORK_MAP: dict[str, tuple[str, Type[Any]]] = { + "megatron_single": ("megatron", MegatronVariantConfig), + "megatron_distributed": ("megatron", MegatronVariantConfig), + "torchtitan_single": ("torchtitan", TorchTitanVariantConfig), + "torchtitan_distributed": ("torchtitan", TorchTitanVariantConfig), + "jaxmaxtext": ("jaxmaxtext", TrainingVariantConfig), + "vllm": ("vllm", VllmVariantConfig), + "atom": ("atom", AtomVariantConfig), + "sglang_single": ("sglang", SglangSingleVariantConfig), +} + + +def _validate_variant_config(raw_config: dict, model_cls: Type[Any]): + """Structural validation for variant JSON (no cluster substitution). + + Threshold files are loaded separately at runtime; when ``thresholds`` is absent + or empty, disable enforcement so shape checks still pass on committed samples. + """ + known = {k: v for k, v in raw_config.items() if k in model_cls.model_fields} + known.setdefault("thresholds", {}) + if not known["thresholds"]: + known["enforce_thresholds"] = False + return model_cls.model_validate(known) + + +def validate_config_file( + config_path: Union[str, Path], config_type: str = "auto" +) -> Union[ + AortaBenchmarkConfigFile, + ClusterConfigFile, + PytorchXditWanConfigFile, + PytorchXditFluxConfigFile, + PreflightConfigFile, + MegatronVariantConfig, + TorchTitanVariantConfig, + TrainingVariantConfig, +]: + """ + Load and validate a configuration file. + + Args: + config_path: Path to configuration file (YAML or JSON) + config_type: Type of config - "aorta", "cluster", "pytorch_xdit_wan", + "pytorch_xdit_flux", "preflight", "megatron", "torchtitan", or + "auto" (detect from content) + + Returns: + Validated Pydantic model + + Raises: + ValueError: If config is invalid with detailed error message + FileNotFoundError: If config file doesn't exist + """ + import json + import yaml + + config_path = Path(config_path) + + if not config_path.exists(): + raise FileNotFoundError(f"Configuration file not found: {config_path}") + + with open(config_path) as f: + if config_path.suffix in ('.yaml', '.yml'): + raw_config = yaml.safe_load(f) + else: + raw_config = json.load(f) + + if raw_config is None: + raise ValueError(f"Configuration file is empty: {config_path}") + + if config_type == "auto": + if "node_dict" in raw_config: + config_type = "cluster" + elif "preflight" in raw_config: + config_type = "preflight" + elif "aorta_path" in raw_config: + config_type = "aorta" + elif raw_config.get("framework") in _VARIANT_FRAMEWORK_MAP: + config_type = _VARIANT_FRAMEWORK_MAP[raw_config["framework"]][0] + elif "config" in raw_config and "benchmark_params" in raw_config: + config_section = raw_config.get("config", {}) + benchmark_section = raw_config.get("benchmark_params", {}) + + if "flux1_dev_t2i" in benchmark_section or "FLUX" in config_section.get("model_repo", ""): + config_type = "pytorch_xdit_flux" + elif "wan22_i2v_a14b" in benchmark_section or "Wan" in config_section.get("model_repo", ""): + config_type = "pytorch_xdit_wan" + else: + config_type = "pytorch_xdit_wan" + else: + raise ValueError( + f"Cannot auto-detect config type for {config_path}. " + f"Specify config_type='aorta', config_type='cluster', " + f"config_type='pytorch_xdit_wan', config_type='pytorch_xdit_flux', " + "config_type='preflight', config_type='megatron', config_type='torchtitan', " + f"config_type='jaxmaxtext', config_type='vllm', config_type='atom', or " + f"config_type='sglang'" + ) + + try: + if config_type == "cluster": + return ClusterConfigFile.model_validate(raw_config) + if config_type == "preflight": + if "preflight" in raw_config: + return PreflightConfigFile.model_validate(raw_config["preflight"]) + raise ValueError("Preflight config must contain 'preflight' section") + if config_type == "aorta": + return AortaBenchmarkConfigFile.model_validate(raw_config) + if config_type == "pytorch_xdit_wan": + return PytorchXditWanConfigFile.model_validate(raw_config) + if config_type == "pytorch_xdit_flux": + return PytorchXditFluxConfigFile.model_validate(raw_config) + if config_type in ("megatron", "torchtitan", "jaxmaxtext", "vllm", "atom", "sglang"): + model_cls = { + "megatron": MegatronVariantConfig, + "torchtitan": TorchTitanVariantConfig, + "jaxmaxtext": TrainingVariantConfig, + "vllm": VllmVariantConfig, + "atom": AtomVariantConfig, + "sglang": SglangSingleVariantConfig, + }[config_type] + return _validate_variant_config(raw_config, model_cls) + raise ValueError(f"Unknown config_type: {config_type}") + except Exception as e: + raise ValueError(f"Invalid configuration in {config_path}:\n{e}") from e diff --git a/cvs/tests/benchmark/test_aorta.py b/cvs/tests/benchmark/test_aorta.py index c7f6f6767..a24ce859f 100644 --- a/cvs/tests/benchmark/test_aorta.py +++ b/cvs/tests/benchmark/test_aorta.py @@ -28,12 +28,9 @@ from cvs.runners._base_runner import RunStatus from cvs.parsers.aorta_report import AortaReportParser from cvs.parsers.tracelens import TraceLensParser -from cvs.parsers.schemas import ( - ParseStatus, - # Config validation schemas - ClusterConfigFile, - AortaBenchmarkConfigFile, -) +from cvs.parsers.schemas import ParseStatus +from cvs.schema.config_file.aorta.benchmark import AortaBenchmarkConfigFile +from cvs.schema.cluster_file.cluster import ClusterConfigFile from cvs.lib import globals from cvs.lib.utils_lib import ( diff --git a/cvs/tests/inference/xdit/pytorch_xdit_flux_dev_single.py b/cvs/tests/inference/xdit/pytorch_xdit_flux_dev_single.py index 2aa50f878..8f1909937 100644 --- a/cvs/tests/inference/xdit/pytorch_xdit_flux_dev_single.py +++ b/cvs/tests/inference/xdit/pytorch_xdit_flux_dev_single.py @@ -26,7 +26,8 @@ ) from cvs.lib import docker_lib from cvs.lib import globals -from cvs.parsers.schemas import ClusterConfigFile, PytorchXditFluxConfigFile +from cvs.schema.cluster_file.cluster import ClusterConfigFile +from cvs.schema.config_file.inference.pytorch_xdit.config import PytorchXditFluxConfigFile from cvs.lib.inference.xdit.pytorch_xdit_model_verify import ( build_diffusers_local_model_required_checks, verify_required_checks_on_nodes, diff --git a/cvs/tests/inference/xdit/pytorch_xdit_wan22_14b_single.py b/cvs/tests/inference/xdit/pytorch_xdit_wan22_14b_single.py index 832af87c5..42461e0bd 100644 --- a/cvs/tests/inference/xdit/pytorch_xdit_wan22_14b_single.py +++ b/cvs/tests/inference/xdit/pytorch_xdit_wan22_14b_single.py @@ -29,7 +29,8 @@ ) from cvs.lib import docker_lib from cvs.lib import globals -from cvs.parsers.schemas import ClusterConfigFile, PytorchXditWanConfigFile +from cvs.schema.cluster_file.cluster import ClusterConfigFile +from cvs.schema.config_file.inference.pytorch_xdit.config import PytorchXditWanConfigFile from cvs.lib.inference.xdit.pytorch_xdit_model_verify import ( build_diffusers_local_model_required_checks, resolve_wan_local_model_required_checks, diff --git a/cvs/tests/preflight/README.md b/cvs/tests/preflight/README.md index 80e3aae58..b2d8ec3fb 100644 --- a/cvs/tests/preflight/README.md +++ b/cvs/tests/preflight/README.md @@ -525,7 +525,7 @@ When adding new preflight checks: 1. Add the check class or function under `cvs/lib/preflight/` 2. Add the test function to `preflight_checks.py` 3. Update `report.py` summary and HTML generation -4. Add configuration parameters in `cvs/parsers/schemas.py` and `preflight_config.json` +4. Add configuration parameters in `cvs/schema/config_file/preflight/config.py` and `preflight_config.json` 5. Add unit tests under the module's `unittests/` directory 6. Update documentation in this README and `README_preflight_config.md` diff --git a/cvs/tests/preflight/preflight_checks.py b/cvs/tests/preflight/preflight_checks.py index 0e3f9e542..95f4588d9 100644 --- a/cvs/tests/preflight/preflight_checks.py +++ b/cvs/tests/preflight/preflight_checks.py @@ -28,11 +28,11 @@ from cvs.lib.parallel.config import ParallelConfig from cvs.lib.utils_lib import * from cvs.lib.verify_lib import * -from cvs.parsers.schemas import ( +from cvs.schema.config_file.preflight.config import ( normalize_legacy_preflight_node_smoke_sections, normalize_legacy_preflight_rdma_config, - validate_config_file, ) +from cvs.schema.validate import validate_config_file from cvs.lib import globals