From 7a1537c311b64e399086c1ac98527d84d74167db Mon Sep 17 00:00:00 2001 From: Mikhail Kuznetsov Date: Thu, 6 Aug 2026 16:53:17 +0000 Subject: [PATCH 1/7] feat(manifest): validate the manifest schema and settle one source of truth The manifest was checked for three top-level keys and nothing else, so a typo in a transport variable, a nested env value, or a node count that disagreed with itself surfaced as a failed multi-node job minutes later rather than as a startup error. Separately, the deployment target could be read from two places: a top-level slurm block selected the target while the values that took effect were the ones under deployment_config, so a manifest could be configured and ignored at the same time. Declare the shape in src/madengine/schemas/build_manifest.schema.json and validate it where the manifest is loaded, reporting the first violation with its JSON pointer. Unknown keys stay allowed so a manifest can carry consumer metadata. Two cross-field checks a schema cannot express are included: built_models entries must have a matching built_images entry, and slurm.nodes must equal distributed.nnodes. A top-level deployment block is folded into deployment_config with a warning, which leaves one place to read the target from. The schema accepts the nulls madengine itself writes: `madengine build` copies optional model fields straight from models.json, so a model declaring `"timeout": null` reaches the manifest as null rather than as the default the code appears to supply. Optional scalars, the tag list and the env/slurm/distributed blocks therefore allow null, while required fields and value types are unchanged -- a null block is a model that declared nothing, a nested null inside an env map is still an error. The blanket *.json ignore, there to keep model JSONs out of a dev checkout, would have swallowed the schema: runs from a checkout would work while the wheel shipped without it, since hatchling selects wheel contents by VCS status. src/madengine/ schemas/*.json is package data, so it is exempt. --- .gitignore | 4 + pyproject.toml | 1 + src/madengine/deployment/base.py | 7 +- .../orchestration/run_orchestrator.py | 8 +- src/madengine/schemas/__init__.py | 173 +++++++++++++++ .../schemas/build_manifest.schema.json | 150 +++++++++++++ tests/unit/test_manifest_schema.py | 201 ++++++++++++++++++ tests/unit/test_orchestration.py | 3 + 8 files changed, 545 insertions(+), 2 deletions(-) create mode 100644 src/madengine/schemas/__init__.py create mode 100644 src/madengine/schemas/build_manifest.schema.json create mode 100644 tests/unit/test_manifest_schema.py diff --git a/.gitignore b/.gitignore index c824efdf..3ab73834 100644 --- a/.gitignore +++ b/.gitignore @@ -131,6 +131,10 @@ venv/ docker/ scripts/ *.json +# ... except the schemas madengine ships and loads at runtime. The blanket rule above is +# for the model JSONs that land in a dev checkout; these are package data, and hatchling +# picks files by VCS status, so an ignored schema would also be missing from the wheel. +!src/madengine/schemas/*.json .*_env/ .vscode/ diff --git a/pyproject.toml b/pyproject.toml index dd9c7566..83b2978c 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -25,6 +25,7 @@ dependencies = [ "rich>=13.0.0", "click>=8.0.0", "jinja2>=3.0.0", + "jsonschema>=4.0.0", "pyyaml>=6.0", "kubernetes>=28.0.0", "pytest>=7.0", diff --git a/src/madengine/deployment/base.py b/src/madengine/deployment/base.py index 69e1367e..9381f729 100644 --- a/src/madengine/deployment/base.py +++ b/src/madengine/deployment/base.py @@ -19,6 +19,8 @@ from jinja2 import Environment, FileSystemLoader from rich.console import Console +from madengine.schemas import validate_build_manifest + # Regex for parsing "performance: " log lines. # Value: optional sign, integer/decimal, scientific notation (e or E). @@ -122,8 +124,8 @@ def __init__(self, config: DeploymentConfig): config: Deployment configuration """ self.config = config - self.manifest = self._load_manifest(config.manifest_file) self.console = Console() + self.manifest = self._load_manifest(config.manifest_file) def _load_manifest(self, manifest_file: str) -> Dict: """ @@ -152,6 +154,9 @@ def _load_manifest(self, manifest_file: str) -> Dict: if missing: raise ValueError(f"Invalid manifest, missing: {missing}") + for warning in validate_build_manifest(manifest, source=str(manifest_path)): + self.console.print(f"[yellow]⚠ {warning}[/yellow]") + return manifest # Template Method - defines workflow diff --git a/src/madengine/orchestration/run_orchestrator.py b/src/madengine/orchestration/run_orchestrator.py index 148c953c..5ad2ee77 100644 --- a/src/madengine/orchestration/run_orchestrator.py +++ b/src/madengine/orchestration/run_orchestrator.py @@ -31,6 +31,7 @@ ExecutionError, create_error_context, ) +from madengine.schemas import validate_build_manifest from madengine.utils.session_tracker import SessionTracker from madengine.orchestration.image_filtering import ( filter_images_by_gpu_compatibility as _filter_by_gpu_compat, @@ -221,7 +222,12 @@ def execute( # (with optional runtime override) with open(manifest_file) as f: manifest = json.load(f) - + + # Fails fast on a malformed manifest, and folds any top-level deployment + # block into deployment_config so the target is read from one place. + for warning in validate_build_manifest(manifest, source=str(manifest_file)): + self.rich_console.print(f"[yellow]⚠ {warning}[/yellow]") + deployment_config = manifest.get("deployment_config", {}) # Update additional_context with deployment_config for deployment layer diff --git a/src/madengine/schemas/__init__.py b/src/madengine/schemas/__init__.py new file mode 100644 index 00000000..ae78c6e7 --- /dev/null +++ b/src/madengine/schemas/__init__.py @@ -0,0 +1,173 @@ +#!/usr/bin/env python3 +""" +Declared shapes for the files madengine reads, and the validator that enforces them. + +The build manifest used to be checked for three top-level keys and nothing else, so a +typo in a transport variable or a node count that disagreed with itself surfaced as a +failed multi-node job minutes later instead of as a startup error. `build_manifest.schema.json` +declares the shape; `validate_build_manifest` reports the first violation with its JSON +pointer, plus the cross-field checks a schema cannot express. + +The schema is also where the deployment target is defined to live: under +`deployment_config`. A manifest that carries a top-level `slurm`/`k8s`/`distributed` block +is migrated into `deployment_config` with a warning, so the two used to disagree silently +and now cannot. + +Copyright (c) Advanced Micro Devices, Inc. All rights reserved. +""" + +import json +from pathlib import Path +from typing import Any, Dict, List, Optional + +from madengine.core.errors import ErrorContext, ValidationError + +SCHEMA_DIR = Path(__file__).parent + +#: Blocks that describe *where* a run is deployed. They belong under `deployment_config`. +DEPLOYMENT_BLOCKS = ("slurm", "k8s", "kubernetes", "distributed") + + +def load_schema(name: str = "build_manifest.schema.json") -> Dict[str, Any]: + """Load a bundled JSON Schema by file name.""" + with open(SCHEMA_DIR / name) as f: + return json.load(f) + + +def _pointer(path) -> str: + """Render a jsonschema error path as a JSON pointer (RFC 6901).""" + return "/" + "/".join(str(part) for part in path) if path else "/" + + +def migrate_top_level_deployment_blocks(manifest: Dict[str, Any]) -> List[str]: + """ + Move top-level deployment blocks under `deployment_config`, in place. + + Historically a manifest could carry a top-level `slurm` block that selected the SLURM + target while the values that took effect were the ones under `deployment_config`. Fold + the former into the latter so there is one place to read. + + Args: + manifest: parsed manifest, modified in place + + Returns: + List[str]: human-readable warnings, one per migrated or ignored block + """ + warnings: List[str] = [] + for block in DEPLOYMENT_BLOCKS: + if block not in manifest: + continue + value = manifest.pop(block) + deployment_config = manifest.setdefault("deployment_config", {}) + if block in deployment_config: + warnings.append( + f"manifest carries both '{block}' at the top level and under " + f"'deployment_config'; the top-level block was ignored" + ) + else: + deployment_config[block] = value + warnings.append( + f"manifest carries '{block}' at the top level; it belongs under " + f"'deployment_config' and was moved there" + ) + return warnings + + +def _semantic_warnings(manifest: Dict[str, Any]) -> List[str]: + """Cross-field checks that are advisory rather than fatal.""" + warnings: List[str] = [] + + images = manifest.get("built_images") or {} + models = manifest.get("built_models") or {} + unused = sorted(set(images) - set(models)) + if unused: + warnings.append( + f"built_images entries with no matching built_models entry: {', '.join(unused)}" + ) + + deployment_config = manifest.get("deployment_config") or {} + target = deployment_config.get("target") + if target == "slurm" and "slurm" not in deployment_config: + warnings.append( + "deployment_config.target is 'slurm' but there is no deployment_config.slurm " + "block; cluster defaults will be used" + ) + return warnings + + +def _semantic_errors(manifest: Dict[str, Any]) -> List[str]: + """Cross-field checks a JSON Schema cannot express, and that break a run.""" + errors: List[str] = [] + + images = manifest.get("built_images") or {} + models = manifest.get("built_models") or {} + orphans = sorted(set(models) - set(images)) + if orphans: + errors.append( + f"built_models entries with no matching built_images entry: " + f"{', '.join(orphans)}. The two dicts are joined by key." + ) + + deployment_config = manifest.get("deployment_config") or {} + nodes = (deployment_config.get("slurm") or {}).get("nodes") + nnodes = (deployment_config.get("distributed") or {}).get("nnodes") + if nodes is not None and nnodes is not None and nodes != nnodes: + errors.append( + f"deployment_config.slurm.nodes ({nodes}) != " + f"deployment_config.distributed.nnodes ({nnodes}); sbatch and the launcher " + f"would disagree on the world size" + ) + return errors + + +def validate_build_manifest( + manifest: Dict[str, Any], + source: Optional[str] = None, + migrate: bool = True, +) -> List[str]: + """ + Validate a build manifest against the bundled schema, fail-fast on the first error. + + Args: + manifest: parsed manifest; modified in place when `migrate` is set + source: manifest path, for error messages + migrate: fold top-level deployment blocks into `deployment_config` + + Returns: + List[str]: non-fatal warnings + + Raises: + ValidationError: on the first schema violation or failed cross-field check + """ + import jsonschema + + warnings = migrate_top_level_deployment_blocks(manifest) if migrate else [] + + where = f" in {source}" if source else "" + validator = jsonschema.Draft202012Validator(load_schema()) + first = next(iter(sorted(validator.iter_errors(manifest), key=lambda e: list(e.path))), None) + if first is not None: + raise ValidationError( + f"Invalid manifest{where}: {_pointer(first.absolute_path)}: {first.message}", + context=ErrorContext( + operation="manifest validation", + component="schemas", + file_path=source, + ), + suggestions=[ + "Check the field against src/madengine/schemas/build_manifest.schema.json", + ], + ) + + errors = _semantic_errors(manifest) + if errors: + raise ValidationError( + f"Invalid manifest{where}: {errors[0]}", + context=ErrorContext( + operation="manifest validation", + component="schemas", + file_path=source, + ), + ) + + return warnings + _semantic_warnings(manifest) diff --git a/src/madengine/schemas/build_manifest.schema.json b/src/madengine/schemas/build_manifest.schema.json new file mode 100644 index 00000000..459aa2f2 --- /dev/null +++ b/src/madengine/schemas/build_manifest.schema.json @@ -0,0 +1,150 @@ +{ + "$schema": "https://json-schema.org/draft/2020-12/schema", + "$id": "https://github.com/ROCm/madengine/schemas/build_manifest.schema.json", + "title": "madengine build manifest", + "description": "Shape of the manifest produced by `madengine build` and consumed by `madengine run`. Unknown keys are accepted everywhere so a manifest can carry consumer-specific metadata; the constraints below cover the fields madengine itself reads.", + "type": "object", + "required": ["built_images", "built_models", "context"], + "properties": { + "built_images": { + "description": "One entry per built image, keyed by image key. The same keys must appear in built_models.", + "type": "object", + "additionalProperties": { "$ref": "#/$defs/builtImage" } + }, + "built_models": { + "description": "One entry per model to run, keyed by the matching built_images key.", + "type": "object", + "additionalProperties": { "$ref": "#/$defs/builtModel" } + }, + "context": { "$ref": "#/$defs/context" }, + "deployment_config": { "$ref": "#/$defs/deploymentConfig" }, + "credentials_required": { "type": "array", "items": { "type": "string" } }, + "summary": { "type": "object" } + }, + "$defs": { + "envMap": { + "description": "Environment variables. Values are rendered into a shell command, so they must be scalars. The block itself may be null: that is what a model declaring no variables writes.", + "type": ["object", "null"], + "additionalProperties": { "type": ["string", "number", "boolean"] } + }, + "builtImage": { + "description": "An optional field carrying null means the model declared it as null in models.json; that is what madengine writes, so the schema accepts it.", + "type": "object", + "properties": { + "model": { "type": ["string", "null"] }, + "docker_image": { "type": ["string", "null"] }, + "dockerfile": { "type": ["string", "null"] }, + "base_docker": { "type": ["string", "null"] }, + "local_image": { "type": ["boolean", "null"] }, + "registry_image": { "type": ["string", "null"] }, + "registry": { "type": ["string", "null"] }, + "gpu_vendor": { "type": ["string", "null"] }, + "build_duration": { "type": ["number", "string", "null"] } + } + }, + "builtModel": { + "type": "object", + "required": ["name"], + "properties": { + "name": { "type": "string" }, + "url": { "type": ["string", "null"] }, + "dockerfile": { "type": ["string", "null"] }, + "scripts": { + "description": "Run script, resolved under MODEL_DIR.", + "type": ["string", "null"] + }, + "n_gpus": { "type": ["integer", "string", "null"] }, + "owner": { "type": ["string", "null"] }, + "training_precision": { "type": ["string", "null"] }, + "multiple_results": { + "description": "Name of the per-model CSV madengine collects instead of a single metric.", + "type": ["string", "null"] + }, + "tags": { "type": ["array", "null"], "items": { "type": "string" } }, + "timeout": { "type": ["integer", "string", "null"] }, + "args": { "type": ["string", "null"] }, + "additional_docker_run_options": { "type": ["string", "null"] }, + "env_vars": { "$ref": "#/$defs/envMap" }, + "slurm": { "$ref": "#/$defs/slurm" }, + "distributed": { "$ref": "#/$defs/distributed" } + } + }, + "context": { + "type": "object", + "properties": { + "docker_env_vars": { "$ref": "#/$defs/envMap" }, + "docker_mounts": { + "description": "Bind mounts keyed {container_path: host_path}; madengine renders each as -v :.", + "type": ["object", "null"], + "additionalProperties": { "type": "string" } + }, + "docker_build_arg": { "$ref": "#/$defs/envMap" }, + "docker_gpus": { "type": ["string", "integer", "null"] }, + "gpu_vendor": { "type": ["string", "null"] }, + "guest_os": { "type": ["string", "null"] }, + "skip_perf_collection": { "type": ["boolean", "null"] } + } + }, + "deploymentConfig": { + "description": "The single source of truth for where and how the run is deployed. Deployment blocks belong here, not at the top level of the manifest.", + "type": "object", + "properties": { + "target": { + "type": "string", + "enum": ["local", "docker", "slurm", "k8s", "kubernetes"] + }, + "slurm": { "$ref": "#/$defs/slurm" }, + "k8s": { "type": ["object", "null"] }, + "kubernetes": { "type": ["object", "null"] }, + "distributed": { "$ref": "#/$defs/distributed" }, + "env_vars": { "$ref": "#/$defs/envMap" }, + "env_file": { + "description": "Path to a shell env file (e.g. mad.env) that madengine sources before the run.", + "type": "string" + }, + "debug": { "type": "boolean" }, + "docker_gpus": { "type": ["string", "integer"] }, + "gpus_per_node": { "type": "integer", "minimum": 1 } + } + }, + "slurm": { + "type": ["object", "null"], + "properties": { + "partition": { "type": "string" }, + "account": { "type": "string" }, + "qos": { "type": "string" }, + "reservation": { "type": "string" }, + "nodelist": { "type": "string" }, + "exclude": { "type": "string" }, + "constraint": { "type": "string" }, + "nodes": { "type": "integer", "minimum": 1 }, + "gpus_per_node": { "type": "integer", "minimum": 1 }, + "skip_gpus_directive": { + "description": "Omit #SBATCH --gpus-per-node, for clusters that advertise no GPU GRES.", + "type": "boolean" + }, + "time": { "type": "string" }, + "output_dir": { "type": "string" }, + "results_dir": { "type": "string" }, + "shared_workspace": { "type": "string" }, + "exclusive": { "type": "boolean" }, + "network_interface": { "type": "string" }, + "cluster_profile": { + "description": "Name of a per-cluster preset profile under deployment/presets/slurm/clusters/.", + "type": "string" + }, + "modules": { "type": "array", "items": { "type": "string" } } + } + }, + "distributed": { + "type": ["object", "null"], + "properties": { + "launcher": { "type": "string" }, + "backend": { "type": "string" }, + "port": { "type": "integer", "minimum": 1, "maximum": 65535 }, + "nnodes": { "type": "integer", "minimum": 1 }, + "nproc_per_node": { "type": "integer", "minimum": 1 } + } + } + } +} diff --git a/tests/unit/test_manifest_schema.py b/tests/unit/test_manifest_schema.py new file mode 100644 index 00000000..9ad36d48 --- /dev/null +++ b/tests/unit/test_manifest_schema.py @@ -0,0 +1,201 @@ +#!/usr/bin/env python3 +""" +Unit tests for build-manifest schema validation. + +The manifest used to be checked for three top-level keys, so the failures below all +surfaced minutes later as a failed multi-node job. Each test pins one of them to a +startup error instead: + +1. A field with the wrong type is reported with its JSON pointer. +2. `built_models` and `built_images` are joined by key, so an orphan model is fatal. +3. `slurm.nodes` and `distributed.nnodes` must agree or sbatch and the launcher + disagree on the world size. +4. A top-level deployment block is folded into `deployment_config`, which is the one + place the target is read from. + +Copyright (c) Advanced Micro Devices, Inc. All rights reserved. +""" + +import copy + +import pytest + +from madengine.core.errors import ValidationError +from madengine.schemas import ( + load_schema, + migrate_top_level_deployment_blocks, + validate_build_manifest, +) + + +VALID_MANIFEST = { + "built_images": { + "img": { + "model": "dummy", + "docker_image": "rocm/dummy:latest", + "dockerfile": "docker/dummy.ubuntu.amd.Dockerfile", + "local_image": True, + } + }, + "built_models": { + "img": { + "name": "dummy", + "scripts": "scripts/dummy/run.sh", + "n_gpus": "-1", + "tags": ["pyt", "training"], + "timeout": -1, + "args": "--model_repo dummy", + "multiple_results": "perf_dummy.csv", + } + }, + "context": { + "docker_env_vars": {"NCCL_DEBUG": "INFO"}, + "docker_mounts": {"/dev/infiniband": "/dev/infiniband"}, + "docker_build_arg": {}, + "gpu_vendor": "AMD", + "guest_os": "UBUNTU", + "docker_gpus": "0,1,2,3,4,5,6,7", + }, + "deployment_config": { + "target": "slurm", + "slurm": {"partition": "meta64", "nodes": 2, "gpus_per_node": 8}, + "distributed": {"launcher": "torchrun", "nnodes": 2, "nproc_per_node": 8}, + "env_vars": {}, + }, +} + + +@pytest.fixture +def manifest(): + return copy.deepcopy(VALID_MANIFEST) + + +class TestSchemaIsWellFormed: + def test_schema_loads(self): + schema = load_schema() + assert schema["$schema"].startswith("https://json-schema.org/draft/2020-12") + assert set(schema["required"]) == {"built_images", "built_models", "context"} + + def test_valid_manifest_passes_without_warnings(self, manifest): + assert validate_build_manifest(manifest) == [] + + +class TestSchemaViolations: + def test_wrong_type_is_reported_with_json_pointer(self, manifest): + manifest["deployment_config"]["slurm"]["nodes"] = "2" + with pytest.raises(ValidationError) as excinfo: + validate_build_manifest(manifest) + assert "/deployment_config/slurm/nodes" in str(excinfo.value) + + def test_nested_env_var_value_is_rejected(self, manifest): + """Env values are rendered into a shell command, so they must be scalars.""" + manifest["context"]["docker_env_vars"]["NCCL_IB_HCA"] = {"device": "bnxt_re0"} + with pytest.raises(ValidationError) as excinfo: + validate_build_manifest(manifest) + assert "/context/docker_env_vars/NCCL_IB_HCA" in str(excinfo.value) + + def test_unknown_deployment_target_is_rejected(self, manifest): + manifest["deployment_config"]["target"] = "sluurm" + with pytest.raises(ValidationError) as excinfo: + validate_build_manifest(manifest) + assert "/deployment_config/target" in str(excinfo.value) + + def test_unknown_keys_are_allowed(self, manifest): + """A manifest may carry consumer-specific metadata madengine does not read.""" + manifest["built_models"]["img"]["consumer_note"] = "keep me" + manifest["rccl_ci"] = {"branch": "develop"} + assert validate_build_manifest(manifest) == [] + + def test_null_optional_fields_are_accepted(self, manifest): + """`madengine build` copies a null from models.json straight through. + + The schema describes what madengine writes, so rejecting its own output would only + break runs that worked: a model declaring `"timeout": null` is not a broken manifest. + """ + manifest["built_models"]["img"].update( + { + "timeout": None, + "scripts": None, + "n_gpus": None, + "args": None, + "multiple_results": None, + "tags": None, + "env_vars": None, + "slurm": None, + "distributed": None, + } + ) + manifest["built_images"]["img"]["dockerfile"] = None + manifest["context"]["docker_gpus"] = None + + assert validate_build_manifest(manifest) == [] + + def test_error_names_the_source_file(self, manifest): + manifest["deployment_config"]["distributed"]["port"] = 70000 + with pytest.raises(ValidationError) as excinfo: + validate_build_manifest(manifest, source="build_manifest.json") + assert "build_manifest.json" in str(excinfo.value) + + +class TestCrossFieldChecks: + def test_model_without_matching_image_is_fatal(self, manifest): + manifest["built_models"]["orphan"] = {"name": "orphan"} + with pytest.raises(ValidationError) as excinfo: + validate_build_manifest(manifest) + assert "orphan" in str(excinfo.value) + + def test_node_count_mismatch_is_fatal(self, manifest): + manifest["deployment_config"]["distributed"]["nnodes"] = 4 + with pytest.raises(ValidationError) as excinfo: + validate_build_manifest(manifest) + assert "world size" in str(excinfo.value) + + def test_image_without_model_is_a_warning(self, manifest): + manifest["built_images"]["spare"] = {"docker_image": "rocm/spare:latest"} + warnings = validate_build_manifest(manifest) + assert any("spare" in w for w in warnings) + + def test_slurm_target_without_slurm_block_is_a_warning(self, manifest): + del manifest["deployment_config"]["slurm"] + warnings = validate_build_manifest(manifest) + assert any("cluster defaults" in w for w in warnings) + + +class TestSingleSourceOfTruthForDeploymentTarget: + """A top-level deployment block used to select the target while its values were ignored.""" + + def test_top_level_block_is_moved_under_deployment_config(self, manifest): + top_level = {"partition": "amd-rccl", "nodes": 2} + manifest["slurm"] = top_level + del manifest["deployment_config"]["slurm"] + + warnings = validate_build_manifest(manifest) + + assert "slurm" not in manifest + assert manifest["deployment_config"]["slurm"] == top_level + assert any("belongs under" in w for w in warnings) + + def test_deployment_config_wins_when_both_are_present(self, manifest): + manifest["slurm"] = {"partition": "ignored-partition", "nodes": 2} + + warnings = validate_build_manifest(manifest) + + assert "slurm" not in manifest + assert manifest["deployment_config"]["slurm"]["partition"] == "meta64" + assert any("was ignored" in w for w in warnings) + + def test_migration_can_be_disabled(self, manifest): + manifest["slurm"] = {"partition": "amd-rccl", "nodes": 2} + validate_build_manifest(manifest, migrate=False) + assert "slurm" in manifest + + def test_every_deployment_block_is_migrated(self): + manifest = { + "built_images": {}, + "built_models": {}, + "context": {}, + "distributed": {"launcher": "torchrun", "nnodes": 1}, + "kubernetes": {"namespace": "mad"}, + } + migrate_top_level_deployment_blocks(manifest) + assert set(manifest["deployment_config"]) == {"distributed", "kubernetes"} diff --git a/tests/unit/test_orchestration.py b/tests/unit/test_orchestration.py index 54a11bd5..b188be10 100644 --- a/tests/unit/test_orchestration.py +++ b/tests/unit/test_orchestration.py @@ -209,6 +209,7 @@ def test_skip_after_build_calls_execute_local(self, mock_cleanup, tmp_path): "deployment_config": {"target": "local"}, "context": {}, "built_images": {}, + "built_models": {}, } ) ) @@ -253,6 +254,7 @@ def test_skip_run_only_still_calls_execute_local( "deployment_config": {"target": "local"}, "context": {}, "built_images": {}, + "built_models": {}, } ) ) @@ -286,6 +288,7 @@ def test_skip_model_run_calls_execute_local(self, mock_cleanup, tmp_path): "deployment_config": {"target": "local"}, "context": {}, "built_images": {}, + "built_models": {}, } ) ) From 5ee553349a5ad3911de354ed3bff05afbf63e42a Mon Sep 17 00:00:00 2001 From: Mikhail Kuznetsov Date: Thu, 6 Aug 2026 16:53:38 +0000 Subject: [PATCH 2/7] feat(config): give the environment and build-dir contract a manifest field A cluster run depends on host variables that say where things live: MODEL_DIR, the cache roots, MAD_DOCKER_BUILDS. The only way to supply them was to source mad.env in the shell that launched madengine, and forgetting produced failures far from the cause -- an empty MODEL_DIR resolves the run script path to nothing, and a MAD_DOCKER_BUILDS off shared storage makes every worker fail to find the image. deployment_config.env_file now names that file and madengine sources it itself, on the submit node while rendering the job and on every worker during the run. Relative paths resolve against the manifest, so a run directory stays movable. A missing file is fatal at startup instead of an empty string mid-run. Only the names of the applied variables are logged; an env file may carry tokens. MAD_DOCKER_BUILDS, until now only visible in the code, is documented alongside it. --- docs/configuration.md | 5 + docs/deployment.md | 41 ++++ src/madengine/core/env_file.py | 150 ++++++++++++++ src/madengine/deployment/base.py | 12 ++ .../orchestration/run_orchestrator.py | 14 +- tests/unit/test_env_file.py | 195 ++++++++++++++++++ 6 files changed, 416 insertions(+), 1 deletion(-) create mode 100644 src/madengine/core/env_file.py create mode 100644 tests/unit/test_env_file.py diff --git a/docs/configuration.md b/docs/configuration.md index 4831cc4f..e2f0af3e 100644 --- a/docs/configuration.md +++ b/docs/configuration.md @@ -252,6 +252,11 @@ This allows you to rebuild only changed models while maintaining references to e ### Environment Variables +Host-side variables (`MODEL_DIR`, cache roots, `MAD_DOCKER_BUILDS`) can be collected in a +shell env file that the manifest names via `deployment_config.env_file`; see +[Environment file](deployment.md#environment-file-env_file). The variables below are the +container's, and are set from the run configuration. + Pass environment variables to containers: ```json diff --git a/docs/deployment.md b/docs/deployment.md index fa03e7f5..88b3cd04 100644 --- a/docs/deployment.md +++ b/docs/deployment.md @@ -252,6 +252,47 @@ The deployment target is automatically detected from the `slurm` key in the conf See [examples/slurm-configs/](../examples/slurm-configs/) for complete examples. +### Environment file (`env_file`) + +A cluster run usually depends on a handful of host variables that say where things live — +the model directory, the cache roots, the shared image store. Instead of requiring every +operator to `source mad.env` in the shell that launches madengine, name the file in the +build manifest and madengine loads it itself, on the submit node and on every worker: + +```json +{ + "deployment_config": { + "target": "slurm", + "env_file": "mad.env" + } +} +``` + +A relative path is resolved against the manifest's directory, so a run directory holding +the manifest and its `mad.env` side by side stays movable. The file is sourced with +`bash`, exactly as `source mad.env` would, so `${VAR:-default}`, command substitution and +conditionals all work — and, like the manifest that names it, the file is trusted input. +Values in the file override what madengine inherited from the launching shell. madengine +logs the names of the variables it applied, never their values, so an env file may carry +tokens. + +A missing `env_file` is a fatal error at startup rather than a variable that silently +resolves to the empty string mid-run. + +### Shared image store (`MAD_DOCKER_BUILDS`) + +For a multi-node run every worker needs the same image. Set `MAD_DOCKER_BUILDS` to a +directory on shared storage (usually from the `env_file` above) and madengine saves the +built image there once, then loads it from the tar on workers whose local image ID differs: + +```bash +export MAD_DOCKER_BUILDS=/shared/MADstorage/docker_builds +``` + +Leaving it unset is supported but leaves image distribution to the operator: workers that +do not already have the image cannot reconcile it, and the run fails on those nodes with a +message saying so. + ### Multi-Node Training For distributed training across SLURM nodes: diff --git a/src/madengine/core/env_file.py b/src/madengine/core/env_file.py new file mode 100644 index 00000000..22281875 --- /dev/null +++ b/src/madengine/core/env_file.py @@ -0,0 +1,150 @@ +#!/usr/bin/env python3 +""" +Load a shell env file (the `mad.env` convention) into the run's environment. + +Multi-node runs depend on a set of variables that describe *where things live* on the +cluster: `MODEL_DIR`, the cache roots, and `MAD_DOCKER_BUILDS`. Until now the only way to +supply them was to `source mad.env` in the same shell before every `madengine run`, and +forgetting produced failures far from the cause — an empty `MODEL_DIR` makes the run +script path resolve to nothing, and a `MAD_DOCKER_BUILDS` that is not on shared storage +makes every worker rebuild the image or fail to find it. + +A manifest can now name the file (`deployment_config.env_file`) and madengine loads it +itself. The file is executed by `bash`, exactly as sourcing it would, so the usual shell +constructs (`${VAR:-default}`, `$(cat ~/.token)`, conditionals) behave the same — which +also means an env file is trusted input, on par with the manifest that names it. + +Copyright (c) Advanced Micro Devices, Inc. All rights reserved. +""" + +import os +import shlex +import subprocess +from pathlib import Path +from typing import Dict, Optional + +from madengine.core.errors import ErrorContext, ValidationError + +#: A shell that hangs (waiting on a prompt, say) must not hang the run. +SOURCE_TIMEOUT_SECONDS = 120 + +#: Separates the before/after environment dumps in the helper shell's output. Both dumps +#: come from the same shell, so anything bash sets on its own (COLUMNS, SHLVL, ...) +#: appears in both and is not mistaken for something the env file did. +_BOUNDARY = "__madengine_env_file_boundary__" + +#: Bash's own bookkeeping, which differs between the two dumps for reasons unrelated to +#: the file's contents. +_SHELL_BOOKKEEPING = frozenset({"_", "SHLVL", "PWD", "OLDPWD"}) + + +def load_env_file(env_file: str, base_dir: Optional[str] = None) -> Dict[str, str]: + """ + Source an env file with bash and return the variables it sets or changes. + + Args: + env_file: path to the file; relative paths resolve against `base_dir` + base_dir: directory to resolve a relative `env_file` against (typically the + manifest's directory), defaults to the current working directory + + Returns: + Dict[str, str]: variables the file added or changed, relative to the environment + madengine is running with + + Raises: + ValidationError: the file is missing, or bash failed to source it + """ + path = Path(env_file) + if not path.is_absolute() and base_dir: + path = Path(base_dir) / path + + context = ErrorContext( + operation="env_file loading", component="core.env_file", file_path=str(path) + ) + if not path.is_file(): + raise ValidationError( + f"env_file not found: {path}", + context=context, + suggestions=[ + "deployment_config.env_file is resolved relative to the manifest", + ], + ) + + # `set -a` is what makes plain `KEY=value` lines exported, matching what an operator + # gets from `source mad.env` in a shell configured the usual way. Sourcing is checked + # explicitly: a syntax error makes `.` fail but would otherwise be masked by the + # `env` that follows it. + script = ( + f"env -0; printf '%s\\0' {shlex.quote(_BOUNDARY)}; " + f"set -a; . {shlex.quote(str(path))} || exit 42; env -0" + ) + try: + result = subprocess.run( + ["bash", "-c", script], + capture_output=True, + text=True, + timeout=SOURCE_TIMEOUT_SECONDS, + check=False, + ) + except subprocess.TimeoutExpired as exc: + raise ValidationError( + f"Timed out after {SOURCE_TIMEOUT_SECONDS}s sourcing env_file: {path}", + context=context, + cause=exc, + ) from exc + + if result.returncode != 0: + detail = result.stderr.strip() or f"bash exited with {result.returncode}" + raise ValidationError( + f"Failed to source env_file {path}: {detail}", + context=context, + ) + + before_dump, _, after_dump = result.stdout.partition(f"{_BOUNDARY}\0") + before = _parse_env_dump(before_dump) + + loaded: Dict[str, str] = {} + for key, value in _parse_env_dump(after_dump).items(): + # A variable that already held this value is not something the file changed. + if key in _SHELL_BOOKKEEPING or before.get(key) == value: + continue + loaded[key] = value + return loaded + + +def _parse_env_dump(dump: str) -> Dict[str, str]: + """ + Parse NUL-delimited `env -0` output into a mapping. + + Args: + dump: raw `env -0` output + + Returns: + Dict[str, str]: the variables in the dump + """ + parsed: Dict[str, str] = {} + for entry in dump.split("\0"): + if not entry or "=" not in entry: + continue + key, value = entry.split("=", 1) + parsed[key] = value + return parsed + + +def apply_env_file(env_file: str, base_dir: Optional[str] = None) -> Dict[str, str]: + """ + Load an env file and apply it to `os.environ`, as sourcing it would. + + The file wins over the inherited environment, so behaviour matches what the operator + gets by sourcing it before the run. + + Args: + env_file: path to the file; relative paths resolve against `base_dir` + base_dir: directory to resolve a relative `env_file` against + + Returns: + Dict[str, str]: the variables that were applied + """ + loaded = load_env_file(env_file, base_dir) + os.environ.update(loaded) + return loaded diff --git a/src/madengine/deployment/base.py b/src/madengine/deployment/base.py index 9381f729..2414fb1a 100644 --- a/src/madengine/deployment/base.py +++ b/src/madengine/deployment/base.py @@ -19,6 +19,7 @@ from jinja2 import Environment, FileSystemLoader from rich.console import Console +from madengine.core.env_file import apply_env_file from madengine.schemas import validate_build_manifest @@ -157,6 +158,17 @@ def _load_manifest(self, manifest_file: str) -> Dict: for warning in validate_build_manifest(manifest, source=str(manifest_path)): self.console.print(f"[yellow]⚠ {warning}[/yellow]") + env_file = manifest.get("deployment_config", {}).get("env_file") + if env_file: + # The submit side needs these too: MAD_DOCKER_BUILDS and the cache roots are + # read while rendering the job script, before any node starts. + applied = apply_env_file(env_file, base_dir=str(manifest_path.resolve().parent)) + # Names only: an env file legitimately carries secrets. + self.console.print( + f"[cyan]Loaded env_file {env_file}: " + f"{', '.join(sorted(applied)) or '(no new variables)'}[/cyan]" + ) + return manifest # Template Method - defines workflow diff --git a/src/madengine/orchestration/run_orchestrator.py b/src/madengine/orchestration/run_orchestrator.py index 5ad2ee77..ad894c69 100644 --- a/src/madengine/orchestration/run_orchestrator.py +++ b/src/madengine/orchestration/run_orchestrator.py @@ -25,6 +25,7 @@ from madengine.core.auth import load_credentials from madengine.core.context import Context from madengine.core.dataprovider import Data +from madengine.core.env_file import apply_env_file from madengine.core.errors import ( BuildError, ConfigurationError, @@ -229,7 +230,18 @@ def execute( self.rich_console.print(f"[yellow]⚠ {warning}[/yellow]") deployment_config = manifest.get("deployment_config", {}) - + + if deployment_config.get("env_file"): + applied = apply_env_file( + deployment_config["env_file"], + base_dir=str(Path(manifest_file).resolve().parent), + ) + # Names only: an env file legitimately carries secrets (MAD_SECRETS_*). + self.rich_console.print( + f"[cyan]Loaded env_file {deployment_config['env_file']}: " + f"{', '.join(sorted(applied)) or '(no new variables)'}[/cyan]" + ) + # Update additional_context with deployment_config for deployment layer if not self.additional_context: self.additional_context = {} diff --git a/tests/unit/test_env_file.py b/tests/unit/test_env_file.py new file mode 100644 index 00000000..ee3a21ce --- /dev/null +++ b/tests/unit/test_env_file.py @@ -0,0 +1,195 @@ +#!/usr/bin/env python3 +""" +Tests for `deployment_config.env_file` loading. + +The contract these lock down is the one an operator already relies on when they run +`source mad.env` by hand: the file is real shell, the values it sets win over whatever +was inherited, and a path that does not exist stops the run at submit time instead of +surfacing as an empty `MODEL_DIR` twenty minutes into an allocation. + +Copyright (c) Advanced Micro Devices, Inc. All rights reserved. +""" + +import os + +import pytest + +from madengine.core.env_file import apply_env_file, load_env_file +from madengine.core.errors import ValidationError + + +@pytest.fixture +def env_dir(tmp_path): + """Directory holding env files, standing in for a run directory.""" + return tmp_path + + +@pytest.fixture(autouse=True) +def restore_environ(): + """Keep `apply_env_file` from leaking into the rest of the suite.""" + saved = os.environ.copy() + yield + os.environ.clear() + os.environ.update(saved) + + +class TestLoadEnvFile: + """Reading values out of an env file.""" + + def test_plain_assignments_are_returned(self, env_dir): + """`KEY=value` lines come back without needing an explicit export.""" + env_file = env_dir / "mad.env" + env_file.write_text("MODEL_DIR=/shared/models\nMAD_DOCKER_BUILDS=/shared/builds\n") + + loaded = load_env_file(str(env_file)) + + assert loaded["MODEL_DIR"] == "/shared/models" + assert loaded["MAD_DOCKER_BUILDS"] == "/shared/builds" + + def test_shell_constructs_are_evaluated(self, env_dir): + """The file is sourced, so expansion and defaults behave as in a shell.""" + env_file = env_dir / "mad.env" + env_file.write_text( + 'MAD_STORAGE=/shared/MADstorage\n' + 'MAD_DOCKER_BUILDS="$MAD_STORAGE/docker_builds"\n' + 'MAD_PARTITION="${MAD_PARTITION:-meta64}"\n' + ) + + loaded = load_env_file(str(env_file)) + + assert loaded["MAD_DOCKER_BUILDS"] == "/shared/MADstorage/docker_builds" + assert loaded["MAD_PARTITION"] == "meta64" + + def test_inherited_value_wins_over_default(self, env_dir, monkeypatch): + """A `${VAR:-default}` respects what the caller already exported.""" + monkeypatch.setenv("MAD_PARTITION", "debug") + env_file = env_dir / "mad.env" + env_file.write_text('MAD_PARTITION="${MAD_PARTITION:-meta64}"\n') + + loaded = load_env_file(str(env_file)) + + # Unchanged relative to the current environment, so nothing to report. + assert "MAD_PARTITION" not in loaded + + def test_unchanged_variables_are_not_reported(self, env_dir, monkeypatch): + """Only what the file actually changes is returned.""" + monkeypatch.setenv("MODEL_DIR", "/shared/models") + env_file = env_dir / "mad.env" + env_file.write_text("MODEL_DIR=/shared/models\nEXTRA=1\n") + + loaded = load_env_file(str(env_file)) + + assert loaded == {"EXTRA": "1"} + + def test_bash_bookkeeping_is_dropped(self, env_dir): + """`_`, `SHLVL` and friends are the subshell's, not the file's.""" + env_file = env_dir / "mad.env" + env_file.write_text("MODEL_DIR=/shared/models\n") + + loaded = load_env_file(str(env_file)) + + assert set(loaded) == {"MODEL_DIR"} + + def test_values_with_newlines_survive(self, env_dir): + """NUL-delimited output keeps multi-line values intact.""" + env_file = env_dir / "mad.env" + env_file.write_text('MAD_BANNER="line one\nline two"\n') + + loaded = load_env_file(str(env_file)) + + assert loaded["MAD_BANNER"] == "line one\nline two" + + def test_relative_path_resolves_against_base_dir(self, env_dir): + """A manifest names its env file relative to itself, not to the CWD.""" + (env_dir / "mad.env").write_text("MODEL_DIR=/shared/models\n") + + loaded = load_env_file("mad.env", base_dir=str(env_dir)) + + assert loaded["MODEL_DIR"] == "/shared/models" + + def test_absolute_path_ignores_base_dir(self, env_dir, tmp_path): + """An absolute path is taken as given.""" + env_file = env_dir / "mad.env" + env_file.write_text("MODEL_DIR=/shared/models\n") + + loaded = load_env_file(str(env_file), base_dir=str(tmp_path / "elsewhere")) + + assert loaded["MODEL_DIR"] == "/shared/models" + + def test_path_with_spaces_is_quoted(self, env_dir): + """The path reaches bash as one word.""" + directory = env_dir / "run dir" + directory.mkdir() + env_file = directory / "mad.env" + env_file.write_text("MODEL_DIR=/shared/models\n") + + loaded = load_env_file(str(env_file)) + + assert loaded["MODEL_DIR"] == "/shared/models" + + +class TestEnvFileFailures: + """A bad env file has to stop the run where it can still be explained.""" + + def test_missing_file_names_the_resolved_path(self, env_dir): + """The error shows where madengine looked, not just what the manifest said.""" + with pytest.raises(ValidationError) as exc_info: + load_env_file("absent.env", base_dir=str(env_dir)) + + assert str(env_dir / "absent.env") in str(exc_info.value) + + def test_directory_is_not_a_file(self, env_dir): + """Pointing at a directory fails the same way a missing file does.""" + with pytest.raises(ValidationError): + load_env_file(str(env_dir)) + + def test_shell_error_is_reported(self, env_dir): + """A non-zero exit from bash carries stderr into the message.""" + env_file = env_dir / "mad.env" + env_file.write_text("exit 3\n") + + with pytest.raises(ValidationError) as exc_info: + load_env_file(str(env_file)) + + assert "Failed to source" in str(exc_info.value) + + def test_syntax_error_is_reported(self, env_dir): + """Malformed shell is a fatal, named error.""" + env_file = env_dir / "mad.env" + env_file.write_text('MODEL_DIR="/unterminated\n') + + with pytest.raises(ValidationError): + load_env_file(str(env_file)) + + +class TestApplyEnvFile: + """Applying the file to the running process.""" + + def test_variables_land_in_environ(self, env_dir): + """What the file sets is visible to everything downstream.""" + env_file = env_dir / "mad.env" + env_file.write_text("MAD_DOCKER_BUILDS=/shared/builds\n") + + apply_env_file(str(env_file)) + + assert os.environ["MAD_DOCKER_BUILDS"] == "/shared/builds" + + def test_file_wins_over_inherited_value(self, env_dir, monkeypatch): + """Sourcing overwrites, and so does this.""" + monkeypatch.setenv("MAD_DOCKER_BUILDS", "/tmp/stale") + env_file = env_dir / "mad.env" + env_file.write_text("MAD_DOCKER_BUILDS=/shared/builds\n") + + apply_env_file(str(env_file)) + + assert os.environ["MAD_DOCKER_BUILDS"] == "/shared/builds" + + def test_unrelated_variables_are_left_alone(self, env_dir, monkeypatch): + """Loading an env file is additive, not a replacement of the environment.""" + monkeypatch.setenv("MAD_KEEP_ME", "yes") + env_file = env_dir / "mad.env" + env_file.write_text("MODEL_DIR=/shared/models\n") + + apply_env_file(str(env_file)) + + assert os.environ["MAD_KEEP_ME"] == "yes" From f65584c0682c12705da39d5627b479dcb72959f7 Mon Sep 17 00:00:00 2001 From: Mikhail Kuznetsov Date: Thu, 6 Aug 2026 16:53:39 +0000 Subject: [PATCH 3/7] refactor(reporting): make the result schema a contract, not a column list Three writers create perf.csv -- the container runner, the deployment aggregation path and the Kubernetes results mixin -- and each carried its own copy of the 29-column header under a comment asking the reader to keep them in sync. Nothing enforced that, and a consumer reading the CSV had no declaration to read at all. perf_csv.schema.json now declares the row: names, order, types and a description per column, with the writers reading the header from it. Columns that restate a manifest field carry a JSON pointer to it, so the result row and the build manifest are two views of one shape; a test walks those pointers and fails if a manifest field is renamed out from under the result contract. --- src/madengine/deployment/base.py | 11 +- src/madengine/deployment/k8s_results.py | 12 +- src/madengine/reporting/update_perf_csv.py | 13 +- src/madengine/schemas/__init__.py | 50 ++++- src/madengine/schemas/perf_csv.schema.json | 140 ++++++++++++++ tests/unit/test_perf_csv_schema.py | 204 +++++++++++++++++++++ 6 files changed, 402 insertions(+), 28 deletions(-) create mode 100644 src/madengine/schemas/perf_csv.schema.json create mode 100644 tests/unit/test_perf_csv_schema.py diff --git a/src/madengine/deployment/base.py b/src/madengine/deployment/base.py index 2414fb1a..be72ddb2 100644 --- a/src/madengine/deployment/base.py +++ b/src/madengine/deployment/base.py @@ -20,7 +20,7 @@ from rich.console import Console from madengine.core.env_file import apply_env_file -from madengine.schemas import validate_build_manifest +from madengine.schemas import perf_csv_header, validate_build_manifest # Regex for parsing "performance: " log lines. @@ -607,14 +607,7 @@ def _ensure_perf_csv_exists(self) -> None: perf_csv_path = Path("perf.csv") if perf_csv_path.exists(): return - standard_header = ( - "model,n_gpus,nnodes,gpus_per_node,training_precision,pipeline,args,tags," - "docker_file,base_docker,docker_sha,docker_image,git_commit,machine_name," - "deployment_type,launcher,gpu_architecture,performance,metric,relative_change," - "status,build_duration,test_duration,dataname,data_provider_type,data_size," - "data_download_duration,build_number,additional_docker_run_options" - ) - perf_csv_path.write_text(standard_header + "\n", encoding="utf-8") + perf_csv_path.write_text(perf_csv_header() + "\n", encoding="utf-8") def _write_to_perf_csv(self, perf_data: Dict[str, Any]) -> None: """ diff --git a/src/madengine/deployment/k8s_results.py b/src/madengine/deployment/k8s_results.py index 6da189b5..8567433d 100644 --- a/src/madengine/deployment/k8s_results.py +++ b/src/madengine/deployment/k8s_results.py @@ -16,6 +16,7 @@ from typing import Any, Dict, List, Optional from .common import normalize_launcher +from madengine.schemas import perf_csv_header from madengine.utils.path_utils import scripts_base_dir_from from madengine.utils.run_details import flatten_tags_in_place, get_build_number, get_pipeline @@ -48,15 +49,6 @@ def collector_pod_name(deployment_id: str) -> str: class KubernetesResultsMixin: """Results collection and performance reporting for Kubernetes deployments.""" - # Standard perf.csv header (must match container_runner.ensure_perf_csv_exists) - _PERF_CSV_HEADER = ( - "model,n_gpus,nnodes,gpus_per_node,training_precision,pipeline,args,tags," - "docker_file,base_docker,docker_sha,docker_image,git_commit,machine_name," - "deployment_type,launcher,gpu_architecture,performance,metric,relative_change," - "status,build_duration,test_duration,dataname,data_provider_type,data_size," - "data_download_duration,build_number,additional_docker_run_options" - ) - def collect_results(self, deployment_id: str) -> Dict[str, Any]: """ Enhanced results collection from K8s pods following vLLM multi-node best practices. @@ -977,7 +969,7 @@ def _ensure_perf_csv_exists(self) -> None: """Ensure perf.csv exists with standard header (same as Docker container_runner).""" perf_csv_path = Path("perf.csv") if not perf_csv_path.exists(): - perf_csv_path.write_text(self._PERF_CSV_HEADER + "\n", encoding="utf-8") + perf_csv_path.write_text(perf_csv_header() + "\n", encoding="utf-8") self.console.print("[dim]Created perf.csv with standard header[/dim]") def _build_perf_entry_from_aggregated( diff --git a/src/madengine/reporting/update_perf_csv.py b/src/madengine/reporting/update_perf_csv.py index f298efa2..d0122af4 100644 --- a/src/madengine/reporting/update_perf_csv.py +++ b/src/madengine/reporting/update_perf_csv.py @@ -14,14 +14,11 @@ # third-party imports import pandas as pd -# Standard header for perf CSV; must match ContainerRunner.ensure_perf_csv_exists() -PERF_CSV_HEADER = ( - "model,n_gpus,nnodes,gpus_per_node,training_precision,pipeline,args,tags," - "docker_file,base_docker,docker_sha,docker_image,git_commit,machine_name," - "deployment_type,launcher,gpu_architecture,performance,metric,relative_change," - "status,build_duration,test_duration,dataname,data_provider_type,data_size," - "data_download_duration,build_number,additional_docker_run_options" -) +from madengine.schemas import perf_csv_header + +# The columns are declared in schemas/perf_csv.schema.json, which every writer reads, so +# the four copies of this string can no longer drift apart. +PERF_CSV_HEADER = perf_csv_header() def df_strip_columns(df: pd.DataFrame) -> pd.DataFrame: diff --git a/src/madengine/schemas/__init__.py b/src/madengine/schemas/__init__.py index ae78c6e7..45c787da 100644 --- a/src/madengine/schemas/__init__.py +++ b/src/madengine/schemas/__init__.py @@ -13,12 +13,19 @@ is migrated into `deployment_config` with a warning, so the two used to disagree silently and now cannot. +The result side is declared the same way. `perf_csv.schema.json` gives the columns of +`perf.csv` in order, replacing the header string that each of the three writers spelled +out for itself under a comment asking the reader to keep them in sync. Columns that +restate a manifest field point at it, so the row and the manifest are two views of one +shape rather than two lists that happen to agree today. + Copyright (c) Advanced Micro Devices, Inc. All rights reserved. """ +import functools import json from pathlib import Path -from typing import Any, Dict, List, Optional +from typing import Any, Dict, List, Optional, Tuple from madengine.core.errors import ErrorContext, ValidationError @@ -34,6 +41,47 @@ def load_schema(name: str = "build_manifest.schema.json") -> Dict[str, Any]: return json.load(f) +@functools.lru_cache(maxsize=1) +def perf_csv_columns() -> Tuple[str, ...]: + """ + Column names of `perf.csv`, in order. + + Order is the declaration order in `perf_csv.schema.json`: a row is written positionally + into an existing file, so the schema owns the order as well as the names. + + Returns: + Tuple[str, ...]: the column names + """ + return tuple(load_schema("perf_csv.schema.json")["properties"]) + + +@functools.lru_cache(maxsize=1) +def perf_csv_header() -> str: + """ + The `perf.csv` header line, without a trailing newline. + + Returns: + str: comma-separated column names + """ + return ",".join(perf_csv_columns()) + + +def unknown_perf_columns(row: Dict[str, Any]) -> List[str]: + """ + Report keys of a result row that no column accepts. + + A row is written with `extrasaction="ignore"`, so a key the schema does not declare is + dropped on the floor rather than reported. Callers that care can ask. + + Args: + row: a result row keyed by column name + + Returns: + List[str]: keys that are not declared columns, sorted + """ + return sorted(set(row) - set(perf_csv_columns())) + + def _pointer(path) -> str: """Render a jsonschema error path as a JSON pointer (RFC 6901).""" return "/" + "/".join(str(part) for part in path) if path else "/" diff --git a/src/madengine/schemas/perf_csv.schema.json b/src/madengine/schemas/perf_csv.schema.json new file mode 100644 index 00000000..32a272fe --- /dev/null +++ b/src/madengine/schemas/perf_csv.schema.json @@ -0,0 +1,140 @@ +{ + "$schema": "https://json-schema.org/draft/2020-12/schema", + "$id": "https://github.com/ROCm/madengine/schemas/perf_csv.schema.json", + "title": "madengine performance result row", + "description": "Shape of one row of perf.csv. Property order is column order. Columns that restate a manifest field carry x-manifest-source, a JSON pointer into build_manifest.schema.json, so the two shapes cannot be renamed apart without a test failing.", + "type": "object", + "properties": { + "model": { + "description": "Model name as it appears in the manifest.", + "type": ["string", "null"], + "x-manifest-source": "/$defs/builtModel/properties/name" + }, + "n_gpus": { + "description": "GPUs the model asked for; -1 means all visible GPUs.", + "type": ["integer", "string", "null"], + "x-manifest-source": "/$defs/builtModel/properties/n_gpus" + }, + "nnodes": { + "description": "Nodes the run spanned.", + "type": ["integer", "string", "null"], + "x-manifest-source": "/$defs/distributed/properties/nnodes" + }, + "gpus_per_node": { + "description": "GPUs used per node.", + "type": ["integer", "string", "null"], + "x-manifest-source": "/$defs/slurm/properties/gpus_per_node" + }, + "training_precision": { + "description": "Precision the model was run at, e.g. bf16, fp8.", + "type": ["string", "null"], + "x-manifest-source": "/$defs/builtModel/properties/training_precision" + }, + "pipeline": { + "description": "Pipeline the model belongs to, from its model card.", + "type": ["string", "null"] + }, + "args": { + "description": "Arguments passed to the run script.", + "type": ["string", "null"], + "x-manifest-source": "/$defs/builtModel/properties/args" + }, + "tags": { + "description": "Model card tags, flattened to a comma-separated string.", + "type": ["string", "null"], + "x-manifest-source": "/$defs/builtModel/properties/tags" + }, + "docker_file": { + "description": "Dockerfile the image was built from.", + "type": ["string", "null"], + "x-manifest-source": "/$defs/builtImage/properties/dockerfile" + }, + "base_docker": { + "description": "Base image the Dockerfile started from.", + "type": ["string", "null"], + "x-manifest-source": "/$defs/builtImage/properties/base_docker" + }, + "docker_sha": { + "description": "Image digest, as reported by the container runtime at run time.", + "type": ["string", "null"] + }, + "docker_image": { + "description": "Image tag the run used.", + "type": ["string", "null"], + "x-manifest-source": "/$defs/builtImage/properties/docker_image" + }, + "git_commit": { + "description": "Commit of the model repository the run script came from.", + "type": ["string", "null"] + }, + "machine_name": { + "description": "Host the run reported from; the head node for a multi-node run.", + "type": ["string", "null"] + }, + "deployment_type": { + "description": "Where the run was deployed.", + "type": ["string", "null"], + "x-manifest-source": "/$defs/deploymentConfig/properties/target" + }, + "launcher": { + "description": "Distributed launcher, e.g. torchrun, mpirun.", + "type": ["string", "null"], + "x-manifest-source": "/$defs/distributed/properties/launcher" + }, + "gpu_architecture": { + "description": "GPU architecture detected at run time, e.g. gfx950.", + "type": ["string", "null"] + }, + "performance": { + "description": "The metric value. Empty when the workload produced no metric.", + "type": ["number", "string", "null"] + }, + "metric": { + "description": "Unit of the performance value, e.g. tokens/s/GPU.", + "type": ["string", "null"] + }, + "relative_change": { + "description": "Change against the recorded baseline, when one exists.", + "type": ["number", "string", "null"] + }, + "status": { + "description": "Outcome of the run. SUCCESS only when the workload succeeded and a metric was collected.", + "type": ["string", "null"] + }, + "build_duration": { + "description": "Seconds spent building the image.", + "type": ["number", "string", "null"], + "x-manifest-source": "/$defs/builtImage/properties/build_duration" + }, + "test_duration": { + "description": "Seconds spent running the workload.", + "type": ["number", "string", "null"] + }, + "dataname": { + "description": "Dataset the run used.", + "type": ["string", "null"] + }, + "data_provider_type": { + "description": "How the dataset was made available, e.g. nas, minio, custom.", + "type": ["string", "null"] + }, + "data_size": { + "description": "Dataset size as reported by the provider.", + "type": ["number", "string", "null"] + }, + "data_download_duration": { + "description": "Seconds spent fetching the dataset.", + "type": ["number", "string", "null"] + }, + "build_number": { + "description": "CI build number, when the run came from CI.", + "type": ["integer", "string", "null"] + }, + "additional_docker_run_options": { + "description": "Extra options passed to docker run.", + "type": ["string", "null"], + "x-manifest-source": "/$defs/builtModel/properties/additional_docker_run_options" + } + }, + "additionalProperties": true +} diff --git a/tests/unit/test_perf_csv_schema.py b/tests/unit/test_perf_csv_schema.py new file mode 100644 index 00000000..1336019a --- /dev/null +++ b/tests/unit/test_perf_csv_schema.py @@ -0,0 +1,204 @@ +#!/usr/bin/env python3 +""" +Tests for the declared shape of a perf.csv row. + +Three writers create perf.csv — the container runner, the SLURM/base deployment path and +the Kubernetes results mixin — and each used to carry its own copy of the header under a +comment asking the reader to keep them in sync. These tests hold the writers to the +schema, and hold the schema to the manifest it borrows field names from. + +Copyright (c) Advanced Micro Devices, Inc. All rights reserved. +""" + +import json + +import pytest + +from madengine.schemas import ( + load_schema, + perf_csv_columns, + perf_csv_header, + unknown_perf_columns, +) + + +@pytest.fixture(scope="module") +def perf_schema(): + """The declared result-row shape.""" + return load_schema("perf_csv.schema.json") + + +@pytest.fixture(scope="module") +def manifest_schema(): + """The declared manifest shape the result row borrows names from.""" + return load_schema("build_manifest.schema.json") + + +def _resolve_pointer(document, pointer): + """Resolve an RFC 6901 JSON pointer, returning None when it does not exist.""" + node = document + for part in pointer.lstrip("/").split("/"): + part = part.replace("~1", "/").replace("~0", "~") + if not isinstance(node, dict) or part not in node: + return None + node = node[part] + return node + + +class TestPerfCsvColumns: + """The column list itself.""" + + def test_header_is_the_columns_in_order(self): + """The header line is the declaration order, not an independently kept string.""" + assert perf_csv_header() == ",".join(perf_csv_columns()) + + def test_columns_match_the_published_order(self): + """Rows are appended positionally into existing files; the order is the contract.""" + assert perf_csv_columns() == ( + "model", + "n_gpus", + "nnodes", + "gpus_per_node", + "training_precision", + "pipeline", + "args", + "tags", + "docker_file", + "base_docker", + "docker_sha", + "docker_image", + "git_commit", + "machine_name", + "deployment_type", + "launcher", + "gpu_architecture", + "performance", + "metric", + "relative_change", + "status", + "build_duration", + "test_duration", + "dataname", + "data_provider_type", + "data_size", + "data_download_duration", + "build_number", + "additional_docker_run_options", + ) + + def test_every_column_is_documented(self, perf_schema): + """A column nobody can explain is a column nobody can consume.""" + undocumented = [ + name + for name, spec in perf_schema["properties"].items() + if not spec.get("description") + ] + assert undocumented == [] + + def test_no_duplicate_columns(self): + """JSON object keys make this hard to get wrong; say so anyway.""" + assert len(perf_csv_columns()) == len(set(perf_csv_columns())) + + def test_schema_is_valid(self, perf_schema): + """The schema itself has to be a schema.""" + import jsonschema + + jsonschema.Draft202012Validator.check_schema(perf_schema) + + +class TestManifestLinkage: + """Columns that restate a manifest field say which one.""" + + def test_manifest_pointers_resolve(self, perf_schema, manifest_schema): + """Renaming a manifest field without updating the result contract fails here.""" + dangling = { + column: spec["x-manifest-source"] + for column, spec in perf_schema["properties"].items() + if "x-manifest-source" in spec + and _resolve_pointer(manifest_schema, spec["x-manifest-source"]) is None + } + assert dangling == {} + + def test_linked_columns_cover_the_obvious_ones(self, perf_schema): + """The fields an operator reads off a manifest are the ones that must stay linked.""" + linked = { + column + for column, spec in perf_schema["properties"].items() + if "x-manifest-source" in spec + } + assert {"model", "docker_image", "nnodes", "launcher", "deployment_type"} <= linked + + def test_pointers_reach_a_declared_property(self, perf_schema, manifest_schema): + """A pointer must land on a property definition, not on some intermediate node.""" + for column, spec in perf_schema["properties"].items(): + pointer = spec.get("x-manifest-source") + if pointer is None: + continue + target = _resolve_pointer(manifest_schema, pointer) + assert isinstance(target, dict), column + assert "type" in target or "$ref" in target, column + + +class TestWritersUseTheSchema: + """No writer keeps its own copy of the header.""" + + def test_reporting_header_comes_from_the_schema(self): + """update_perf_csv exports the header other code imports.""" + from madengine.reporting.update_perf_csv import PERF_CSV_HEADER + + assert PERF_CSV_HEADER == perf_csv_header() + + def test_deployment_writes_the_schema_header(self, tmp_path, monkeypatch): + """The SLURM/base aggregation path creates perf.csv with the declared columns.""" + from madengine.deployment.base import BaseDeployment + + monkeypatch.chdir(tmp_path) + # The method touches no instance state, and BaseDeployment is abstract. + BaseDeployment._ensure_perf_csv_exists(None) + + assert (tmp_path / "perf.csv").read_text().strip() == perf_csv_header() + + def test_existing_file_is_left_alone(self, tmp_path, monkeypatch): + """A perf.csv from an earlier run keeps its own column order.""" + from madengine.deployment.base import BaseDeployment + + monkeypatch.chdir(tmp_path) + (tmp_path / "perf.csv").write_text("model,performance\n") + BaseDeployment._ensure_perf_csv_exists(None) + + assert (tmp_path / "perf.csv").read_text() == "model,performance\n" + + def test_no_writer_hardcodes_the_column_list(self): + """The header string should exist once, in the schema.""" + from pathlib import Path + + import madengine + + src = Path(madengine.__file__).parent + needle = "data_download_duration,build_number" + offenders = [ + str(path.relative_to(src)) + for path in src.rglob("*.py") + if needle in path.read_text(encoding="utf-8", errors="replace") + ] + assert offenders == [] + + +class TestUnknownColumns: + """Rows are written with extrasaction='ignore'; callers can still ask what was dropped.""" + + def test_declared_keys_are_not_reported(self): + """A row made of columns reports nothing.""" + row = {name: "" for name in perf_csv_columns()} + + assert unknown_perf_columns(row) == [] + + def test_undeclared_keys_are_reported(self): + """A typo'd key would otherwise vanish silently into the CSV writer.""" + row = {"model": "llama", "perfromance": 1.0, "nnodes": 2} + + assert unknown_perf_columns(row) == ["perfromance"] + + def test_partial_rows_are_fine(self): + """Missing columns are normal: not every run has data provider fields.""" + assert unknown_perf_columns({"model": "llama"}) == [] From a4922e9ddb422d30c3e1ab267e1667f650be7d08 Mon Sep 17 00:00:00 2001 From: Mikhail Kuznetsov Date: Thu, 6 Aug 2026 16:54:04 +0000 Subject: [PATCH 4/7] feat(slurm): make the per-cluster presets record the cluster's facts The SLURM presets describe how big a job is and for how long. Nothing described the cluster it lands on, so cluster facts were baked into the shape: the multi-node preset sets NCCL_IB_DISABLE=1 and NCCL_SOCKET_IFNAME=eth0, putting every multi-node run on TCP over an interface many clusters do not have, on a fabric that is often RoCE. Fixing that per cluster meant editing a shipped preset or repeating transport variables in every manifest. slurm.cluster_profile now names one or more fact profiles, merged after the shape presets and before the user, so a profile can correct a bad default and the user can still correct the profile. Facts are orthogonal, so several profiles may be named -- a fabric and "this scheduler advertises no GPU GRES" are separate statements -- and a null value removes a variable an earlier layer set, because an interface name that does not exist here is worse than none. Bundled profiles name hardware archetypes (Broadcom Thor2 and ConnectX-7 RoCE, Mellanox InfiniBand, plain TCP, no-GRES schedulers). Anything site-specific belongs in a site profile, found by path or via MADENGINE_CLUSTER_PROFILES. An unknown profile stops the run: a silent fallback to the wrong fabric costs more than a startup error. The profile shape is declared in cluster_profile.schema.json and validated on load, and the .gitignore exemption is extended so preset JSON reaches the wheel. --- .gitignore | 1 + docs/deployment.md | 45 +++ src/madengine/deployment/config_loader.py | 22 +- .../deployment/presets/cluster_profiles.py | 221 +++++++++++++++ .../deployment/presets/slurm/__init__.py | 5 +- .../presets/slurm/clusters/ethernet-tcp.json | 14 + .../slurm/clusters/infiniband-mellanox.json | 16 ++ .../presets/slurm/clusters/no-gpu-gres.json | 8 + .../slurm/clusters/roce-broadcom-thor2.json | 16 ++ .../slurm/clusters/roce-mellanox-cx7.json | 16 ++ .../schemas/build_manifest.schema.json | 5 +- .../schemas/cluster_profile.schema.json | 58 ++++ tests/unit/test_cluster_profiles.py | 266 ++++++++++++++++++ 13 files changed, 685 insertions(+), 8 deletions(-) create mode 100644 src/madengine/deployment/presets/cluster_profiles.py create mode 100644 src/madengine/deployment/presets/slurm/clusters/ethernet-tcp.json create mode 100644 src/madengine/deployment/presets/slurm/clusters/infiniband-mellanox.json create mode 100644 src/madengine/deployment/presets/slurm/clusters/no-gpu-gres.json create mode 100644 src/madengine/deployment/presets/slurm/clusters/roce-broadcom-thor2.json create mode 100644 src/madengine/deployment/presets/slurm/clusters/roce-mellanox-cx7.json create mode 100644 src/madengine/schemas/cluster_profile.schema.json create mode 100644 tests/unit/test_cluster_profiles.py diff --git a/.gitignore b/.gitignore index 3ab73834..a038d685 100644 --- a/.gitignore +++ b/.gitignore @@ -135,6 +135,7 @@ scripts/ # for the model JSONs that land in a dev checkout; these are package data, and hatchling # picks files by VCS status, so an ignored schema would also be missing from the wheel. !src/madengine/schemas/*.json +!src/madengine/deployment/presets/**/*.json .*_env/ .vscode/ diff --git a/docs/deployment.md b/docs/deployment.md index 88b3cd04..858bfa3c 100644 --- a/docs/deployment.md +++ b/docs/deployment.md @@ -279,6 +279,51 @@ tokens. A missing `env_file` is a fatal error at startup rather than a variable that silently resolves to the empty string mid-run. +### Cluster profiles (`cluster_profile`) + +The SLURM presets describe the shape of a job — how many nodes, for how long. What they +cannot describe is the cluster it lands on, so cluster facts used to be baked into the +shape: the built-in multi-node preset sets `NCCL_IB_DISABLE=1` and +`NCCL_SOCKET_IFNAME=eth0`, which puts a run on TCP over an interface that may not exist, +on a fabric that may well be RoCE. + +A cluster profile carries those facts separately: + +```json +{ + "slurm": { + "nodes": 2, + "cluster_profile": ["roce-broadcom-thor2", "no-gpu-gres"] + } +} +``` + +Profiles merge after the shape presets and before your own configuration, so a profile +fixes a bad default and you can still override the profile. Several may be named because +the facts are orthogonal: the fabric and whether the scheduler advertises GPU GRES are +separate statements. A value of `null` in a profile removes a variable an earlier layer +set — an interface name that does not exist on this cluster is worse than none at all. + +Bundled archetypes (`madengine/deployment/presets/slurm/clusters/`): + +| Profile | What it asserts | +| --- | --- | +| `roce-broadcom-thor2` | RoCEv2 over `bnxt_re` devices | +| `roce-mellanox-cx7` | RoCEv2 over `mlx5` devices | +| `infiniband-mellanox` | InfiniBand over `mlx5` devices | +| `ethernet-tcp` | No RDMA fabric; collectives over TCP | +| `no-gpu-gres` | `sbatch` rejects GPU directives (`GresTypes=(null)`) | + +Archetypes name hardware, never a particular cluster: anything site-specific — the +partition, the account, the management interface — belongs in a site profile. Write one as +a JSON file with the same shape, then name it by path, or by name after pointing +`MADENGINE_CLUSTER_PROFILES` at the directory holding it. A site profile shadows a bundled +one of the same name. `MADENGINE_CLUSTER_PROFILE` selects a profile for runs whose manifest +names none. + +A profile that does not exist stops the run: a silent fallback to the wrong fabric costs +more than a startup error. + ### Shared image store (`MAD_DOCKER_BUILDS`) For a multi-node run every worker needs the same image. Set `MAD_DOCKER_BUILDS` to a diff --git a/src/madengine/deployment/config_loader.py b/src/madengine/deployment/config_loader.py index 06d8a1b1..b7c4cdf1 100644 --- a/src/madengine/deployment/config_loader.py +++ b/src/madengine/deployment/config_loader.py @@ -15,6 +15,11 @@ from typing import Any, Callable, Dict, Optional from copy import deepcopy +from madengine.deployment.presets.cluster_profiles import ( + apply_cluster_profiles, + selected_profiles, +) + def apply_deployment_config(config: Any, load_fn: Callable[[Dict[str, Any]], Dict[str, Any]]) -> Dict[str, Any]: """Apply deployment defaults via a loader and set config.additional_context. @@ -192,14 +197,18 @@ def load_slurm_config(cls, user_config: Dict[str, Any]) -> Dict[str, Any]: Layers: 1. Base SLURM defaults - 2. Profile preset (single-node/multi-node) - 3. User configuration (already merged from file + CLI) + 2. Profile preset (single-node/multi-node) — the shape of the job + 3. Cluster profile(s) — facts about the cluster it lands on + 4. User configuration (already merged from file + CLI) Args: user_config: User-provided configuration Returns: Complete configuration with defaults applied + + Raises: + ValidationError: A named cluster profile is missing or invalid """ # Layer 1: Base defaults config = cls.load_preset("slurm/defaults.json") @@ -218,8 +227,13 @@ def load_slurm_config(cls, user_config: Dict[str, Any]) -> Dict[str, Any]: else: profile_preset = cls.load_preset("slurm/profiles/single-node.json") config = cls.deep_merge(config, profile_preset) - - # Layer 3: User configuration (highest priority) + + # Layer 3: cluster facts. These come after the shape presets, which carry transport + # defaults that only fit one kind of cluster, and before the user, who is always + # entitled to override a fact we got wrong. + config = apply_cluster_profiles(config, selected_profiles(temp_config)) + + # Layer 4: User configuration (highest priority) config = cls.deep_merge(config, user_config) return config diff --git a/src/madengine/deployment/presets/cluster_profiles.py b/src/madengine/deployment/presets/cluster_profiles.py new file mode 100644 index 00000000..36a27079 --- /dev/null +++ b/src/madengine/deployment/presets/cluster_profiles.py @@ -0,0 +1,221 @@ +#!/usr/bin/env python3 +""" +Per-cluster fact profiles for SLURM deployments. + +The existing SLURM presets describe the *shape* of a job — how many nodes, how long, which +launcher. What they cannot describe is the cluster the job lands on, so cluster facts ended +up baked into the shape: `slurm/profiles/multi-node.json` sets `NCCL_IB_DISABLE=1` and +`NCCL_SOCKET_IFNAME=eth0`, which quietly puts every multi-node run on TCP over an interface +that may not exist, on a fabric that may well be RoCE. + +A cluster profile carries the facts instead: whether the scheduler advertises GPU GRES, +which NICs carry the collectives, the transport variables that make RDMA work, and the node +facts (GPU vendor, GPUs per node, architecture) that a submit node without GPUs cannot +discover for itself. Profiles are selected by `slurm.cluster_profile`, merge after the shape +profile and before the user's own configuration, and a site can keep its own profile outside +the repository — a bundled profile names a hardware archetype, never someone's cluster. + +Copyright (c) Advanced Micro Devices, Inc. All rights reserved. +""" + +import json +import os +from pathlib import Path +from typing import Any, Dict, List, Optional, Sequence, Union + +from madengine.core.errors import ErrorContext, ValidationError + +CLUSTER_DIR = Path(__file__).parent / "slurm" / "clusters" + +#: Colon-separated directories holding site profiles, searched before the bundled ones. +PROFILE_PATH_ENV = "MADENGINE_CLUSTER_PROFILES" + +#: Set when no profile is named in the configuration. +PROFILE_NAME_ENV = "MADENGINE_CLUSTER_PROFILE" + + +def available_profiles() -> List[str]: + """ + Names of the profiles that can be selected, from the search path and the bundle. + + Returns: + List[str]: profile names, sorted, without the .json suffix + """ + names = set() + for directory in _search_dirs(): + if directory.is_dir(): + names.update(path.stem for path in directory.glob("*.json")) + return sorted(names) + + +def _search_dirs() -> List[Path]: + """Directories to look in, site-first so a site can shadow a bundled archetype.""" + dirs = [Path(p) for p in os.environ.get(PROFILE_PATH_ENV, "").split(os.pathsep) if p] + dirs.append(CLUSTER_DIR) + return dirs + + +def resolve_profile_path(name: str) -> Path: + """ + Find the file backing a profile reference. + + A reference that looks like a path (contains a separator or ends in `.json`) is taken as + one, so a site can point at a profile it keeps next to its manifests. Anything else is a + name looked up in the search path and then in the bundled archetypes. + + Args: + name: profile name or path + + Returns: + Path: the profile file + + Raises: + ValidationError: no such profile + """ + context = ErrorContext( + operation="cluster profile lookup", component="deployment.presets" + ) + + if os.sep in name or name.endswith(".json"): + path = Path(name).expanduser() + if not path.is_file(): + raise ValidationError( + f"Cluster profile not found: {path}", + context=context, + suggestions=[f"Bundled profiles: {', '.join(available_profiles())}"], + ) + return path + + for directory in _search_dirs(): + candidate = directory / f"{name}.json" + if candidate.is_file(): + return candidate + + raise ValidationError( + f"Unknown cluster profile: {name}", + context=context, + suggestions=[ + f"Available profiles: {', '.join(available_profiles())}", + f"Point {PROFILE_PATH_ENV} at a directory of site profiles, or give a path", + ], + ) + + +def load_profile(name: str) -> Dict[str, Any]: + """ + Load and validate one cluster profile. + + Args: + name: profile name or path + + Returns: + Dict[str, Any]: the profile + + Raises: + ValidationError: the profile is missing, unparseable, or does not match the schema + """ + import jsonschema + + from madengine.schemas import load_schema + + path = resolve_profile_path(name) + context = ErrorContext( + operation="cluster profile loading", + component="deployment.presets", + file_path=str(path), + ) + + try: + with open(path) as f: + profile = json.load(f) + except json.JSONDecodeError as exc: + raise ValidationError( + f"Cluster profile {path} is not valid JSON: {exc}", context=context, cause=exc + ) from exc + + validator = jsonschema.Draft202012Validator(load_schema("cluster_profile.schema.json")) + first = next(iter(sorted(validator.iter_errors(profile), key=lambda e: list(e.path))), None) + if first is not None: + pointer = "/" + "/".join(str(part) for part in first.absolute_path) + raise ValidationError( + f"Invalid cluster profile {path}: {pointer}: {first.message}", context=context + ) + + return profile + + +def _merge(base: Dict[str, Any], profile: Dict[str, Any]) -> Dict[str, Any]: + """ + Merge one profile over a configuration. + + Nested dicts merge; a null value removes the key, which is how a profile says a variable + inherited from a shape preset does not apply here — an interface name that does not exist + on this cluster is worse than no interface name at all. + + Args: + base: configuration so far + profile: profile to apply + + Returns: + Dict[str, Any]: the merged configuration + """ + result = dict(base) + for key, value in profile.items(): + if value is None: + result.pop(key, None) + elif isinstance(value, dict) and isinstance(result.get(key), dict): + result[key] = _merge(result[key], value) + else: + result[key] = value + return result + + +def selected_profiles(config: Dict[str, Any]) -> List[str]: + """ + Profile references named by a configuration, in merge order. + + Args: + config: SLURM configuration, before user overrides are applied + + Returns: + List[str]: profile names or paths; empty when none is selected + """ + selection: Union[str, Sequence[str], None] = (config.get("slurm") or {}).get( + "cluster_profile" + ) + if not selection: + selection = os.environ.get(PROFILE_NAME_ENV) or None + if not selection: + return [] + if isinstance(selection, str): + return [selection] + return list(selection) + + +def apply_cluster_profiles( + config: Dict[str, Any], selection: Optional[Sequence[str]] = None +) -> Dict[str, Any]: + """ + Merge the selected cluster profiles into a configuration. + + Several profiles may be named, and they merge left to right: cluster facts are + orthogonal, so "this fabric" and "this scheduler advertises no GPU GRES" are separate + statements rather than a combinatorial set of files. + + Args: + config: configuration to merge into + selection: profile references; taken from the configuration when omitted + + Returns: + Dict[str, Any]: the configuration with profiles applied + + Raises: + ValidationError: a named profile is missing or invalid + """ + references = list(selection) if selection is not None else selected_profiles(config) + for reference in references: + profile = load_profile(reference) + # Documentation keys describe the file, not the cluster. + facts = {k: v for k, v in profile.items() if not k.startswith("_")} + config = _merge(config, facts) + return config diff --git a/src/madengine/deployment/presets/slurm/__init__.py b/src/madengine/deployment/presets/slurm/__init__.py index 9d11608c..a799c794 100644 --- a/src/madengine/deployment/presets/slurm/__init__.py +++ b/src/madengine/deployment/presets/slurm/__init__.py @@ -4,7 +4,9 @@ Layered configuration system: 1. defaults.json - Base SLURM defaults 2. profiles/*.json - Workload-specific profiles (single-node, multi-node) -3. User configuration - Highest priority +3. clusters/*.json - Facts about the cluster the job lands on, selected by + slurm.cluster_profile; see ../cluster_profiles.py +4. User configuration - Highest priority Convention over Configuration: - Presence of "slurm" field → SLURM deployment @@ -12,4 +14,3 @@ Copyright (c) Advanced Micro Devices, Inc. All rights reserved. """ - diff --git a/src/madengine/deployment/presets/slurm/clusters/ethernet-tcp.json b/src/madengine/deployment/presets/slurm/clusters/ethernet-tcp.json new file mode 100644 index 00000000..dc66a511 --- /dev/null +++ b/src/madengine/deployment/presets/slurm/clusters/ethernet-tcp.json @@ -0,0 +1,14 @@ +{ + "_description": "No RDMA fabric: collectives fall back to TCP.", + "_comment": "Name the interface in a site profile; the default that used to ship, eth0, does not exist on many clusters.", + + "facts": { + "fabric": "ethernet" + }, + + "env_vars": { + "NCCL_IB_DISABLE": "1", + "NCCL_IB_HCA": null, + "NCCL_IB_GID_INDEX": null + } +} diff --git a/src/madengine/deployment/presets/slurm/clusters/infiniband-mellanox.json b/src/madengine/deployment/presets/slurm/clusters/infiniband-mellanox.json new file mode 100644 index 00000000..f8396fcd --- /dev/null +++ b/src/madengine/deployment/presets/slurm/clusters/infiniband-mellanox.json @@ -0,0 +1,16 @@ +{ + "_description": "Collectives over InfiniBand on Mellanox HCAs (mlx5 devices).", + "_comment": "InfiniBand needs no GID index: that is a RoCE concern.", + + "facts": { + "fabric": "infiniband" + }, + + "env_vars": { + "NCCL_IB_DISABLE": "0", + "NCCL_IB_HCA": "mlx5", + "NCCL_IB_GID_INDEX": null, + "NCCL_SOCKET_IFNAME": null, + "GLOO_SOCKET_IFNAME": null + } +} diff --git a/src/madengine/deployment/presets/slurm/clusters/no-gpu-gres.json b/src/madengine/deployment/presets/slurm/clusters/no-gpu-gres.json new file mode 100644 index 00000000..3834792c --- /dev/null +++ b/src/madengine/deployment/presets/slurm/clusters/no-gpu-gres.json @@ -0,0 +1,8 @@ +{ + "_description": "Scheduler advertises no GPU GRES, so GPU directives are rejected by sbatch.", + "_comment": "A cluster where scontrol reports GresTypes=(null): nodes hand out all their GPUs with the allocation. Combine with a fabric profile, e.g. \"cluster_profile\": [\"roce-broadcom-thor2\", \"no-gpu-gres\"].", + + "slurm": { + "skip_gpus_directive": true + } +} diff --git a/src/madengine/deployment/presets/slurm/clusters/roce-broadcom-thor2.json b/src/madengine/deployment/presets/slurm/clusters/roce-broadcom-thor2.json new file mode 100644 index 00000000..85f9f109 --- /dev/null +++ b/src/madengine/deployment/presets/slurm/clusters/roce-broadcom-thor2.json @@ -0,0 +1,16 @@ +{ + "_description": "Collectives over RoCEv2 on Broadcom Thor2 NICs (bnxt_re devices).", + "_comment": "The control plane still runs over TCP; pin NCCL_SOCKET_IFNAME to the management interface in a site profile, since its name varies by cluster.", + + "facts": { + "fabric": "roce" + }, + + "env_vars": { + "NCCL_IB_DISABLE": "0", + "NCCL_IB_HCA": "bnxt_re", + "NCCL_IB_GID_INDEX": "3", + "NCCL_SOCKET_IFNAME": null, + "GLOO_SOCKET_IFNAME": null + } +} diff --git a/src/madengine/deployment/presets/slurm/clusters/roce-mellanox-cx7.json b/src/madengine/deployment/presets/slurm/clusters/roce-mellanox-cx7.json new file mode 100644 index 00000000..198d8407 --- /dev/null +++ b/src/madengine/deployment/presets/slurm/clusters/roce-mellanox-cx7.json @@ -0,0 +1,16 @@ +{ + "_description": "Collectives over RoCEv2 on Mellanox ConnectX-7 NICs (mlx5 devices).", + "_comment": "GID index 3 is the usual RoCEv2 entry with an IPv4-mapped address; confirm with show_gids on the cluster.", + + "facts": { + "fabric": "roce" + }, + + "env_vars": { + "NCCL_IB_DISABLE": "0", + "NCCL_IB_HCA": "mlx5", + "NCCL_IB_GID_INDEX": "3", + "NCCL_SOCKET_IFNAME": null, + "GLOO_SOCKET_IFNAME": null + } +} diff --git a/src/madengine/schemas/build_manifest.schema.json b/src/madengine/schemas/build_manifest.schema.json index 459aa2f2..8c16317d 100644 --- a/src/madengine/schemas/build_manifest.schema.json +++ b/src/madengine/schemas/build_manifest.schema.json @@ -130,8 +130,9 @@ "exclusive": { "type": "boolean" }, "network_interface": { "type": "string" }, "cluster_profile": { - "description": "Name of a per-cluster preset profile under deployment/presets/slurm/clusters/.", - "type": "string" + "description": "Cluster fact profile(s) to apply: a bundled name from deployment/presets/slurm/clusters/, a path to a site profile, or a list of either, merged in order.", + "type": ["string", "array"], + "items": { "type": "string" } }, "modules": { "type": "array", "items": { "type": "string" } } } diff --git a/src/madengine/schemas/cluster_profile.schema.json b/src/madengine/schemas/cluster_profile.schema.json new file mode 100644 index 00000000..dceac9ad --- /dev/null +++ b/src/madengine/schemas/cluster_profile.schema.json @@ -0,0 +1,58 @@ +{ + "$schema": "https://json-schema.org/draft/2020-12/schema", + "$id": "https://github.com/ROCm/madengine/schemas/cluster_profile.schema.json", + "title": "madengine cluster profile", + "description": "Facts about a cluster, as opposed to the shape of a job. Merged after the SLURM shape presets and before the user's own configuration. A null value removes an inherited key.", + "type": "object", + "properties": { + "_description": { "type": "string" }, + "_comment": { "type": "string" }, + "facts": { + "description": "Node facts a submit node cannot always discover for itself, e.g. when the login node has no GPUs.", + "type": "object", + "properties": { + "gpu_vendor": { "type": "string", "enum": ["AMD", "NVIDIA"] }, + "gpus_per_node": { "type": "integer", "minimum": 1 }, + "gpu_architecture": { + "description": "e.g. gfx942; used in place of probing when the submit node has no GPU.", + "type": "string" + }, + "gpu_renderer": { "type": "string" }, + "shared_filesystem": { + "description": "Filesystem types df reports for shared storage on this cluster.", + "type": "array", + "items": { "type": "string" } + }, + "fabric": { + "description": "Interconnect the collectives run over, e.g. roce, infiniband, ethernet.", + "type": "string" + } + }, + "additionalProperties": true + }, + "slurm": { + "description": "Scheduler facts, e.g. whether GPU GRES is advertised. Site policy such as partition or account belongs in the user's configuration, not in a shared profile.", + "type": "object", + "properties": { + "skip_gpus_directive": { "type": ["boolean", "null"] }, + "gpus_per_node": { "type": ["integer", "null"], "minimum": 1 }, + "network_interface": { "type": ["string", "null"] }, + "modules": { + "type": ["array", "null"], + "items": { "type": "string" } + } + }, + "additionalProperties": true + }, + "distributed": { + "type": "object", + "additionalProperties": true + }, + "env_vars": { + "description": "Transport and runtime variables. Values are rendered into a shell command, so they must be scalars; null removes a variable a shape preset set.", + "type": "object", + "additionalProperties": { "type": ["string", "number", "boolean", "null"] } + } + }, + "additionalProperties": false +} diff --git a/tests/unit/test_cluster_profiles.py b/tests/unit/test_cluster_profiles.py new file mode 100644 index 00000000..11941eda --- /dev/null +++ b/tests/unit/test_cluster_profiles.py @@ -0,0 +1,266 @@ +#!/usr/bin/env python3 +""" +Tests for per-cluster fact profiles. + +The behaviour under test is the separation the profiles exist for: the shape presets say +how big the job is, the cluster profile says what the cluster is, and the user has the last +word on both. The regression that motivated this is concrete — the shipped multi-node preset +puts every run on TCP over eth0, which is wrong on any RoCE cluster and on any cluster whose +management interface is not called eth0. + +Copyright (c) Advanced Micro Devices, Inc. All rights reserved. +""" + +import json + +import pytest + +from madengine.core.errors import ValidationError +from madengine.deployment.config_loader import ConfigLoader +from madengine.deployment.presets.cluster_profiles import ( + CLUSTER_DIR, + PROFILE_NAME_ENV, + PROFILE_PATH_ENV, + apply_cluster_profiles, + available_profiles, + load_profile, + resolve_profile_path, + selected_profiles, +) + + +@pytest.fixture +def site_profiles(tmp_path, monkeypatch): + """A directory of site profiles on the search path.""" + monkeypatch.setenv(PROFILE_PATH_ENV, str(tmp_path)) + return tmp_path + + +def write_profile(directory, name, body): + """Write a profile file and return its path.""" + path = directory / f"{name}.json" + path.write_text(json.dumps(body)) + return path + + +class TestBundledProfiles: + """The archetypes madengine ships.""" + + def test_bundled_profiles_are_discoverable(self): + """A user who mistypes a name gets the list, so the list has to be real.""" + assert "roce-broadcom-thor2" in available_profiles() + assert "no-gpu-gres" in available_profiles() + + @pytest.mark.parametrize( + "name", sorted(path.stem for path in CLUSTER_DIR.glob("*.json")) + ) + def test_bundled_profile_matches_the_schema(self, name): + """Every shipped profile validates; a typo here breaks a cluster, not a test.""" + assert load_profile(name) + + def test_roce_profile_enables_rdma(self): + """The point of the RoCE archetype is that IB is not disabled.""" + profile = load_profile("roce-broadcom-thor2") + + assert profile["env_vars"]["NCCL_IB_DISABLE"] == "0" + assert profile["env_vars"]["NCCL_IB_HCA"] == "bnxt_re" + + def test_no_gres_profile_is_scheduler_only(self): + """Whether GRES exists says nothing about the fabric, so the profile says nothing.""" + profile = load_profile("no-gpu-gres") + + assert profile["slurm"]["skip_gpus_directive"] is True + assert "env_vars" not in profile + + def test_bundled_profiles_name_no_cluster(self): + """Archetypes describe hardware; a site's own cluster stays out of the repository.""" + for path in CLUSTER_DIR.glob("*.json"): + body = path.read_text() + assert "partition" not in body, path.name + assert "account" not in body, path.name + + +class TestProfileResolution: + """Finding the file behind a reference.""" + + def test_site_directory_is_searched_first(self, site_profiles): + """A site can shadow an archetype it disagrees with.""" + override = write_profile(site_profiles, "roce-broadcom-thor2", {"facts": {"fabric": "roce"}}) + + assert resolve_profile_path("roce-broadcom-thor2") == override + + def test_path_reference_is_taken_as_a_path(self, tmp_path): + """A profile kept next to the manifests needs no installation.""" + path = write_profile(tmp_path, "our-cluster", {"facts": {"gpus_per_node": 8}}) + + assert resolve_profile_path(str(path)) == path + + def test_unknown_name_lists_what_exists(self): + """The error a mistyped name produces should be the answer to it.""" + with pytest.raises(ValidationError) as exc_info: + load_profile("no-such-cluster") + + message = str(exc_info.value) + assert "no-such-cluster" in message + + def test_missing_path_is_reported(self, tmp_path): + """A path that does not exist fails as a path, not as an unknown name.""" + with pytest.raises(ValidationError) as exc_info: + load_profile(str(tmp_path / "absent.json")) + + assert "absent.json" in str(exc_info.value) + + def test_malformed_profile_is_reported(self, site_profiles): + """Broken JSON names the file it could not parse.""" + (site_profiles / "broken.json").write_text("{not json") + + with pytest.raises(ValidationError) as exc_info: + load_profile("broken") + + assert "broken.json" in str(exc_info.value) + + def test_profile_violating_the_schema_is_rejected(self, site_profiles): + """A profile is validated before it can misconfigure a run.""" + write_profile(site_profiles, "bad", {"env_vars": {"NCCL_IB_HCA": ["mlx5"]}}) + + with pytest.raises(ValidationError) as exc_info: + load_profile("bad") + + assert "NCCL_IB_HCA" in str(exc_info.value) + + def test_unknown_top_level_key_is_rejected(self, site_profiles): + """Facts land where the loader looks for them, or not at all.""" + write_profile(site_profiles, "bad", {"enviroment": {"NCCL_IB_HCA": "mlx5"}}) + + with pytest.raises(ValidationError): + load_profile("bad") + + +class TestSelection: + """Which profiles a configuration asks for.""" + + def test_single_name(self): + """The common case.""" + config = {"slurm": {"cluster_profile": "ethernet-tcp"}} + + assert selected_profiles(config) == ["ethernet-tcp"] + + def test_list_of_names(self): + """Orthogonal facts are separate profiles rather than a combinatorial file set.""" + config = {"slurm": {"cluster_profile": ["roce-broadcom-thor2", "no-gpu-gres"]}} + + assert selected_profiles(config) == ["roce-broadcom-thor2", "no-gpu-gres"] + + def test_environment_names_a_profile(self, monkeypatch): + """A site can point every run at its profile without touching manifests.""" + monkeypatch.setenv(PROFILE_NAME_ENV, "ethernet-tcp") + + assert selected_profiles({"slurm": {}}) == ["ethernet-tcp"] + + def test_manifest_wins_over_environment(self, monkeypatch): + """An explicit choice in the manifest is not overridden by the environment.""" + monkeypatch.setenv(PROFILE_NAME_ENV, "ethernet-tcp") + config = {"slurm": {"cluster_profile": "infiniband-mellanox"}} + + assert selected_profiles(config) == ["infiniband-mellanox"] + + def test_nothing_selected(self): + """Profiles are opt-in; without one, nothing changes.""" + assert selected_profiles({"slurm": {}}) == [] + + +class TestMerging: + """How facts land on a configuration.""" + + def test_profile_overrides_shape_preset(self, site_profiles): + """This is the regression: a RoCE cluster must not inherit NCCL_IB_DISABLE=1.""" + write_profile(site_profiles, "ours", {"env_vars": {"NCCL_IB_DISABLE": "0"}}) + config = {"env_vars": {"NCCL_IB_DISABLE": "1", "NCCL_DEBUG": "WARN"}} + + merged = apply_cluster_profiles(config, ["ours"]) + + assert merged["env_vars"]["NCCL_IB_DISABLE"] == "0" + assert merged["env_vars"]["NCCL_DEBUG"] == "WARN" + + def test_null_removes_an_inherited_variable(self, site_profiles): + """An interface name that does not exist here is worse than none.""" + write_profile(site_profiles, "ours", {"env_vars": {"NCCL_SOCKET_IFNAME": None}}) + config = {"env_vars": {"NCCL_SOCKET_IFNAME": "eth0"}} + + merged = apply_cluster_profiles(config, ["ours"]) + + assert "NCCL_SOCKET_IFNAME" not in merged["env_vars"] + + def test_profiles_merge_left_to_right(self, site_profiles): + """Later profiles refine earlier ones.""" + write_profile(site_profiles, "fabric", {"env_vars": {"NCCL_IB_HCA": "mlx5"}}) + write_profile(site_profiles, "site", {"env_vars": {"NCCL_IB_HCA": "mlx5_0"}}) + + merged = apply_cluster_profiles({}, ["fabric", "site"]) + + assert merged["env_vars"]["NCCL_IB_HCA"] == "mlx5_0" + + def test_documentation_keys_do_not_leak_into_config(self, site_profiles): + """`_description` describes the file, not the cluster.""" + write_profile(site_profiles, "ours", {"_description": "ours", "facts": {"fabric": "roce"}}) + + merged = apply_cluster_profiles({}, ["ours"]) + + assert "_description" not in merged + + def test_input_configuration_is_not_mutated(self, site_profiles): + """Merging returns a new configuration; callers keep theirs.""" + write_profile(site_profiles, "ours", {"env_vars": {"NCCL_IB_DISABLE": "0"}}) + config = {"env_vars": {"NCCL_IB_DISABLE": "1"}} + + apply_cluster_profiles(config, ["ours"]) + + assert config["env_vars"]["NCCL_IB_DISABLE"] == "1" + + +class TestConfigLoaderIntegration: + """The layer as the deployment path sees it.""" + + def test_cluster_facts_beat_the_multi_node_preset(self): + """A two-node RoCE run keeps RDMA on, which the shipped preset would have disabled.""" + config = ConfigLoader.load_slurm_config( + {"slurm": {"nodes": 2, "cluster_profile": "roce-broadcom-thor2"}} + ) + + assert config["env_vars"]["NCCL_IB_DISABLE"] == "0" + assert config["env_vars"]["NCCL_IB_HCA"] == "bnxt_re" + assert "NCCL_SOCKET_IFNAME" not in config["env_vars"] + + def test_user_configuration_still_wins(self): + """A profile is a default, not a policy.""" + config = ConfigLoader.load_slurm_config( + { + "slurm": {"nodes": 2, "cluster_profile": "roce-broadcom-thor2"}, + "env_vars": {"NCCL_IB_HCA": "bnxt_re0"}, + } + ) + + assert config["env_vars"]["NCCL_IB_HCA"] == "bnxt_re0" + + def test_scheduler_fact_reaches_the_slurm_block(self): + """skip_gpus_directive is what keeps sbatch from rejecting the job.""" + config = ConfigLoader.load_slurm_config( + {"slurm": {"nodes": 2, "cluster_profile": ["roce-broadcom-thor2", "no-gpu-gres"]}} + ) + + assert config["slurm"]["skip_gpus_directive"] is True + assert config["env_vars"]["NCCL_IB_DISABLE"] == "0" + + def test_no_profile_leaves_behaviour_unchanged(self): + """Existing deployments see exactly what they saw before.""" + config = ConfigLoader.load_slurm_config({"slurm": {"nodes": 2}}) + + assert config["env_vars"]["NCCL_IB_DISABLE"] == "1" + assert config["env_vars"]["NCCL_SOCKET_IFNAME"] == "eth0" + + def test_unknown_profile_fails_the_run(self): + """Better a startup error than a silent fallback to the wrong fabric.""" + with pytest.raises(ValidationError): + ConfigLoader.load_slurm_config( + {"slurm": {"nodes": 2, "cluster_profile": "not-a-cluster"}} + ) From f573f940bd95332e1c8035c2ed2ceb4301f18bfc Mon Sep 17 00:00:00 2001 From: Mikhail Kuznetsov Date: Thu, 6 Aug 2026 16:54:20 +0000 Subject: [PATCH 5/7] feat(context): support a headless submit node Runtime context initialisation probes the local node for its GPU vendor, architecture, count and render devices. On a SLURM login node there are none of those, so it raised "Unable to determine gpu vendor" and the run stopped before it reached the scheduler -- on a node that was never going to run the workload anyway. A cluster profile facts block now answers those questions when the local node cannot. Facts are a fallback, not an override: a node that can answer for itself still does, so on a heterogeneous partition the architecture recorded in results is the one that ran the work. A probe that reports nothing rather than failing -- "0 GPUs" from a machine with none -- counts as no answer too. --- docs/deployment.md | 23 +++ src/madengine/core/context.py | 121 ++++++++++++-- .../schemas/cluster_profile.schema.json | 9 +- tests/unit/test_headless_context.py | 158 ++++++++++++++++++ 4 files changed, 297 insertions(+), 14 deletions(-) create mode 100644 tests/unit/test_headless_context.py diff --git a/docs/deployment.md b/docs/deployment.md index 858bfa3c..65910cd6 100644 --- a/docs/deployment.md +++ b/docs/deployment.md @@ -324,6 +324,29 @@ names none. A profile that does not exist stops the run: a silent fallback to the wrong fabric costs more than a startup error. +#### Submit nodes without GPUs + +A profile's `facts` block describes the compute nodes, which is what lets a login node with +no GPUs, no ROCm and no `/dev/dri` prepare a job for nodes that have all three: + +```json +{ + "_description": "Our cluster.", + "facts": { + "gpu_vendor": "AMD", + "gpus_per_node": 8, + "gpu_architecture": "gfx942", + "gpu_product_name": "AMD Instinct MI300X", + "hip_version": "6.4" + } +} +``` + +Without this, runtime context initialisation on a login node fails with `Unable to +determine gpu vendor`. Facts are a fallback rather than an override: a node that can answer +for itself still does, so on a heterogeneous partition the architecture recorded in results +is the one that ran the work, not the one the profile expected. + ### Shared image store (`MAD_DOCKER_BUILDS`) For a multi-node run every worker needs the same image. Set `MAD_DOCKER_BUILDS` to a diff --git a/src/madengine/core/context.py b/src/madengine/core/context.py index 7b22443d..83174d68 100644 --- a/src/madengine/core/context.py +++ b/src/madengine/core/context.py @@ -101,6 +101,9 @@ def __init__( self.console = Console() self._gpu_context_initialized = False self._build_only_mode = build_only_mode + # Set in init_gpu_context when a cluster profile describes the compute nodes and + # this one is not expected to look like them. + self._headless = False self._detect_local_gpu_arch = detect_local_gpu_arch self._system_context_initialized = False self._gpu_tool_manager = None # Lazy initialization @@ -254,25 +257,101 @@ def init_system_context(self) -> None: f"System context detection failed on runtime node: {e}" ) + def _cluster_facts(self) -> typing.Dict[str, typing.Any]: + """Node facts from the selected cluster profile, if any. + + A submit node is often a login node with no GPUs, no ROCm and no /dev/dri, yet the + job it prepares runs on nodes that have all three. The cluster profile states what + those nodes are, which is the only honest answer available here: probing the local + node would describe the wrong machine even when it happens to succeed. + + Returns: + Dict[str, Any]: the profile's `facts` block, empty when no profile is selected + """ + from madengine.deployment.presets.cluster_profiles import ( + apply_cluster_profiles, + selected_profiles, + ) + + selection = self.ctx.get("cluster_profile") or (self.ctx.get("slurm") or {}).get( + "cluster_profile" + ) + references = selected_profiles({"slurm": {"cluster_profile": selection}}) + if not references: + return {} + return apply_cluster_profiles({}, references).get("facts") or {} + + def _probe(self, description: str, probe: typing.Callable, fallback=None): + """Run a local probe, tolerating its absence when a cluster profile stands in. + + Args: + description: what is being detected, for the warning + probe: the detection callable + fallback: value to use when detection fails in headless mode + + Returns: + The probed value, or `fallback` in headless mode when probing fails. + + Raises: + Exception: whatever the probe raised, when not running headless + """ + try: + detected = probe() + except Exception as exc: + if not self._headless: + raise + print( + f"Warning: cannot detect {description} on this node ({exc}); " + f"continuing with {fallback!r} from the cluster profile" + ) + return fallback + + # A probe that reports nothing rather than failing — "0 GPUs" from a node with none + # — is not an answer about the compute nodes either. + if self._headless and not detected and fallback: + print( + f"Warning: this node reports no {description}; " + f"continuing with {fallback!r} from the cluster profile" + ) + return fallback + return detected + def init_gpu_context(self) -> None: """Initialize GPU-specific context for runtime. This method detects GPU configuration and sets up environment variables - needed for container execution. Should only be called on GPU nodes. + needed for container execution. On a node with GPUs the values are probed; + on a submit node they come from the selected cluster profile's facts, which + describe the compute nodes the job will actually land on. User-provided GPU contexts will not be overridden. Raises: - RuntimeError: If GPU detection fails. + RuntimeError: If GPU detection fails and no cluster profile supplies the facts. """ if self._gpu_context_initialized: return print("Detecting GPU configuration...") + # A node answers for itself whenever it can: on a compute node the local values + # describe the machine that will run the container, and a cluster-wide profile + # cannot know which architecture a heterogeneous partition handed out. The facts + # are a fallback for the node that has nothing to say. + facts = self._cluster_facts() + self._headless = bool(facts.get("gpu_vendor")) + try: # GPU vendor detection - only if not provided by user if "gpu_vendor" not in self.ctx: - self.ctx["gpu_vendor"] = self.get_gpu_vendor() + vendor = self._probe("the GPU vendor", self.get_gpu_vendor) + if vendor not in ("AMD", "NVIDIA") and facts.get("gpu_vendor"): + print( + f"No GPU on this node; taking the compute nodes' facts from the " + f"cluster profile: " + f"{', '.join(f'{k}={v}' for k, v in sorted(facts.items()))}" + ) + vendor = facts["gpu_vendor"] + self.ctx["gpu_vendor"] = vendor print(f"Detected GPU vendor: {self.ctx['gpu_vendor']}") else: print(f"Using provided GPU vendor: {self.ctx['gpu_vendor']}") @@ -288,24 +367,34 @@ def init_gpu_context(self) -> None: # normalizes it at run time. Auto-detection runs in finalize when absent. if "MAD_SYSTEM_NGPUS" not in self.ctx["docker_env_vars"]: - self.ctx["docker_env_vars"][ - "MAD_SYSTEM_NGPUS" - ] = self.get_system_ngpus() + self.ctx["docker_env_vars"]["MAD_SYSTEM_NGPUS"] = self._probe( + "the GPU count", self.get_system_ngpus, facts.get("gpus_per_node", 0) + ) if "MAD_SYSTEM_GPU_ARCHITECTURE" not in self.ctx["docker_env_vars"]: self.ctx["docker_env_vars"][ "MAD_SYSTEM_GPU_ARCHITECTURE" - ] = self.get_system_gpu_architecture() + ] = self._probe( + "the GPU architecture", + self.get_system_gpu_architecture, + facts.get("gpu_architecture", ""), + ) if "MAD_SYSTEM_HIP_VERSION" not in self.ctx["docker_env_vars"]: - self.ctx["docker_env_vars"][ - "MAD_SYSTEM_HIP_VERSION" - ] = self.get_system_hip_version() + self.ctx["docker_env_vars"]["MAD_SYSTEM_HIP_VERSION"] = self._probe( + "the HIP version", + self.get_system_hip_version, + facts.get("hip_version", ""), + ) if "MAD_SYSTEM_GPU_PRODUCT_NAME" not in self.ctx["docker_env_vars"]: self.ctx["docker_env_vars"][ "MAD_SYSTEM_GPU_PRODUCT_NAME" - ] = self.get_system_gpu_product_name() + ] = self._probe( + "the GPU product name", + self.get_system_gpu_product_name, + facts.get("gpu_product_name", ""), + ) # Also add to build args (for runtime builds) - only if not already set if "MAD_SYSTEM_GPU_ARCHITECTURE" not in self.ctx["docker_build_arg"]: @@ -315,10 +404,16 @@ def init_gpu_context(self) -> None: # Docker GPU configuration - only if not already set if "docker_gpus" not in self.ctx: - self.ctx["docker_gpus"] = self.get_docker_gpus() + self.ctx["docker_gpus"] = self._probe( + "the Docker GPU selection", self.get_docker_gpus + ) if "gpu_renderDs" not in self.ctx: - self.ctx["gpu_renderDs"] = self.get_gpu_renderD_nodes() + # Render nodes are per-machine device numbers; a submit node's would be + # meaningless even if it had any. + self.ctx["gpu_renderDs"] = self._probe( + "the GPU render nodes", self.get_gpu_renderD_nodes + ) self._gpu_context_initialized = True diff --git a/src/madengine/schemas/cluster_profile.schema.json b/src/madengine/schemas/cluster_profile.schema.json index dceac9ad..eb2f647b 100644 --- a/src/madengine/schemas/cluster_profile.schema.json +++ b/src/madengine/schemas/cluster_profile.schema.json @@ -17,7 +17,14 @@ "description": "e.g. gfx942; used in place of probing when the submit node has no GPU.", "type": "string" }, - "gpu_renderer": { "type": "string" }, + "gpu_product_name": { + "description": "e.g. AMD Instinct MI300X; reported in results when probing is not possible.", + "type": "string" + }, + "hip_version": { + "description": "HIP/CUDA version of the compute nodes, e.g. 6.4.", + "type": ["string", "number"] + }, "shared_filesystem": { "description": "Filesystem types df reports for shared storage on this cluster.", "type": "array", diff --git a/tests/unit/test_headless_context.py b/tests/unit/test_headless_context.py new file mode 100644 index 00000000..13dc5856 --- /dev/null +++ b/tests/unit/test_headless_context.py @@ -0,0 +1,158 @@ +#!/usr/bin/env python3 +""" +Tests for building a runtime context on a node without GPUs. + +A SLURM submit node is usually a login node: no GPUs, no ROCm, no /dev/dri. Runtime +context initialisation used to raise there — "Unable to determine gpu vendor" — even though +the job being prepared runs somewhere else entirely. When a cluster profile states what the +compute nodes are, those facts are the honest answer, and probing the login node is not. + +Copyright (c) Advanced Micro Devices, Inc. All rights reserved. +""" + +import json + +import pytest + +from madengine.core.context import Context +from madengine.deployment.presets.cluster_profiles import PROFILE_PATH_ENV + + +@pytest.fixture +def profile_dir(tmp_path, monkeypatch): + """A directory of site profiles on the search path.""" + monkeypatch.setenv(PROFILE_PATH_ENV, str(tmp_path)) + return tmp_path + + +@pytest.fixture +def headless_profile(profile_dir): + """A profile describing compute nodes this machine is not one of.""" + (profile_dir / "compute.json").write_text( + json.dumps( + { + "facts": { + "gpu_vendor": "AMD", + "gpus_per_node": 8, + "gpu_architecture": "gfx942", + "gpu_product_name": "AMD Instinct MI300X", + "hip_version": "6.4", + } + } + ) + ) + return "compute" + + +@pytest.fixture +def no_local_gpu(monkeypatch): + """A machine where every GPU probe fails, like a login node.""" + def unavailable(*args, **kwargs): + raise RuntimeError("no GPU on this node") + + for name in ( + "get_gpu_vendor", + "get_system_ngpus", + "get_system_gpu_architecture", + "get_system_gpu_product_name", + "get_system_hip_version", + "get_docker_gpus", + "get_gpu_renderD_nodes", + ): + monkeypatch.setattr(Context, name, unavailable) + + +def build_context(profile, **kwargs): + """Construct a runtime Context that selects a cluster profile.""" + return Context(additional_context=repr({"cluster_profile": profile, **kwargs})) + + +class TestHeadlessSubmitNode: + """A node with no GPUs can still prepare a job for nodes that have them.""" + + def test_context_builds_without_local_gpus(self, headless_profile, no_local_gpu): + """This is the failure the profile exists to remove.""" + context = build_context(headless_profile) + + assert context.ctx["gpu_vendor"] == "AMD" + + def test_facts_populate_the_container_environment(self, headless_profile, no_local_gpu): + """The values a container would have been given by probing come from the profile.""" + context = build_context(headless_profile) + env = context.ctx["docker_env_vars"] + + assert env["MAD_SYSTEM_NGPUS"] == 8 + assert env["MAD_SYSTEM_GPU_ARCHITECTURE"] == "gfx942" + assert env["MAD_SYSTEM_GPU_PRODUCT_NAME"] == "AMD Instinct MI300X" + assert env["MAD_SYSTEM_HIP_VERSION"] == "6.4" + + def test_architecture_reaches_build_args(self, headless_profile, no_local_gpu): + """A build kicked off from the submit node targets the compute nodes' architecture.""" + context = build_context(headless_profile) + + assert context.ctx["docker_build_arg"]["MAD_SYSTEM_GPU_ARCHITECTURE"] == "gfx942" + + def test_machine_specific_probes_are_not_invented(self, headless_profile, no_local_gpu): + """Render node numbers describe one machine, so the submit node reports none.""" + context = build_context(headless_profile) + + assert context.ctx["gpu_renderDs"] is None + + def test_profile_from_the_slurm_block(self, profile_dir, headless_profile, no_local_gpu): + """A manifest names the profile under slurm, which is where deployments read it.""" + context = Context( + additional_context=repr({"slurm": {"cluster_profile": headless_profile}}) + ) + + assert context.ctx["gpu_vendor"] == "AMD" + + def test_partial_facts_still_probe_what_they_omit(self, profile_dir, no_local_gpu): + """A profile that names only the vendor does not silently invent an architecture.""" + (profile_dir / "vendor-only.json").write_text( + json.dumps({"facts": {"gpu_vendor": "AMD"}}) + ) + + context = build_context("vendor-only") + + assert context.ctx["docker_env_vars"]["MAD_SYSTEM_GPU_ARCHITECTURE"] == "" + + +class TestProbingIsStillTheDefault: + """Nothing changes for a node that can answer for itself.""" + + def test_no_profile_means_detection_failure_is_fatal(self, no_local_gpu): + """Without a profile there is no second source, so the run stops.""" + with pytest.raises(RuntimeError, match="GPU detection failed"): + Context() + + def test_probed_values_win_when_the_node_has_gpus(self, headless_profile, monkeypatch): + """On a compute node the local answer describes the machine that will run the work.""" + monkeypatch.setattr(Context, "get_gpu_vendor", lambda self: "AMD") + monkeypatch.setattr(Context, "get_system_ngpus", lambda self: 4) + monkeypatch.setattr(Context, "get_system_gpu_architecture", lambda self: "gfx950") + monkeypatch.setattr(Context, "get_system_gpu_product_name", lambda self: "local") + monkeypatch.setattr(Context, "get_system_hip_version", lambda self: "7.0") + monkeypatch.setattr(Context, "get_docker_gpus", lambda self: "all") + monkeypatch.setattr(Context, "get_gpu_renderD_nodes", lambda self: [128, 129]) + + context = build_context(headless_profile) + + # The profile speaks for the cluster, this node speaks for itself, and it is the + # one about to run the container — including on a heterogeneous partition where + # the profile's architecture would be wrong. + assert context.ctx["docker_env_vars"]["MAD_SYSTEM_GPU_ARCHITECTURE"] == "gfx950" + assert context.ctx["docker_env_vars"]["MAD_SYSTEM_NGPUS"] == 4 + assert context.ctx["gpu_renderDs"] == [128, 129] + + def test_user_context_still_wins_over_facts(self, headless_profile, no_local_gpu): + """An explicit override beats both the profile and the node.""" + context = build_context(headless_profile, gpu_vendor="NVIDIA") + + assert context.ctx["gpu_vendor"] == "NVIDIA" + + def test_unknown_profile_is_reported(self, profile_dir, no_local_gpu): + """A mistyped profile name must not look like "no profile".""" + from madengine.core.errors import ValidationError + + with pytest.raises((ValidationError, RuntimeError)): + build_context("not-a-profile") From ec55966425f70218141c28af29024e547c2d8ff5 Mon Sep 17 00:00:00 2001 From: Mikhail Kuznetsov Date: Thu, 6 Aug 2026 16:54:49 +0000 Subject: [PATCH 6/7] fix(slurm): report multinode failures honestly A two-node run finished with the scheduler recording exit code 3 and madengine printing "All model executions completed successfully!". Three defects lined up to make that possible: per-node exit codes were collected nowhere, a metric parsed from a log was treated as proof the whole run succeeded, and a deployment the scheduler called failed produced an empty failure list when its results could not be parsed. Each node now records its exit code, host and rank in a marker the submit node reads back, and result collection consults those markers and the scheduler's own verdict before it believes any number. The task script runs under `set -e`, so a non-zero madengine used to end it on that very line, leaving the failing node with no marker and no artifacts -- exactly the evidence a failed node has to leave behind; errexit is now off across that one call and restored once the exit code is captured. What a non-zero node means is decided the way the single-node path already states it in resolve_run_status: a metric is the strongest evidence a run did the work. Primus/Megatron reports throughput from the last global rank only, so every other node finds no metric locally and exits non-zero even when the training was perfect -- both accepted baseline runs on the reference cluster ended with node 0 at exit 3 and node 1 at 0. With a metric, the node outcomes and the scheduler's verdict are warnings and the numbers stand: a row with a number in it is a measurement whatever happened around it, so the rows stay in successful_runs and the verdict travels separately, naming what did not finish. With no metric anywhere the node evidence is all there is and it decides: a node that exited non-zero or never reported gives a failure naming it, a clean set of exit codes gives NO_METRIC and exit 5, because a broken result contract is not the same event as a crash, and neither is success. The summary is built in _summarise_deployment so the rule can be tested without standing up a deployment. For the same reason srun --kill-on-bad-exit is opt-in through slurm.kill_on_bad_exit rather than the default: the rank exiting non-zero may be the peer of the one holding the results, and killing the step there would throw them away. The multiple_results diagnostics now name the paths searched and the variable the container was given, so a model script that does not hold up its end of that contract can be fixed without reading madengine's source. --- docs/cli-reference.md | 15 + docs/deployment.md | 25 ++ src/madengine/cli/commands/run.py | 29 +- src/madengine/cli/constants.py | 3 + src/madengine/deployment/slurm.py | 219 +++++++++++- .../deployment/templates/slurm/job.sh.j2 | 23 +- src/madengine/execution/container_runner.py | 29 +- .../orchestration/run_orchestrator.py | 108 +++++- .../schemas/build_manifest.schema.json | 4 + tests/unit/test_multinode_failure.py | 316 ++++++++++++++++++ 10 files changed, 748 insertions(+), 23 deletions(-) create mode 100644 tests/unit/test_multinode_failure.py diff --git a/docs/cli-reference.md b/docs/cli-reference.md index 842e0fdf..71c95aee 100644 --- a/docs/cli-reference.md +++ b/docs/cli-reference.md @@ -522,9 +522,24 @@ madengine uses standard exit codes so scripts and CI (e.g. Jenkins) can detect s | `2` | `BUILD_FAILURE` | One or more image builds failed (e.g. Docker build error) | | `3` | `RUN_FAILURE` | One or more model executions failed | | `4` | `INVALID_ARGS` | Invalid command-line arguments or configuration | +| `5` | `NO_METRIC` | The workload ran to completion but reported no performance metric | **Failure recording:** Pre-run failures (e.g. image pull, setup) and run failures are recorded in the performance table (`perf.csv`) with status `FAILURE`, so all attempted models appear in the CSV. The file is created automatically if missing. +**`NO_METRIC` versus `RUN_FAILURE`:** a crashed workload and a workload that finished +without producing a number are different problems — the first is a broken run, the second a +broken contract between the model script and madengine — so they get different exit codes +and the row in `perf.csv` reads `NO_METRIC` rather than `FAILURE`. + +**Multi-node verdicts:** every node records its exit code, and the submit node reads them +all back. A metric outweighs an exit code, the same way it does for a single node: where a +framework reports throughput from one rank only, the nodes that collected nothing exit +non-zero on a perfectly healthy run, so a run that produced a metric is reported as a +success with the node outcomes listed as warnings. With no metric anywhere, the node +evidence is all there is and it decides: a node that exited non-zero or never reported at +all gives `RUN_FAILURE` naming the node, while a clean set of exit codes and no metric +gives `NO_METRIC`. + **Example usage in scripts / CI:** ```bash diff --git a/docs/deployment.md b/docs/deployment.md index 65910cd6..9052159f 100644 --- a/docs/deployment.md +++ b/docs/deployment.md @@ -361,6 +361,31 @@ Leaving it unset is supported but leaves image distribution to the operator: wor do not already have the image cannot reconcile it, and the run fails on those nodes with a message saying so. +### How a multi-node run is judged + +Every node writes its exit code to `///node_/node.status` +before it copies anything else, so the outcome survives even when the artifacts do not. The +submit node reads them all back, and weighs them the way a single-node run does — a metric +outweighs an exit code: + +| What came back | Verdict | +| --- | --- | +| A metric, all nodes zero | success | +| A metric, some node non-zero or silent | success, with the node outcomes as warnings | +| No metric, some node non-zero or silent | `RUN_FAILURE`, naming the node | +| No metric, all nodes zero | `NO_METRIC` | + +The warning row is not a technicality. Where a framework reports throughput from one rank +only — Primus/Megatron reports from the last global rank — every other node finds no metric +locally and exits non-zero on a completely healthy run. A verdict that trusted exit codes +over results would fail every such run, so the exit codes are reported and the numbers are +kept. + +For a workload where every rank really is expected to exit zero, set +`slurm.kill_on_bad_exit` to tear the step down on the first bad exit instead of letting the +survivors block on a peer that will never answer. It is off by default for the reason above: +the node exiting non-zero may be the one whose peer holds the numbers. + ### Multi-Node Training For distributed training across SLURM nodes: diff --git a/src/madengine/cli/commands/run.py b/src/madengine/cli/commands/run.py index c961cbe5..9c99743b 100644 --- a/src/madengine/cli/commands/run.py +++ b/src/madengine/cli/commands/run.py @@ -280,16 +280,35 @@ def run( save_summary_with_feedback(execution_summary, summary_output, "Execution") failed_runs = len(execution_summary.get("failed_runs", [])) - if failed_runs == 0: + no_metric_runs = len(execution_summary.get("no_metric_runs", [])) + incomplete = execution_summary.get("incomplete") + for warning in execution_summary.get("warnings", []): + console.print(f"⚠️ [yellow]{warning}[/yellow]") + if incomplete: console.print( - "🎉 [bold green]All model executions completed successfully![/bold green]" + f"💥 [bold red]Run did not complete and produced no metric: " + f"{incomplete.get('reason', 'unknown reason')}[/bold red]" ) - raise typer.Exit(ExitCode.SUCCESS) - else: + raise typer.Exit(ExitCode.RUN_FAILURE) + if failed_runs: console.print( - f"💥 [bold red]Execution failed for {failed_runs} models[/bold red]" + f"💥 [bold red]Execution failed for {failed_runs} " + f"{'run' if failed_runs == 1 else 'runs'}[/bold red]" ) raise typer.Exit(ExitCode.RUN_FAILURE) + if no_metric_runs: + # The workload ran to completion and produced nothing to measure. That is + # not a successful benchmark, and it is not the same failure as a crash: + # what broke is the contract between the model script and madengine. + console.print( + f"📉 [bold yellow]{no_metric_runs} models ran but reported no " + f"performance metric[/bold yellow]" + ) + raise typer.Exit(ExitCode.NO_METRIC) + console.print( + "🎉 [bold green]All model executions completed successfully![/bold green]" + ) + raise typer.Exit(ExitCode.SUCCESS) else: # MAD_CONTAINER_IMAGE handling is done in RunOrchestrator diff --git a/src/madengine/cli/constants.py b/src/madengine/cli/constants.py index b437fa30..bac819bf 100644 --- a/src/madengine/cli/constants.py +++ b/src/madengine/cli/constants.py @@ -17,6 +17,9 @@ class ExitCode(IntEnum): BUILD_FAILURE = 2 RUN_FAILURE = 3 INVALID_ARGS = 4 + #: The workload finished but produced no performance metric. Distinct from + #: RUN_FAILURE so a caller can tell a crashed run from a broken result contract. + NO_METRIC = 5 # Valid values for validation diff --git a/src/madengine/deployment/slurm.py b/src/madengine/deployment/slurm.py index 088a3fb2..af723c47 100644 --- a/src/madengine/deployment/slurm.py +++ b/src/madengine/deployment/slurm.py @@ -80,6 +80,11 @@ def __init__(self, config: DeploymentConfig): self.reservation = self.slurm_config.get("reservation", None) # Some clusters expose no GPU GRES, so sbatch rejects --gpus-per-node. self.skip_gpus_directive = self.slurm_config.get("skip_gpus_directive", False) + # Off by default: in a run whose framework reports throughput from one rank only, + # the node that collected nothing exits non-zero on a healthy run, and tearing the + # step down there would kill the node that holds the numbers. Turn it on for a + # workload where any non-zero rank really does mean the run is over. + self.kill_on_bad_exit = self.slurm_config.get("kill_on_bad_exit", False) # Setup Jinja2 template engine template_dir = Path(__file__).parent / "templates" / "slurm" @@ -694,6 +699,7 @@ def debug(self, msg): "nodes": self.nodes, "gpus_per_node": resolved_gpus_per_node, # Use resolved GPU count "skip_gpus_directive": self.skip_gpus_directive, + "kill_on_bad_exit": self.kill_on_bad_exit, "time_limit": self.time_limit, "output_dir": str(self.output_dir), "master_port": master_port, @@ -1728,6 +1734,157 @@ def _select_best_multiple_results_csv(self, candidates: List[Path]) -> Optional[ return best_candidate + def _job_exit_state(self, job_id: str) -> Optional[str]: + """SLURM's own verdict on the job. + + Args: + job_id: the job to ask about + + Returns: + Optional[str]: the state, e.g. COMPLETED or FAILED, or None when sacct cannot + answer — in which case nothing is claimed on its behalf + """ + try: + result = subprocess.run( + ["sacct", "-j", job_id, "-n", "-X", "-o", "State"], + capture_output=True, + text=True, + timeout=10, + ) + except Exception: + return None + if result.returncode != 0: + return None + first_line = (result.stdout or "").strip().splitlines() + return first_line[0].strip().upper() if first_line else None + + def _read_node_statuses(self, job_dir: Path) -> List[Dict[str, Any]]: + """Read the per-node status markers the job script leaves behind. + + Args: + job_dir: collection directory for this job + + Returns: + List[Dict[str, Any]]: one entry per node found, sorted by node rank, each with + `node`, `host` and `exit_code` + """ + statuses: List[Dict[str, Any]] = [] + for status_path in sorted(job_dir.glob("node_*/node.status")): + fields: Dict[str, Any] = {} + try: + for line in status_path.read_text(encoding="utf-8", errors="ignore").splitlines(): + key, _, value = line.partition("=") + if key: + fields[key.strip()] = value.strip() + except OSError: + continue + try: + fields["node"] = int(fields.get("node", status_path.parent.name.replace("node_", ""))) + fields["exit_code"] = int(fields.get("exit_code", 1)) + except ValueError: + continue + statuses.append(fields) + return sorted(statuses, key=lambda s: s["node"]) + + def _report_node_statuses(self, statuses: List[Dict[str, Any]]) -> List[Dict[str, Any]]: + """Print the per-node outcome and return the nodes that failed. + + Args: + statuses: entries from `_read_node_statuses` + + Returns: + List[Dict[str, Any]]: the entries whose exit code is non-zero + """ + failures = [s for s in statuses if s["exit_code"] != 0] + missing = [n for n in range(self.nodes) if n not in {s["node"] for s in statuses}] + + if failures: + self.console.print("[red]Nodes that failed:[/red]") + for status in failures: + self.console.print( + f"[red] node {status['node']} ({status.get('host', 'unknown host')}): " + f"exit code {status['exit_code']}[/red]" + ) + if missing and statuses: + # A node that never wrote a marker did not reach the end of the task script: + # it was killed, or the node itself went away. Only meaningful once at least + # one node has reported, since a job may predate these markers entirely. + self.console.print( + f"[red]Nodes that reported no status (killed or lost): " + f"{', '.join(str(n) for n in missing)}[/red]" + ) + failures = failures + [ + {"node": n, "host": "", "exit_code": None, "missing": True} for n in missing + ] + return failures + + @staticmethod + def _incomplete_reason( + node_failures: List[Dict[str, Any]], job_state: Optional[str] + ) -> str: + """Say why the run is incomplete, in the words of whoever noticed. + + Args: + node_failures: entries from `_report_node_statuses` + job_state: the state sacct reports, if it could be read + + Returns: + str: a phrase naming the failing nodes, or SLURM's own verdict + """ + if node_failures: + parts = [] + for failure in node_failures: + if failure.get("missing"): + parts.append(f"node {failure['node']} reported no status (killed or lost)") + else: + parts.append(f"node {failure['node']} exited {failure['exit_code']}") + return ", ".join(parts) + if job_state: + return f"SLURM reports the job as {job_state}" + return "the job did not complete" + + def _record_verdict( + self, + results: Dict[str, Any], + model: str, + node_failures: List[Dict[str, Any]], + job_state: Optional[str], + rows_collected: int, + ) -> None: + """Record what a run that did not end cleanly amounts to. + + A metric outweighs an exit code, the way it already does for a single node + (`resolve_run_status`): a node exiting non-zero is routine in a multi-node run + whose framework reports throughput from one rank only, so the node that collected + nothing fails locally while the run as a whole did the work. What such a node + earns is a warning, not a verdict. With no metric anywhere, the node evidence is + all there is, and it decides. + + Args: + results: collection results, updated in place + model: model the run belongs to + node_failures: entries from `_report_node_statuses` + job_state: the state sacct reports, if it could be read + rows_collected: how many metric rows the run produced + """ + reason = self._incomplete_reason(node_failures, job_state) + if rows_collected: + warning = ( + f"{reason}, while the run produced {rows_collected} metric row(s); " + f"reported as measurements, check the node logs before trusting them" + ) + results.setdefault("warnings", []).append(warning) + self.console.print(f"[yellow]⚠ {warning}[/yellow]") + return + results["incomplete"] = { + "model": model, + "reason": reason, + "rows_collected": 0, + } + self.console.print( + f"[red]✗ Job did not complete and produced no metric: {reason}[/red]" + ) + def collect_results(self, deployment_id: str) -> Dict[str, Any]: """Collect performance results from SLURM output files. @@ -1751,6 +1908,7 @@ def collect_results(self, deployment_id: str) -> Dict[str, Any]: "logs": [], "successful_runs": [], "failed_runs": [], + "no_metric_runs": [], "session_start_row": session_start_row, } @@ -1786,6 +1944,22 @@ def collect_results(self, deployment_id: str) -> Dict[str, Any]: job_dir = self.output_dir / model_name_for_path / deployment_id job_dir.mkdir(parents=True, exist_ok=True) + # Per-node outcomes come first: a metric parsed out of a log says nothing about + # whether the other nodes finished, and a run where one rank died is not a result. + node_statuses = self._read_node_statuses(job_dir) + node_failures = self._report_node_statuses(node_statuses) + job_state = self._job_exit_state(deployment_id) + # SLURM's verdict counts even when no node marker says so — a step killed by the + # scheduler leaves no marker at all. + job_failed = bool(node_failures) or ( + job_state is not None and "COMPLETED" not in job_state + ) + if job_failed and not node_failures and job_state: + self.console.print(f"[red]SLURM reports job {deployment_id} as {job_state}[/red]") + results["node_statuses"] = node_statuses + results["node_failures"] = node_failures + results["job_state"] = job_state + # Gather log content per node: from job_dir/node_N/ (new) or flat output_dir .out files per_node_log_contents: List[tuple] = [] flat_out_files = sorted(self.output_dir.glob(f"madengine-*_{deployment_id}_*.out")) @@ -1921,12 +2095,18 @@ def collect_results(self, deployment_id: str) -> Dict[str, Any]: ) results["perf_files"] = [str(Path("perf.csv").resolve())] import csv as _csv + # A row with a number in it is a measurement, and it stays one even when the + # run around it did not finish: a model that reports four precisions and + # crashes on two still measured the other two. The verdict on the run as a + # whole travels separately, in results["incomplete"]. + rows_collected = 0 try: with open(resolved_csv, "r", encoding="utf-8", errors="ignore") as f: reader = _csv.DictReader(f) for row in reader: row = {k.strip(): v for k, v in row.items() if k} if row.get("performance") and row.get("metric"): + rows_collected += 1 results["successful_runs"].append({ "model": model_info_for_entry.get("name", "") + "_" + row.get("model", ""), "status": "SUCCESS", @@ -1939,9 +2119,18 @@ def collect_results(self, deployment_id: str) -> Dict[str, Any]: }) except Exception: pass - self.console.print( - f"[green]✓ Updated perf.csv, perf_super.* from multiple_results (Docker-compatible)[/green]" - ) + if job_failed: + self._record_verdict( + results, + model_info_for_entry.get("name", model_name), + node_failures, + job_state, + rows_collected, + ) + else: + self.console.print( + f"[green]✓ Updated perf.csv, perf_super.* from multiple_results (Docker-compatible)[/green]" + ) return results # multiple_results set but CSV not found: fall through to single-result path (may write FAILURE) @@ -2021,6 +2210,28 @@ def collect_results(self, deployment_id: str) -> Dict[str, Any]: ) if run_details_dict is not None: + # Three outcomes, not two. A workload that failed and a workload that ran to + # completion without producing a metric need different answers: the first is a + # broken run, the second is a broken contract between the model script and + # madengine, and reporting both as FAILURE hides which one happened. + has_metric = bool(str(run_details_dict.get("performance") or "").strip()) + if job_failed: + self._record_verdict( + results, + run_details_dict.get("model", model_name), + node_failures, + job_state, + 1 if has_metric else 0, + ) + if not has_metric: + run_details_dict["status"] = "FAILURE" + elif not has_metric: + run_details_dict["status"] = "NO_METRIC" + self.console.print( + "[yellow]⚠ The workload finished on every node but no performance " + "metric was collected; recording the run as NO_METRIC[/yellow]" + ) + perf_entry_path = Path("perf_entry.json") with open(perf_entry_path, "w") as f: json.dump(run_details_dict, f, indent=2) @@ -2065,6 +2276,8 @@ def collect_results(self, deployment_id: str) -> Dict[str, Any]: } if run_details_dict.get("status") == "SUCCESS": results["successful_runs"].append(run_data) + elif run_details_dict.get("status") == "NO_METRIC": + results.setdefault("no_metric_runs", []).append(run_data) else: results["failed_runs"].append(run_data) summary = { diff --git a/src/madengine/deployment/templates/slurm/job.sh.j2 b/src/madengine/deployment/templates/slurm/job.sh.j2 index 4cdddba0..54becfd4 100644 --- a/src/madengine/deployment/templates/slurm/job.sh.j2 +++ b/src/madengine/deployment/templates/slurm/job.sh.j2 @@ -675,6 +675,11 @@ echo "[DEBUG] $(date -Iseconds) Node ${SLURM_PROCID} ($(hostname)): about to run # Run madengine with output redirected to node-specific log files # Use append (2>>) for stderr so debug lines written above are not overwritten # Environment variables (MASTER_ADDR, MAD_MULTI_NODE_RUNNER, etc.) are inherited +# +# errexit is off across this call on purpose: under `set -e` a failing madengine ends the +# task script right here, so the node writes no status marker and copies no artifacts — +# the run loses exactly the evidence a failed node has to leave behind. +set +e $MAD_CLI_COMMAND run \ --manifest-file "$EXEC_MANIFEST" \ --timeout {{ timeout | default(3600) }} \ @@ -683,6 +688,7 @@ $MAD_CLI_COMMAND run \ > "${NODE_LOG_OUT}" 2>> "${NODE_LOG_ERR}" TASK_EXIT=$? +set -e echo "[DEBUG] $(date -Iseconds) Node ${SLURM_PROCID} ($(hostname)): madengine exited with code $TASK_EXIT" >> "${NODE_LOG_ERR}" echo "" echo "Task completed with exit code: $TASK_EXIT" @@ -700,6 +706,16 @@ JOB_COLLECTION_DIR="${SUBMISSION_DIR}/{{ output_dir }}/{{ model_name }}/${SLURM_ NODE_COLLECTION_DIR="${JOB_COLLECTION_DIR}/node_${SLURM_PROCID}" mkdir -p "$NODE_COLLECTION_DIR" +# Record how this node ended, whatever happened. The submit node reads these to tell a +# workload failure from a missing metric; without them a dead worker is only visible as a +# number in someone's terminal. +{ + echo "node=${SLURM_PROCID}" + echo "host=$(hostname)" + echo "exit_code=${TASK_EXIT}" + echo "timestamp=$(date -Iseconds)" +} > "$NODE_COLLECTION_DIR/node.status" + if [ $TASK_EXIT -eq 0 ]; then echo "" echo "========================================================================" @@ -736,7 +752,12 @@ TASK_SCRIPT_EOF chmod +x "$TASK_SCRIPT" echo "Launching tasks on {{ nodes }} nodes..." -srun bash "$TASK_SCRIPT" +# slurm.kill_on_bad_exit tears the step down on the first non-zero rank instead of leaving +# the survivors to block on a peer that will never answer. It is off by default because a +# rank that exits non-zero is not always a dead run: where the framework reports throughput +# from one rank only, the other nodes fail locally on a perfectly healthy run, and killing +# the step there would take the node holding the numbers with it. +srun {% if kill_on_bad_exit %}--kill-on-bad-exit=1 {% endif %}bash "$TASK_SCRIPT" EXIT_CODE=$? # Cleanup task script diff --git a/src/madengine/execution/container_runner.py b/src/madengine/execution/container_runner.py index eab4af7d..0c2d597b 100644 --- a/src/madengine/execution/container_runner.py +++ b/src/madengine/execution/container_runner.py @@ -1706,9 +1706,24 @@ def run_container( multiple_results, model_dir ) if not resolved_path: + # This is a contract, and the model script held up + # its end or it did not: say which paths were + # searched and what the container was told, so the + # next person does not have to read this code. self.rich_console.print( - f"[yellow]Warning: Could not find multiple results file " - f"(tried cwd and {model_dir}/): {multiple_results}[/yellow]" + f"[yellow]Warning: model '{model_info.get('name', '')}' " + f"declares multiple_results='{multiple_results}' but no " + f"such file was produced.[/yellow]" + ) + self.rich_console.print( + f"[yellow] Searched: {os.path.abspath(multiple_results)}" + f" and {os.path.join(os.path.abspath(model_dir), multiple_results)}[/yellow]" + ) + self.rich_console.print( + f"[yellow] The container was given " + f"MAD_OUTPUT_CSV='{multiple_results}'; the run script must " + f"write a CSV there with 'performance' and 'metric' " + f"columns, relative to its working directory.[/yellow]" ) run_results["performance"] = None else: @@ -1724,7 +1739,10 @@ def run_container( # Check if 'performance' column exists if 'performance' not in csv_reader.fieldnames: - print("Error: 'performance' column not found in multiple results file.") + print( + f"Error: {resolved_path} has no 'performance' column; " + f"found: {', '.join(csv_reader.fieldnames or []) or '(no header)'}" + ) run_results["performance"] = None else: # Check if at least one row has a non-empty performance value @@ -1736,7 +1754,10 @@ def run_container( if not has_valid_perf: run_results["performance"] = None - print("Error: Performance metric is empty in all rows of multiple results file.") + print( + f"Error: every row of {resolved_path} has an empty " + f"'performance' value, so the run produced no metric." + ) except Exception as e: self.rich_console.print( f"[yellow]Warning: Could not validate multiple results file: {e}[/yellow]" diff --git a/src/madengine/orchestration/run_orchestrator.py b/src/madengine/orchestration/run_orchestrator.py index ad894c69..2b5f3724 100644 --- a/src/madengine/orchestration/run_orchestrator.py +++ b/src/madengine/orchestration/run_orchestrator.py @@ -16,7 +16,7 @@ import shlex import subprocess from pathlib import Path -from typing import Dict, Optional +from typing import Any, Dict, Optional from rich.console import Console as RichConsole from rich.panel import Panel @@ -767,15 +767,103 @@ def _execute_distributed(self, target: str, manifest_file: str) -> Dict: self.rich_console.print(f"[dim]{'=' * 60}[/dim]\n") - # Return metrics in the format expected by display_results_table - # Extract successful_runs and failed_runs from metrics if available - if result.metrics: - return { - "successful_runs": result.metrics.get("successful_runs", []), - "failed_runs": result.metrics.get("failed_runs", []), - } - else: - return {"successful_runs": [], "failed_runs": []} + return self._summarise_deployment(result, target, manifest_file) + + def _summarise_deployment( + self, result, target: str, manifest_file: str + ) -> Dict[str, Any]: + """Turn a deployment result into the summary the CLI reports on. + + Args: + result: what the deployment layer returned + target: deployment target name, for the records + manifest_file: manifest the run was launched from + + Returns: + Dict[str, Any]: runs by outcome, plus any warnings and the deployment status + """ + metrics = result.metrics or {} + summary = { + "successful_runs": metrics.get("successful_runs", []), + "failed_runs": list(metrics.get("failed_runs", [])), + "no_metric_runs": metrics.get("no_metric_runs", []), + "deployment_status": result.status.value, + } + + # A run that measured something despite a node ending badly is still a run that + # measured something; what the node did is a warning the caller has to see, not a + # verdict that throws the numbers away. + warnings = metrics.get("warnings") + if warnings: + summary["warnings"] = list(warnings) + + # No metric anywhere, on the other hand, leaves the node evidence as the only + # account of what happened, and it decides. + incomplete = metrics.get("incomplete") + if incomplete: + summary["incomplete"] = incomplete + summary["failed_runs"].append( + { + "model": incomplete.get("model", ""), + "status": "FAILURE", + "performance": "", + "metric": "", + "duration": "", + "gpu_arch": "", + "deployment": target, + "machine": result.deployment_id, + "error": incomplete.get("reason", ""), + } + ) + + # A job the scheduler calls failed is a failed run even when nothing could be + # parsed to say so. Without this the caller sees an empty failure list and reports + # success for a job that died. Measurements change that: the scheduler's verdict + # follows from a rank exiting non-zero, which is routine in a multi-node run whose + # framework reports from one rank only, so with results in hand it is a warning. + if not result.is_success and summary["successful_runs"] and not summary["failed_runs"]: + # The deployment layer may already have said which node ended how; only speak + # up when nothing else accounted for the scheduler's verdict. + if not summary.get("warnings"): + summary["warnings"] = [ + f"{target} reported the run as {result.status.value} " + f"({result.message}), while " + f"{len(summary['successful_runs'])} measurement(s) came back" + ] + elif not result.is_success and not summary["failed_runs"]: + summary["failed_runs"].append( + { + "model": self._model_name_from_manifest(manifest_file), + "status": "FAILURE", + "performance": "", + "metric": "", + "duration": "", + "gpu_arch": "", + "deployment": target, + "machine": result.deployment_id, + "error": result.message, + } + ) + + return summary + + @staticmethod + def _model_name_from_manifest(manifest_file: str) -> str: + """Name the run in a failure record, for a job that produced no results of its own. + + Args: + manifest_file: manifest the run was launched from + + Returns: + str: the first model's name, or "unknown" + """ + try: + with open(manifest_file) as f: + models = json.load(f).get("built_models") or {} + except (OSError, ValueError): + return "unknown" + first = next(iter(models.values()), {}) + return first.get("name") or "unknown" def _show_node_info(self): """Show node ROCm information.""" diff --git a/src/madengine/schemas/build_manifest.schema.json b/src/madengine/schemas/build_manifest.schema.json index 8c16317d..5eab124c 100644 --- a/src/madengine/schemas/build_manifest.schema.json +++ b/src/madengine/schemas/build_manifest.schema.json @@ -123,6 +123,10 @@ "description": "Omit #SBATCH --gpus-per-node, for clusters that advertise no GPU GRES.", "type": "boolean" }, + "kill_on_bad_exit": { + "description": "Tear the srun step down as soon as one rank exits non-zero (default false). Only for workloads where every rank is expected to exit zero: where throughput is reported from one rank only, the other ranks exit non-zero on a healthy run.", + "type": "boolean" + }, "time": { "type": "string" }, "output_dir": { "type": "string" }, "results_dir": { "type": "string" }, diff --git a/tests/unit/test_multinode_failure.py b/tests/unit/test_multinode_failure.py new file mode 100644 index 00000000..aeb58ebe --- /dev/null +++ b/tests/unit/test_multinode_failure.py @@ -0,0 +1,316 @@ +#!/usr/bin/env python3 +""" +Tests for reporting the truth about a multi-node run. + +A two-node run that ended with the scheduler recording exit code 3 used to finish with +"All model executions completed successfully!", because the only thing anyone looked at was +whether a number could be parsed out of a log. Three separate defects made that possible: +per-node exit codes were never collected, nothing distinguished a crashed workload from a +missing metric, and a failed deployment whose results could not be parsed produced an empty +failure list. These tests pin all three. + +Copyright (c) Advanced Micro Devices, Inc. All rights reserved. +""" + +from pathlib import Path +from unittest.mock import MagicMock, patch + +import pytest + +from madengine.deployment.slurm import SlurmDeployment + + +@pytest.fixture +def deployment(): + """A SLURM deployment with the scheduler and manifest stubbed out.""" + instance = object.__new__(SlurmDeployment) + instance.nodes = 2 + instance.gpus_per_node = 8 + instance.console = MagicMock() + return instance + + +def write_status(job_dir, node, exit_code, host="node-a"): + """Write the marker a task script leaves behind.""" + node_dir = job_dir / f"node_{node}" + node_dir.mkdir(parents=True, exist_ok=True) + (node_dir / "node.status").write_text( + f"node={node}\nhost={host}\nexit_code={exit_code}\ntimestamp=2026-08-06T12:00:00+00:00\n" + ) + return node_dir + + +class TestReadingNodeStatuses: + """Per-node exit codes, which nothing collected before.""" + + def test_statuses_are_read_in_node_order(self, deployment, tmp_path): + """Node 10 sorts after node 2, unlike its directory name.""" + deployment.nodes = 11 + for node in (0, 2, 10): + write_status(tmp_path, node, 0) + + statuses = deployment._read_node_statuses(tmp_path) + + assert [s["node"] for s in statuses] == [0, 2, 10] + + def test_exit_code_and_host_are_kept(self, deployment, tmp_path): + """Which node failed matters as much as that one did.""" + write_status(tmp_path, 1, 137, host="node-b") + + status = deployment._read_node_statuses(tmp_path)[0] + + assert status["exit_code"] == 137 + assert status["host"] == "node-b" + + def test_missing_markers_are_not_invented(self, deployment, tmp_path): + """A job from before these markers existed reports nothing, not failure.""" + assert deployment._read_node_statuses(tmp_path) == [] + + def test_unparseable_marker_is_skipped(self, deployment, tmp_path): + """A truncated marker must not take down result collection.""" + node_dir = tmp_path / "node_0" + node_dir.mkdir() + (node_dir / "node.status").write_text("exit_code=not-a-number\n") + + assert deployment._read_node_statuses(tmp_path) == [] + + +class TestReportingNodeStatuses: + """Turning markers into a verdict.""" + + def test_all_zero_is_no_failure(self, deployment, tmp_path): + """The happy path stays quiet.""" + write_status(tmp_path, 0, 0) + write_status(tmp_path, 1, 0) + + assert deployment._report_node_statuses(deployment._read_node_statuses(tmp_path)) == [] + + def test_non_zero_exit_is_a_failure(self, deployment, tmp_path): + """The case that used to be reported as success.""" + write_status(tmp_path, 0, 0) + write_status(tmp_path, 1, 3, host="node-b") + + failures = deployment._report_node_statuses(deployment._read_node_statuses(tmp_path)) + + assert [f["node"] for f in failures] == [1] + assert failures[0]["exit_code"] == 3 + + def test_a_node_that_never_reported_is_a_failure(self, deployment, tmp_path): + """A node killed mid-run writes no marker; silence is not success.""" + write_status(tmp_path, 0, 0) + + failures = deployment._report_node_statuses(deployment._read_node_statuses(tmp_path)) + + assert [f["node"] for f in failures] == [1] + assert failures[0]["missing"] is True + + def test_no_markers_at_all_means_no_verdict(self, deployment, tmp_path): + """Without a single marker there is nothing to conclude from their absence.""" + assert deployment._report_node_statuses([]) == [] + + +class TestIncompleteReason: + """A verdict has to say what happened, not just that something did.""" + + def test_a_failed_node_is_named_with_its_exit_code(self, deployment): + """"Execution failed" alone sends the reader back into the logs.""" + reason = deployment._incomplete_reason( + [{"node": 0, "host": "node-a", "exit_code": 3}], "FAILED" + ) + + assert reason == "node 0 exited 3" + + def test_a_lost_node_is_described_as_lost(self, deployment): + """A node with no marker did not exit; it disappeared.""" + reason = deployment._incomplete_reason( + [{"node": 1, "host": "", "exit_code": None, "missing": True}], None + ) + + assert reason == "node 1 reported no status (killed or lost)" + + def test_every_failing_node_is_listed(self, deployment): + """Two nodes down is a different story from one.""" + reason = deployment._incomplete_reason( + [ + {"node": 0, "host": "node-a", "exit_code": 3}, + {"node": 1, "host": "node-b", "exit_code": 137}, + ], + "FAILED", + ) + + assert reason == "node 0 exited 3, node 1 exited 137" + + def test_without_markers_slurm_speaks(self, deployment): + """The step the scheduler killed leaves nothing behind but its state.""" + assert deployment._incomplete_reason([], "TIMEOUT") == ( + "SLURM reports the job as TIMEOUT" + ) + + def test_with_no_evidence_at_all_the_verdict_stays_plain(self, deployment): + """Nothing to quote, so claim nothing specific.""" + assert deployment._incomplete_reason([], None) == "the job did not complete" + + +class TestVerdictWhenANodeEndsBadly: + """A metric outweighs an exit code, as it already does for a single node.""" + + def test_a_measured_run_keeps_its_numbers_and_gains_a_warning(self, deployment): + """The rank that reports no throughput exits non-zero on a healthy run.""" + results: dict = {} + + deployment._record_verdict( + results, "llama", [{"node": 0, "host": "node-a", "exit_code": 3}], "FAILED", 8 + ) + + assert "incomplete" not in results + assert len(results["warnings"]) == 1 + assert "node 0 exited 3" in results["warnings"][0] + assert "8 metric row" in results["warnings"][0] + + def test_a_run_with_nothing_measured_is_a_failure(self, deployment): + """With no metric anywhere the node evidence is all there is.""" + results: dict = {} + + deployment._record_verdict( + results, "llama", [{"node": 1, "host": "node-b", "exit_code": 137}], "FAILED", 0 + ) + + assert "warnings" not in results + assert results["incomplete"]["reason"] == "node 1 exited 137" + assert results["incomplete"]["model"] == "llama" + + def test_a_lost_node_is_named_in_the_warning(self, deployment): + """Whatever the verdict, the operator hears which node went missing.""" + results: dict = {} + + deployment._record_verdict( + results, "llama", [{"node": 1, "exit_code": None, "missing": True}], None, 4 + ) + + assert "reported no status" in results["warnings"][0] + + +class TestSummaryTheCliReportsOn: + """What the orchestrator hands the CLI, for a job the scheduler calls failed.""" + + @staticmethod + def _summarise(metrics, is_success=False): + from madengine.orchestration.run_orchestrator import RunOrchestrator + + orchestrator = object.__new__(RunOrchestrator) + result = MagicMock() + result.metrics = metrics + result.is_success = is_success + result.status.value = "failed" + result.message = "Job 24505 failed: FAILED" + result.deployment_id = "24505" + return orchestrator._summarise_deployment(result, "slurm", "manifest.json") + + def test_measurements_survive_a_failed_job_state(self): + """The scheduler's verdict follows from a rank that exits non-zero by design.""" + summary = self._summarise({"successful_runs": [{"model": "llama"}]}) + + assert summary["failed_runs"] == [] + assert "measurement(s) came back" in summary["warnings"][0] + + def test_the_deployment_does_not_repeat_a_warning_already_made(self): + """The node-level warning is the specific one; two say no more than one.""" + summary = self._summarise( + {"successful_runs": [{"model": "llama"}], "warnings": ["node 0 exited 3"]} + ) + + assert summary["warnings"] == ["node 0 exited 3"] + + def test_a_failed_job_with_nothing_measured_is_still_a_failure(self): + """Otherwise the caller sees an empty failure list and reports success.""" + summary = self._summarise({"successful_runs": []}) + + assert len(summary["failed_runs"]) == 1 + assert summary["failed_runs"][0]["error"] == "Job 24505 failed: FAILED" + + def test_an_incomplete_run_names_the_model_that_did_not_finish(self): + """The deployment layer's own verdict, carried through.""" + summary = self._summarise( + {"successful_runs": [], "incomplete": {"model": "llama", "reason": "node 1 exited 9"}} + ) + + assert summary["incomplete"]["reason"] == "node 1 exited 9" + assert [f["model"] for f in summary["failed_runs"]] == ["llama"] + + +class TestJobExitState: + """SLURM's own verdict, for the failures that leave no marker.""" + + def test_state_is_read_from_sacct(self, deployment): + """The state, upper-cased, as sacct reports it.""" + with patch("madengine.deployment.slurm.subprocess.run") as run: + run.return_value = MagicMock(returncode=0, stdout="COMPLETED\n") + + assert deployment._job_exit_state("123") == "COMPLETED" + + def test_failed_state_is_reported(self, deployment): + """A job the scheduler killed.""" + with patch("madengine.deployment.slurm.subprocess.run") as run: + run.return_value = MagicMock(returncode=0, stdout="CANCELLED by 1001\n") + + assert deployment._job_exit_state("123") == "CANCELLED BY 1001" + + def test_unavailable_sacct_claims_nothing(self, deployment): + """No accounting database is not evidence either way.""" + with patch("madengine.deployment.slurm.subprocess.run") as run: + run.return_value = MagicMock(returncode=1, stdout="") + + assert deployment._job_exit_state("123") is None + + def test_exception_claims_nothing(self, deployment): + """Neither is a missing sacct binary.""" + with patch("madengine.deployment.slurm.subprocess.run", side_effect=OSError): + assert deployment._job_exit_state("123") is None + + +class TestJobScriptRecordsOutcomes: + """What the rendered job script does on the nodes.""" + + @pytest.fixture + def rendered(self, tmp_path): + """The two-node job script, rendered exactly as prepare() renders it.""" + from tests.unit.test_slurm_job_template import ( + MODEL_ENTRY, + _build_deployment, + ) + + deployment = _build_deployment(tmp_path) + context = deployment._prepare_template_context(MODEL_ENTRY) + return deployment.jinja_env.get_template("job.sh.j2").render(**context) + + def test_every_node_records_its_exit_code(self, rendered): + """The marker the submit node reads back.""" + assert "exit_code=${TASK_EXIT}" in rendered + assert "node.status" in rendered + + def test_the_step_is_not_torn_down_on_a_bad_exit_by_default(self, rendered): + """The rank that collects nothing exits non-zero on a healthy multi-node run. + + Tearing the step down there would kill the rank that holds the numbers, so the + teardown is something a workload opts into. + """ + assert "--kill-on-bad-exit" not in rendered.split("srun ")[-1] + + def test_a_failing_run_still_reaches_the_collection_block(self, rendered): + """Under errexit the script died on the madengine line, losing marker and artifacts.""" + madengine_call = rendered.index("$MAD_CLI_COMMAND run") + assert rendered.rindex("set +e", 0, madengine_call) < madengine_call + assert "TASK_EXIT=$?\nset -e" in rendered + + def test_teardown_can_be_turned_on(self, tmp_path): + """A workload where every rank must exit zero can stop holding the allocation.""" + from tests.unit.test_slurm_job_template import ( + MODEL_ENTRY, + _build_deployment, + ) + + deployment = _build_deployment(tmp_path, {"kill_on_bad_exit": True}) + context = deployment._prepare_template_context(MODEL_ENTRY) + rendered = deployment.jinja_env.get_template("job.sh.j2").render(**context) + + assert "srun --kill-on-bad-exit=1" in rendered From 6449337971d113db945232c17f8b979ce25736e9 Mon Sep 17 00:00:00 2001 From: Mikhail Kuznetsov Date: Fri, 7 Aug 2026 10:31:40 +0000 Subject: [PATCH 7/7] feat(reporting): find a model's results CSV by its shape A model card declares where its results CSV lives in multiple_results, but the field is optional and its value has to agree with a path a script in another repository builds by hand. When the two disagreed madengine fell back to scraping 'performance: ' from the log, found nothing there because a CSV-reporting model never prints that line, and wrote a row with an empty performance and status=FAILURE. One yellow warning was the only trace, and a typo looked exactly like an omission. The shape of the file is the sturdier contract. madengine.reporting.result_csv recognises a results CSV by a header carrying model, performance and metric, in any order and with any number of extra columns -- which is what MAD's scripts write, from three columns in scripts/dummy/run_multi.sh to twenty-nine in scripts/large-ep-benchmark/parse_ep_to_csv.py, with one script spacing them out as 'model, performance, metric'. When nothing was declared, or what was declared is not there, the run directory and the workspace root are searched at depth 1 and the best match stands in. A declared file that exists always wins. Two things keep the search honest. madengine's own perf.csv and the perf_super/perf_entry family are excluded: they match by construction, and reading one back would feed a run its own previous verdict. And the workspace root is shared by every model in a run, so a file not written during this model's run is ignored -- fourteen of the fifteen MAD scripts that write a conforming CSV put it in exactly that shared directory. The job templates blocked all of this: the per-node copy sat inside {% if multiple_results %}, so a card without the field left nothing on the node for the collector to find. All four copy blocks now also sweep for shape-matching CSVs, with the same exclusions, from one macro per template. What a run says when it finds nothing changes too. Instead of 'Performance metric not found in expected format' it names both halves of the contract: where it looked, how many CSVs it saw, why each was rejected, and whether MAD_OUTPUT_CSV was exported at all. The ranking that chooses between several candidates moves out of the SLURM collector into the shared module, so the Docker and SLURM paths judge the same file the same way and ties break deterministically. Known limit: 33 multi-node cards write their CSV to /run_logs on shared storage, outside any depth-1 search. They were undiscoverable before this change too. Tests: 675 -> 732 unit tests, all passing. --- src/madengine/deployment/slurm.py | 211 +++++------ .../templates/kubernetes/job.yaml.j2 | 16 + .../deployment/templates/slurm/job.sh.j2 | 21 +- src/madengine/execution/container_runner.py | 199 ++++++---- src/madengine/reporting/result_csv.py | 250 +++++++++++++ tests/unit/test_result_csv_discovery.py | 352 ++++++++++++++++++ tests/unit/test_result_csv_templates.py | 115 ++++++ 7 files changed, 979 insertions(+), 185 deletions(-) create mode 100644 src/madengine/reporting/result_csv.py create mode 100644 tests/unit/test_result_csv_discovery.py create mode 100644 tests/unit/test_result_csv_templates.py diff --git a/src/madengine/deployment/slurm.py b/src/madengine/deployment/slurm.py index af723c47..9c42e90a 100644 --- a/src/madengine/deployment/slurm.py +++ b/src/madengine/deployment/slurm.py @@ -32,6 +32,7 @@ from madengine.utils.gpu_config import resolve_runtime_gpus from madengine.utils.run_details import get_build_number, get_pipeline from madengine.utils.path_utils import scripts_base_dir_from +from madengine.reporting import result_csv import json @@ -1689,47 +1690,17 @@ def _select_best_multiple_results_csv(self, candidates: List[Path]) -> Optional[ values (e.g. master perf empty while a worker has the real numbers). Ranking candidates by the count of non-empty performance rows lets downstream aggregation use the richest data instead of depending on - node-0 winning the race or being non-empty. Header/row keys are stripped - so a leading space in the CSV header (some Primus configs) still matches. - Ties break on total row count; candidates[0] is the ultimate fallback. + node-0 winning the race or being non-empty. The ranking itself lives in + madengine.reporting.result_csv so the Docker path judges the same file + the same way. """ - if not candidates: - return None - if len(candidates) == 1: - return candidates[0] - import csv as _csv - best_candidate: Optional[Path] = None - best_score = -1 - best_rows = -1 - for candidate in candidates: - non_empty_perf = 0 - total_rows = 0 - has_perf_column = False - try: - with open(candidate, "r", encoding="utf-8", errors="ignore") as f: - reader = _csv.DictReader(f) - fieldnames = reader.fieldnames or [] - stripped_fields = [fn.strip() for fn in fieldnames] - has_perf_column = "performance" in stripped_fields - for row in reader: - total_rows += 1 - if has_perf_column: - normalized_row = {(k.strip() if isinstance(k, str) else k): v for k, v in row.items()} - value = (normalized_row.get("performance") or "").strip() - if value: - non_empty_perf += 1 - except Exception: - continue - score = non_empty_perf if has_perf_column else 0 - if score > best_score or (score == best_score and total_rows > best_rows): - best_score = score - best_rows = total_rows - best_candidate = candidate + best_candidate = result_csv.select_best(candidates) if best_candidate is None: - return candidates[0] - if best_score > 0: + return None + with_metric, _ = result_csv.count_rows(best_candidate) + if with_metric > 0 and len(candidates) > 1: self.console.print( - f"[dim] Selected multiple_results CSV with {best_score} non-empty performance rows: {best_candidate}[/dim]" + f"[dim] Selected multiple_results CSV with {with_metric} non-empty performance rows: {best_candidate}[/dim]" ) return best_candidate @@ -2041,10 +2012,12 @@ def collect_results(self, deployment_id: str) -> Dict[str, Any]: model_key, {} ) if model_key else {} - # Multiple results path: resolve CSV from job_dir/node_*, then cwd/run_directory - mult_res = model_info_for_entry.get("multiple_results") + # Results CSV: the declared name where the nodes left it, and, when the card named + # nothing or named something the script did not write, a file whose header says what + # it is. The declared file always wins when it is there. + mult_res = (model_info_for_entry.get("multiple_results") or "").strip() + resolved_csv: Optional[Path] = None if mult_res: - resolved_csv: Optional[Path] = None # Multi-node: gather all node CSVs and pick the one with the most # non-empty performance rows (master CSV may be empty while a worker # holds the real numbers) instead of taking the first node that has @@ -2062,77 +2035,99 @@ def collect_results(self, deployment_id: str) -> Dict[str, Any]: resolved_csv = Path(mult_res) if not resolved_csv and Path("run_directory", mult_res).is_file(): resolved_csv = Path("run_directory", mult_res) - if resolved_csv: - self._ensure_perf_csv_exists() - gpu_arch = "" - if per_node_metrics: - gpu_arch = per_node_metrics[0].get("gpu_architecture", "") or "" - common_info = self._build_common_info_dict( - model_info_for_entry, build_info, deployment_id, gpu_arch + if resolved_csv is None: + # Only this job's own directories: the submission directory is shared with every + # other model, and a CSV there says nothing about which run wrote it. + search_dirs: List[Path] = [job_dir] + search_dirs += [job_dir / f"node_{i}" for i in range(self.nodes)] + search_dirs += [Path("run_directory")] + discovery = result_csv.discover(search_dirs) + if discovery.winner is not None: + resolved_csv = discovery.winner + origin = ( + f"'{mult_res}' was declared but no node produced it" + if mult_res + else "the model card declares no multiple_results" ) - common_info_path = Path("common_info.json") - with open(common_info_path, "w", encoding="utf-8") as f: - json.dump(common_info, f, indent=2) - update_perf_csv( - perf_csv="perf.csv", - multiple_results=str(resolved_csv), - common_info=str(common_info_path), - model_name=model_info_for_entry.get("name", model_name), + self.console.print( + f"[yellow] Reporting from '{resolved_csv}', found by its header " + f"({origin})[/yellow]" ) - scripts_path = model_info_for_entry.get("scripts", "") - scripts_base_dir = scripts_base_dir_from(scripts_path) - num_entries = update_perf_super_json( - perf_super_json="perf_super.json", - multiple_results=str(resolved_csv), - common_info=str(common_info_path), - model_name=model_info_for_entry.get("name", model_name), - scripts_base_dir=scripts_base_dir, + elif mult_res: + for line in result_csv.describe(discovery, limit=3): + self.console.print(f"[dim] {line}[/dim]") + if resolved_csv: + self._ensure_perf_csv_exists() + gpu_arch = "" + if per_node_metrics: + gpu_arch = per_node_metrics[0].get("gpu_architecture", "") or "" + common_info = self._build_common_info_dict( + model_info_for_entry, build_info, deployment_id, gpu_arch + ) + common_info_path = Path("common_info.json") + with open(common_info_path, "w", encoding="utf-8") as f: + json.dump(common_info, f, indent=2) + update_perf_csv( + perf_csv="perf.csv", + multiple_results=str(resolved_csv), + common_info=str(common_info_path), + model_name=model_info_for_entry.get("name", model_name), + ) + scripts_path = model_info_for_entry.get("scripts", "") + scripts_base_dir = scripts_base_dir_from(scripts_path) + num_entries = update_perf_super_json( + perf_super_json="perf_super.json", + multiple_results=str(resolved_csv), + common_info=str(common_info_path), + model_name=model_info_for_entry.get("name", model_name), + scripts_base_dir=scripts_base_dir, + ) + update_perf_super_csv( + perf_super_json="perf_super.json", + perf_super_csv="perf_super.csv", + num_entries=num_entries, + ) + results["perf_files"] = [str(Path("perf.csv").resolve())] + import csv as _csv + # A row with a number in it is a measurement, and it stays one even when the + # run around it did not finish: a model that reports four precisions and + # crashes on two still measured the other two. The verdict on the run as a + # whole travels separately, in results["incomplete"]. + rows_collected = 0 + try: + with open(resolved_csv, "r", encoding="utf-8", errors="ignore") as f: + reader = _csv.DictReader(f) + for row in reader: + row = {k.strip(): v for k, v in row.items() if k} + if row.get("performance") and row.get("metric"): + rows_collected += 1 + results["successful_runs"].append({ + "model": model_info_for_entry.get("name", "") + "_" + row.get("model", ""), + "status": "SUCCESS", + "performance": str(row.get("performance", "")), + "metric": row.get("metric", ""), + "duration": row.get("test_duration", ""), + "gpu_arch": gpu_arch, + "deployment": "slurm", + "machine": deployment_id, + }) + except Exception: + pass + if job_failed: + self._record_verdict( + results, + model_info_for_entry.get("name", model_name), + node_failures, + job_state, + rows_collected, ) - update_perf_super_csv( - perf_super_json="perf_super.json", - perf_super_csv="perf_super.csv", - num_entries=num_entries, + else: + self.console.print( + f"[green]✓ Updated perf.csv, perf_super.* from multiple_results (Docker-compatible)[/green]" ) - results["perf_files"] = [str(Path("perf.csv").resolve())] - import csv as _csv - # A row with a number in it is a measurement, and it stays one even when the - # run around it did not finish: a model that reports four precisions and - # crashes on two still measured the other two. The verdict on the run as a - # whole travels separately, in results["incomplete"]. - rows_collected = 0 - try: - with open(resolved_csv, "r", encoding="utf-8", errors="ignore") as f: - reader = _csv.DictReader(f) - for row in reader: - row = {k.strip(): v for k, v in row.items() if k} - if row.get("performance") and row.get("metric"): - rows_collected += 1 - results["successful_runs"].append({ - "model": model_info_for_entry.get("name", "") + "_" + row.get("model", ""), - "status": "SUCCESS", - "performance": str(row.get("performance", "")), - "metric": row.get("metric", ""), - "duration": row.get("test_duration", ""), - "gpu_arch": gpu_arch, - "deployment": "slurm", - "machine": deployment_id, - }) - except Exception: - pass - if job_failed: - self._record_verdict( - results, - model_info_for_entry.get("name", model_name), - node_failures, - job_state, - rows_collected, - ) - else: - self.console.print( - f"[green]✓ Updated perf.csv, perf_super.* from multiple_results (Docker-compatible)[/green]" - ) - return results - # multiple_results set but CSV not found: fall through to single-result path (may write FAILURE) + return results + # No results CSV, declared or discovered: fall through to the single-result path, + # which reports from the per-node logs and may write FAILURE. if self.nodes > 1 and per_node_metrics: launcher_type = self.distributed_config.get("launcher", "torchrun") diff --git a/src/madengine/deployment/templates/kubernetes/job.yaml.j2 b/src/madengine/deployment/templates/kubernetes/job.yaml.j2 index 320d049f..d510a585 100644 --- a/src/madengine/deployment/templates/kubernetes/job.yaml.j2 +++ b/src/madengine/deployment/templates/kubernetes/job.yaml.j2 @@ -1,3 +1,17 @@ +{%- macro copy_result_csvs(destination, sources) %} + # A model card need not declare its results CSV, and the name it declares need + # not be the one the script wrote, so take anything shaped like one: a header + # carrying model, performance and metric, in any order. madengine's own outputs + # are skipped -- they match by construction, and copying one back would feed the + # run its own verdict. + for _result_csv in {{ sources }}; do + [ -f "$_result_csv" ] || continue + case "$(basename "$_result_csv")" in perf.csv|perf_super*|perf_entry*) continue ;; esac + if head -n 1 "$_result_csv" 2>/dev/null | tr -d ' "\r' | tr 'A-Z' 'a-z' | awk -F, '{m=0;p=0;t=0;for(i=1;i<=NF;i++){if($i=="model")m=1;if($i=="performance")p=1;if($i=="metric")t=1}exit !(m&&p&&t)}'; then + cp "$_result_csv" {{ destination }}/ 2>/dev/null || true + fi + done +{%- endmacro -%} apiVersion: batch/v1 kind: Job metadata: @@ -363,6 +377,7 @@ spec: fi fi {% endif %} +{{ copy_result_csvs('/results/${HOSTNAME}', '/workspace/*.csv /workspace/run_directory/*.csv ./*.csv ../*.csv') }} # Copy environment details if ls *_env.csv 1> /dev/null 2>&1; then @@ -552,6 +567,7 @@ spec: fi fi {% endif %} +{{ copy_result_csvs('/results/${HOSTNAME}', '/workspace/*.csv /workspace/run_directory/*.csv ./*.csv ../*.csv') }} # Copy environment details if ls *_env.csv 1> /dev/null 2>&1; then diff --git a/src/madengine/deployment/templates/slurm/job.sh.j2 b/src/madengine/deployment/templates/slurm/job.sh.j2 index 54becfd4..2f39f1c6 100644 --- a/src/madengine/deployment/templates/slurm/job.sh.j2 +++ b/src/madengine/deployment/templates/slurm/job.sh.j2 @@ -1,3 +1,16 @@ +{%- macro copy_result_csvs(destination, sources) %} + # A model card need not declare its results CSV, and the name it declares need not be + # the one the script wrote, so take anything shaped like one: a header carrying model, + # performance and metric, in any order. madengine's own outputs are skipped -- they + # match by construction, and copying one back would feed the run its own verdict. + for _result_csv in {{ sources }}; do + [ -f "$_result_csv" ] || continue + case "$(basename "$_result_csv")" in perf.csv|perf_super*|perf_entry*) continue ;; esac + if head -n 1 "$_result_csv" 2>/dev/null | tr -d ' "\r' | tr 'A-Z' 'a-z' | awk -F, '{m=0;p=0;t=0;for(i=1;i<=NF;i++){if($i=="model")m=1;if($i=="performance")p=1;if($i=="metric")t=1}exit !(m&&p&&t)}'; then + cp "$_result_csv" {{ destination }}/ 2>/dev/null || true + fi + done +{%- endmacro -%} #!/bin/bash #SBATCH --job-name=madengine-{{ model_name }} #SBATCH --output={{ output_dir }}/madengine-{{ model_name }}_%j_%t.out @@ -732,6 +745,7 @@ if [ $TASK_EXIT -eq 0 ]; then if [ -f "$WORKSPACE/run_directory/{{ multiple_results }}" ]; then cp "$WORKSPACE/run_directory/{{ multiple_results }}" "$NODE_COLLECTION_DIR/" 2>/dev/null || true; fi if [ -f "$WORKSPACE/{{ multiple_results }}" ]; then cp "$WORKSPACE/{{ multiple_results }}" "$NODE_COLLECTION_DIR/" 2>/dev/null || true; fi {% endif %} +{{ copy_result_csvs('"$NODE_COLLECTION_DIR"', '"$WORKSPACE"/*.csv "$WORKSPACE"/run_directory/*.csv') }} if [ -d "$WORKSPACE/rocprof_output" ]; then cp -r "$WORKSPACE/rocprof_output" "$NODE_COLLECTION_DIR/" 2>/dev/null || true; fi echo " ✓ Node ${SLURM_PROCID} artifacts copied" echo "========================================================================" @@ -793,18 +807,19 @@ $MAD_CLI_COMMAND run \ EXIT_CODE=$? -# Single-node: copy multiple_results CSV to job collection dir so collect_results finds it -{% if multiple_results %} +# Single-node: copy the results CSV to the job collection dir so collect_results finds it if [ $EXIT_CODE -eq 0 ]; then SUBMISSION_DIR={{ manifest_file | dirname }} JOB_COLLECTION_DIR="${SUBMISSION_DIR}/{{ output_dir }}/{{ model_name }}/${SLURM_JOB_ID}" NODE_COLLECTION_DIR="${JOB_COLLECTION_DIR}/node_0" mkdir -p "$NODE_COLLECTION_DIR" + {% if multiple_results %} if [ -f "$WORKSPACE/run_directory/{{ multiple_results }}" ]; then cp "$WORKSPACE/run_directory/{{ multiple_results }}" "$NODE_COLLECTION_DIR/" 2>/dev/null || true; fi if [ -f "$WORKSPACE/{{ multiple_results }}" ]; then cp "$WORKSPACE/{{ multiple_results }}" "$NODE_COLLECTION_DIR/" 2>/dev/null || true; fi + {% endif %} +{{ copy_result_csvs('"$NODE_COLLECTION_DIR"', '"$WORKSPACE"/*.csv "$WORKSPACE"/run_directory/*.csv') }} fi {% endif %} -{% endif %} # ============================================================================= # Job Completion diff --git a/src/madengine/execution/container_runner.py b/src/madengine/execution/container_runner.py index 0c2d597b..6ca41d8b 100644 --- a/src/madengine/execution/container_runner.py +++ b/src/madengine/execution/container_runner.py @@ -32,6 +32,7 @@ flatten_tags, ) from madengine.reporting.update_perf_super import update_perf_super_json, update_perf_super_csv +from madengine.reporting import result_csv from madengine.utils.gpu_config import resolve_runtime_gpus from madengine.deployment.common import canonicalize_distributed_launcher from madengine.utils.config_parser import ConfigParser @@ -182,6 +183,69 @@ def _resolve_multiple_results_path(multiple_results: str, model_dir: str) -> typ return None +def _settle_results_csv( + model_info: dict, + model_dir: str, + say: typing.Callable[[str], None], + min_mtime: typing.Optional[float] = None, +) -> typing.Tuple[typing.Optional[str], typing.Optional["result_csv.Discovery"]]: + """The results CSV this run reports from, and the search that had to find it. + + A declared file that exists always wins: the card said where to look and it was right. + Otherwise the file is looked for by its header, because the name in the card is not the + contract anyone actually keeps -- the field is optional, and the script that writes the + file computes its own path in another repository. The three cases a reader has to be + able to tell apart afterwards are a typo, an omission, and nothing found at all, so + each says so on its way past. + + Args: + model_info: the model entry, read for ``multiple_results`` and ``name`` + model_dir: the run directory, searched before the workspace root + say: where messages go, e.g. ``rich_console.print`` + min_mtime: ignore files older than this, so the model that ran before this one in + the same directory does not get its results reported twice + + Returns: + (path to the CSV or None, the discovery that ran or None when the card was right) + """ + declared = (model_info.get("multiple_results") or "").strip() + resolved = _resolve_multiple_results_path(declared, model_dir) if declared else None + if resolved: + return resolved, None + + discovery = result_csv.discover([model_dir, os.getcwd()], min_mtime=min_mtime) + if declared: + # This is a contract, and the model script held up its end or it did not: say + # which paths were searched and what the container was told, so the next person + # does not have to read this code. + say( + f"[yellow]Warning: model '{model_info.get('name', '')}' declares " + f"multiple_results='{declared}' but no such file was produced.[/yellow]" + ) + say( + f"[yellow] Searched: {os.path.abspath(declared)} and " + f"{os.path.join(os.path.abspath(model_dir), declared)}[/yellow]" + ) + say( + f"[yellow] The container was given MAD_OUTPUT_CSV='{declared}'; the run " + f"script must write a CSV there with 'performance' and 'metric' columns, " + f"relative to its working directory.[/yellow]" + ) + if discovery.winner is None: + return None, discovery + + origin = ( + "the declared file was not there" + if declared + else "the model card declares no multiple_results" + ) + say( + f"[yellow] Reporting from '{discovery.winner}', found by its header: it carries " + f"model, performance and metric ({origin}).[/yellow]" + ) + return str(discovery.winner), discovery + + def _docker_image_exists_locally(image: str) -> bool: """Return True if ``docker image inspect`` succeeds for *image* (argv list; no shell).""" try: @@ -1674,6 +1738,11 @@ def run_container( pre_encapsulate_post_scripts["post_scripts"], ) + # The results CSV this run will report from: declared by the card and + # present, or found by its header. Settled once, below, and reused by + # the reporting step so discovery is not lost on the way there. + resolved_results_csv: typing.Optional[str] = None + if skip_model_run: run_results["status"] = "SKIPPED" self.rich_console.print( @@ -1701,68 +1770,27 @@ def run_container( if multiple_results: multiple_results = multiple_results.strip() - if multiple_results: - resolved_path = _resolve_multiple_results_path( - multiple_results, model_dir - ) - if not resolved_path: - # This is a contract, and the model script held up - # its end or it did not: say which paths were - # searched and what the container was told, so the - # next person does not have to read this code. - self.rich_console.print( - f"[yellow]Warning: model '{model_info.get('name', '')}' " - f"declares multiple_results='{multiple_results}' but no " - f"such file was produced.[/yellow]" - ) - self.rich_console.print( - f"[yellow] Searched: {os.path.abspath(multiple_results)}" - f" and {os.path.join(os.path.abspath(model_dir), multiple_results)}[/yellow]" - ) - self.rich_console.print( - f"[yellow] The container was given " - f"MAD_OUTPUT_CSV='{multiple_results}'; the run script must " - f"write a CSV there with 'performance' and 'metric' " - f"columns, relative to its working directory.[/yellow]" + # The workspace root is shared by every model in the run, so + # only a file written during this model's run can be this + # model's result. + resolved_path, discovery = _settle_results_csv( + model_info, + model_dir, + self.rich_console.print, + min_mtime=test_start_time, + ) + + if resolved_path: + resolved_results_csv = resolved_path + run_results["performance"] = resolved_path + # Same reading of the file as discovery does, so a declared + # file and a found one are judged by one rule. + problem = result_csv.metric_rejection_reason(resolved_path) + if problem: + print( + f"Error: {resolved_path} produced no metric: {problem}." ) run_results["performance"] = None - else: - run_results["performance"] = resolved_path - # Validate multiple results file format using proper CSV parsing - try: - import csv - with open(resolved_path, "r") as f: - csv_reader = csv.DictReader(f) - - # Strip whitespace from fieldnames to handle headers like "model, performance, metric" - csv_reader.fieldnames = [f.strip() for f in csv_reader.fieldnames] - - # Check if 'performance' column exists - if 'performance' not in csv_reader.fieldnames: - print( - f"Error: {resolved_path} has no 'performance' column; " - f"found: {', '.join(csv_reader.fieldnames or []) or '(no header)'}" - ) - run_results["performance"] = None - else: - # Check if at least one row has a non-empty performance value - has_valid_perf = False - for row in csv_reader: - if row.get('performance', '').strip(): - has_valid_perf = True - break - - if not has_valid_perf: - run_results["performance"] = None - print( - f"Error: every row of {resolved_path} has an empty " - f"'performance' value, so the run produced no metric." - ) - except Exception as e: - self.rich_console.print( - f"[yellow]Warning: Could not validate multiple results file: {e}[/yellow]" - ) - run_results["performance"] = None else: # Match the actual output format: "performance: 14164 samples_per_second" # Simple pattern to capture number and metric unit @@ -1803,8 +1831,25 @@ def run_container( run_results["metric"] = "samples_per_second" print(f"✓ Extracted performance (HuggingFace format): {run_results['performance']} {run_results['metric']}") else: - # No performance metrics found - print("Warning: Performance metric not found in expected format 'performance: NUMBER METRIC' or 'train_samples_per_second'") + # Nothing measured. There are two ways a + # model can report a number and both came + # back empty, so say which was tried. + print( + "Warning: no metric found. A model reports one either in a " + "results CSV or in its log; neither had one here." + ) + for line in result_csv.describe(discovery, limit=3): + print(f" {line}") + print( + f" No 'performance: NUMBER METRIC' line and no " + f"'train_samples_per_second' in {log_file_path}" + ) + print( + f" MAD_OUTPUT_CSV was exported as '{multiple_results}'" + if multiple_results + else " MAD_OUTPUT_CSV was not exported: the model card " + "declares no multiple_results" + ) run_results["performance"] = None run_results["metric"] = None @@ -1971,13 +2016,10 @@ def run_container( model_info, build_info, run_results ) - # Handle multiple results if specified - multiple_results = model_info.get("multiple_results", None) - resolved_multiple_results = ( - _resolve_multiple_results_path(multiple_results, model_dir) - if multiple_results - else None - ) + # The results CSV was settled during extraction, declared or + # discovered; resolving it again here would drop the + # discovered one and report an empty row instead. + resolved_multiple_results = resolved_results_csv if ( resolved_multiple_results and run_results.get("status") == "SUCCESS" @@ -2098,13 +2140,22 @@ def run_container( # Ignore errors if no profiler/trace output files exist pass - # Copy multiple_results CSV to workspace root before run_directory is removed - # so SLURM single-node copy can find it at $WORKSPACE/{{ multiple_results }} - mult_res = (model_info.get("multiple_results") or "").strip() - if mult_res: + # Copy the results CSV to workspace root before run_directory is + # removed, so the SLURM single-node copy can find it at + # $WORKSPACE/. A discovered file needs this as much as a + # declared one, and it may be the only copy there is. + results_csv_names = [] + declared_name = (model_info.get("multiple_results") or "").strip() + if declared_name: + results_csv_names.append(declared_name) + if resolved_results_csv: + found_name = os.path.basename(resolved_results_csv) + if found_name not in results_csv_names: + results_csv_names.append(found_name) + for name in results_csv_names: try: model_docker.sh( - _cp_model_dir_file_to_cwd_cmd(model_dir, mult_res) + _cp_model_dir_file_to_cwd_cmd(model_dir, name) ) except Exception: pass diff --git a/src/madengine/reporting/result_csv.py b/src/madengine/reporting/result_csv.py new file mode 100644 index 00000000..734cfcdf --- /dev/null +++ b/src/madengine/reporting/result_csv.py @@ -0,0 +1,250 @@ +"""Find the results CSV a model wrote, by the shape of its header. + +A model card can declare where its results CSV lives (``multiple_results``), but the +field is optional and its value has to agree with a path a script in another repository +builds by hand. When the two disagree madengine used to fall back to scraping the log, +find nothing there either, and record an empty FAILURE row for a run that had measured +perfectly well. The shape of the file is the sturdier contract: a results CSV is one +whose header carries ``model``, ``performance`` and ``metric``, in any order and with +any number of extra columns. A declared file that exists still wins; discovery only +answers when the declaration is absent or does not resolve. +""" + +import csv +import os +import typing +from pathlib import Path + +#: A results CSV is recognised by these three columns, whatever else it carries. +REQUIRED_COLUMNS = ("model", "performance", "metric") + +#: madengine's own outputs satisfy the predicate by construction, so discovering one +#: would feed a run its own previous verdict back as input. +OWN_OUTPUT_NAMES = frozenset({"perf.csv"}) +OWN_OUTPUT_PREFIXES = ("perf_super", "perf_entry") + + +def normalise_columns(fieldnames: typing.Optional[typing.Iterable[str]]) -> typing.List[str]: + """Column names as they are compared: stripped of space, quotes and case.""" + return [ + (name or "").strip().strip('"').strip("'").strip().lower() + for name in (fieldnames or []) + ] + + +def read_columns(path: typing.Union[str, Path]) -> typing.Optional[typing.List[str]]: + """The normalised header of *path*, or None when it cannot be read as CSV.""" + try: + # utf-8-sig so a byte-order mark does not hide behind the first column name. + with open(path, "r", newline="", encoding="utf-8-sig", errors="ignore") as handle: + for row in csv.reader(handle): + return normalise_columns(row) + except (OSError, csv.Error): + return None + return None + + +def missing_columns(path: typing.Union[str, Path]) -> typing.Optional[typing.List[str]]: + """Which required columns *path* lacks; None when it is not readable as CSV.""" + columns = read_columns(path) + if columns is None: + return None + return [name for name in REQUIRED_COLUMNS if name not in columns] + + +def has_result_shape(path: typing.Union[str, Path]) -> bool: + """True when *path* parses as CSV and its header carries all required columns.""" + return missing_columns(path) == [] + + +def is_own_output(path: typing.Union[str, Path]) -> bool: + """True for the files madengine writes itself (perf.csv and the super/entry family).""" + name = os.path.basename(str(path)).lower() + return name in OWN_OUTPUT_NAMES or name.startswith(OWN_OUTPUT_PREFIXES) + + +def count_rows(path: typing.Union[str, Path]) -> typing.Tuple[int, int]: + """(rows with a non-empty ``performance``, total rows) in *path*.""" + with_metric = 0 + total = 0 + try: + with open(path, "r", newline="", encoding="utf-8-sig", errors="ignore") as handle: + reader = csv.DictReader(handle) + reader.fieldnames = normalise_columns(reader.fieldnames) + if "performance" not in (reader.fieldnames or []): + return (0, 0) + for row in reader: + total += 1 + value = row.get("performance") or "" + if str(value).strip(): + with_metric += 1 + except (OSError, csv.Error): + return (0, 0) + return (with_metric, total) + + +def metric_rejection_reason(path: typing.Union[str, Path]) -> typing.Optional[str]: + """Why *path* yields no metric, or None when it does. + + This is the question to ask about a file the model card named: it was declared, so it + is trusted to be the results CSV, and all that is left to check is whether a metric + can be read out of it. + """ + columns = read_columns(path) + if columns is None: + return "not readable as CSV" + if "performance" not in columns: + found = ", ".join(columns) if columns else "(no header)" + return f"no 'performance' column; found: {found}" + if count_rows(path)[0] == 0: + return "every row has an empty 'performance' value" + return None + + +def rejection_reason(path: typing.Union[str, Path]) -> typing.Optional[str]: + """Why *path* cannot be *discovered* as a results CSV, or None when it can. + + Stricter than :func:`metric_rejection_reason` on purpose: nobody pointed at this file, + so it has to look like a results CSV on its own -- all three columns, and a number in + at least one row -- before madengine reads a run's verdict out of it. + """ + if is_own_output(path): + return "madengine's own output" + missing = missing_columns(path) + if missing is None: + return "not readable as CSV" + if missing: + columns = read_columns(path) or [] + found = ", ".join(columns) if columns else "(no header)" + return f"header lacks {', '.join(missing)}; found: {found}" + return metric_rejection_reason(path) + + +def rank_candidates( + candidates: typing.Sequence[typing.Union[str, Path]] +) -> typing.List[Path]: + """Candidates best first: most measured rows, then most rows, then newest, then name. + + In a multi-node run every node copies its own CSV into the collection directory and + only some of them observed the final throughput, so the file with the most non-empty + ``performance`` rows is the one carrying the measurement. The remaining keys exist so + that two runs over the same inputs pick the same file: newest first, and failing that + the order the caller offered them in, which is the order the directories were searched. + """ + + def sort_key(entry: typing.Tuple[int, Path]) -> typing.Tuple: + position, candidate = entry + with_metric, total = count_rows(candidate) + try: + mtime = os.path.getmtime(candidate) + except OSError: + mtime = 0.0 + return (-with_metric, -total, -mtime, position) + + numbered = list(enumerate(Path(c) for c in candidates)) + return [candidate for _, candidate in sorted(numbered, key=sort_key)] + + +def select_best( + candidates: typing.Sequence[typing.Union[str, Path]] +) -> typing.Optional[Path]: + """The best of *candidates*, or None when there are none.""" + ranked = rank_candidates(candidates) + return ranked[0] if ranked else None + + +def _staleness_reason( + path: typing.Union[str, Path], min_mtime: typing.Optional[float] +) -> typing.Optional[str]: + """"Written before this run started", when that is what happened.""" + if min_mtime is None: + return None + try: + mtime = os.path.getmtime(path) + except OSError: + return "not readable" + # A second of slack: file timestamps and the clock this run started by do not always + # come from the same source, and a file written in the first moments of the run is + # this run's. + if mtime < min_mtime - 1.0: + return "written before this run started" + return None + + +class Discovery(typing.NamedTuple): + """What a search found, and enough of what it discarded to explain itself.""" + + winner: typing.Optional[Path] + candidates: typing.List[Path] + rejected: typing.List[typing.Tuple[Path, str]] + searched: typing.List[Path] + seen: int + + +def discover( + search_dirs: typing.Sequence[typing.Union[str, Path]], + excluded: typing.Sequence[typing.Union[str, Path]] = (), + min_mtime: typing.Optional[float] = None, +) -> Discovery: + """Look for a results CSV directly inside each of *search_dirs*. + + Depth 1 only: a training run can leave hundreds of CSVs behind and walking the tree + would turn a diagnostic into a scan. Directories are searched in the order given, + duplicates of the same file are collapsed, and *excluded* paths are skipped along + with madengine's own outputs. + + *min_mtime* is how a shared working directory stays safe: several models run one after + another in the same place, and the one that ran before this one left its results CSV + behind. A file that was not written during this run cannot be this run's result. + """ + excluded_real = {os.path.realpath(str(path)) for path in excluded} + searched: typing.List[Path] = [] + seen_files: typing.Dict[str, Path] = {} + + for directory in search_dirs: + if directory is None: + continue + as_path = Path(directory) + if not as_path.is_dir(): + continue + real_dir = os.path.realpath(str(as_path)) + if real_dir in {os.path.realpath(str(d)) for d in searched}: + continue + searched.append(as_path) + for entry in sorted(as_path.glob("*.csv")): + if not entry.is_file(): + continue + real = os.path.realpath(str(entry)) + if real in excluded_real or real in seen_files: + continue + seen_files[real] = entry + + candidates: typing.List[Path] = [] + rejected: typing.List[typing.Tuple[Path, str]] = [] + for entry in seen_files.values(): + reason = _staleness_reason(entry, min_mtime) or rejection_reason(entry) + if reason is None: + candidates.append(entry) + else: + rejected.append((entry, reason)) + + ranked = rank_candidates(candidates) + return Discovery( + winner=ranked[0] if ranked else None, + candidates=ranked, + rejected=rejected, + searched=searched, + seen=len(seen_files), + ) + + +def describe(discovery: Discovery, limit: int = 4) -> typing.List[str]: + """A few lines saying where the search looked and why each file was discarded.""" + where = ", ".join(str(path) for path in discovery.searched) or "(no directory to search)" + lines = [f"Searched for a results CSV in: {where}", f"CSV files seen: {discovery.seen}"] + for path, reason in discovery.rejected[:limit]: + lines.append(f" rejected {path}: {reason}") + remaining = len(discovery.rejected) - limit + if remaining > 0: + lines.append(f" ... and {remaining} more rejected the same way") + return lines diff --git a/tests/unit/test_result_csv_discovery.py b/tests/unit/test_result_csv_discovery.py new file mode 100644 index 00000000..81e79b14 --- /dev/null +++ b/tests/unit/test_result_csv_discovery.py @@ -0,0 +1,352 @@ +"""Finding the results CSV a model wrote, when the card did not say where it is. + +The header fixtures here are copied verbatim out of MAD's own scripts rather than +invented, so a change in what the scripts emit shows up as a failure here: + +- ``model,performance,metric`` scripts/dummy/run_multi.sh:1 +- ``model,performance,metric`` scripts/mochi/run_mochi.sh:70 +- ``model, performance, metric`` scripts/pyt_chai1_inference/run.sh:39 +- ``model,performance,metric,mode,precision,...`` scripts/primus_megatron-lm/ + primus_megatron-lm_benchmark_report.sh:283 +- ``hf_pipeline_tag,model,...,performance,metric,unit`` scripts/atom/run_atom.py:43 + +and the near-misses that must stay rejected: + +- ``Model,xP/yD,ISL,...`` scripts/sglang_disagg/benchmark_parser.py:188 +- ``model_name,model_unique_name,...`` scripts/kvcache_transfer_bench/kv_cache_estimator.py:1319 +""" + +import os +import time +from pathlib import Path + +import pytest + +from madengine.reporting import result_csv + + +# Real headers, as the scripts write them. +DUMMY_HEADER = "model,performance,metric" +CHAI_HEADER = "model, performance, metric" +PRIMUS_HEADER = ( + "model,performance,metric,mode,precision,batch_size,global_batch_size," + "seq_len,device,num_gpus" +) +# The three columns are neither first nor adjacent here. +ATOM_HEADER = ( + "hf_pipeline_tag,model,benchmark,tp,inp,out,kv_cache_dtype,num_prompts," + "max_concurrency,bs,cmd,performance,metric,unit" +) + + +def write_csv(path: Path, header: str, *rows: str) -> Path: + """Write a CSV with *header* and *rows* verbatim, and return the path.""" + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text("\n".join([header, *rows]) + "\n", encoding="utf-8") + return path + + +def make_result_csv(path: Path, header: str = DUMMY_HEADER, count: int = 1) -> Path: + """A results CSV with *count* measured rows.""" + rows = [f"m{i},{100 + i},tokens_per_second" for i in range(count)] + if header == PRIMUS_HEADER: + rows = [ + f"m{i},{100 + i},tok/s/GPU,train,BF16,1,8,4096,gfx950,16" + for i in range(count) + ] + return write_csv(path, header, *rows) + + +class TestHeaderPredicate: + """A results CSV is recognised by its columns, however they are spelled.""" + + @pytest.mark.parametrize("header", [DUMMY_HEADER, CHAI_HEADER, PRIMUS_HEADER]) + def test_real_headers_are_recognised(self, tmp_path, header): + assert result_csv.has_result_shape(make_result_csv(tmp_path / "r.csv", header)) + + def test_the_three_columns_need_not_be_first_or_adjacent(self, tmp_path): + path = write_csv(tmp_path / "r.csv", ATOM_HEADER, "text,m," + ",".join([""] * 9) + ",42,tok/s,x") + assert result_csv.has_result_shape(path) + + @pytest.mark.parametrize( + "header", + [ + # scripts/sglang_disagg/benchmark_parser.py:188 -- a Model column, no metric + "Model,xP/yD,ISL,OSL,Concurrency,Request Throughput (req/s)", + # kv_cache_estimator.py:1319 -- model_name is a different column + "model_name,model_unique_name,concurrency,seq_length,kv_cache_mb", + # sglang_benchmark_report.py:99 -- throughput instead of performance/metric + "model,total_throughput (tok/sec),output_throughput (tok/sec),tp", + ], + ) + def test_near_miss_headers_are_rejected(self, tmp_path, header): + path = write_csv(tmp_path / "r.csv", header, "a,b,c,d") + assert not result_csv.has_result_shape(path) + + def test_a_raw_training_log_named_csv_is_rejected(self, tmp_path): + """benchmark_report.sh:106 tees stdout into a file with a .csv extension.""" + path = tmp_path / "primus-megatron-Megatron-LM-pretrain.csv" + path.write_text("[INFO] iteration 1/100 | throughput 17967.4\n[INFO] done\n") + assert not result_csv.has_result_shape(path) + + def test_column_order_does_not_matter(self, tmp_path): + path = write_csv(tmp_path / "r.csv", "metric,model,performance", "tok/s,m,42") + assert result_csv.has_result_shape(path) + + def test_quoted_and_uppercase_columns_are_recognised(self, tmp_path): + path = write_csv(tmp_path / "r.csv", '"Model","Performance","Metric"', "m,42,tok/s") + assert result_csv.has_result_shape(path) + + def test_byte_order_mark_does_not_hide_the_first_column(self, tmp_path): + path = tmp_path / "r.csv" + path.write_text("\ufeffmodel,performance,metric\nm,42,tok/s\n", encoding="utf-8") + assert result_csv.has_result_shape(path) + + def test_a_csv_without_the_three_columns_is_not_a_result(self, tmp_path): + path = write_csv(tmp_path / "gpu_info.csv", "gpu,power,temperature", "0,300,45") + assert not result_csv.has_result_shape(path) + assert "header lacks" in result_csv.rejection_reason(path) + + def test_a_file_that_is_not_csv_is_reported_as_such(self, tmp_path): + path = tmp_path / "notes.csv" + path.write_bytes(b"\x00\x01\x02") + assert result_csv.rejection_reason(path) is not None + + def test_missing_performance_column_is_named_for_a_declared_file(self, tmp_path): + path = write_csv(tmp_path / "r.csv", "model,metric", "m,tok/s") + reason = result_csv.metric_rejection_reason(path) + assert "no 'performance' column" in reason + assert "model, metric" in reason + + def test_a_declared_file_needs_only_a_performance_column(self, tmp_path): + """A card that named the file is trusted; only the metric has to be there.""" + path = write_csv(tmp_path / "r.csv", "run,performance", "a,42") + assert result_csv.metric_rejection_reason(path) is None + assert not result_csv.has_result_shape(path) + + def test_all_rows_empty_is_not_a_measurement(self, tmp_path): + path = write_csv(tmp_path / "r.csv", DUMMY_HEADER, "m,,tok/s", "m2, ,tok/s") + assert result_csv.metric_rejection_reason(path) == ( + "every row has an empty 'performance' value" + ) + assert result_csv.rejection_reason(path) is not None + + +class TestOwnOutputsAreNeverInput: + """madengine's own files match the predicate by construction.""" + + @pytest.mark.parametrize( + "name", ["perf.csv", "perf_super.csv", "perf_super_1.csv", "perf_entry.csv"] + ) + def test_own_outputs_are_excluded(self, tmp_path, name): + make_result_csv(tmp_path / name) + assert result_csv.is_own_output(tmp_path / name) + assert result_csv.discover([tmp_path]).winner is None + + def test_own_output_does_not_hide_a_real_result(self, tmp_path): + make_result_csv(tmp_path / "perf.csv", count=9) + make_result_csv(tmp_path / "perf_dummy.csv", count=1) + assert result_csv.discover([tmp_path]).winner == tmp_path / "perf_dummy.csv" + + +class TestRanking: + """Which candidate wins, and the same one every time.""" + + def test_the_file_with_more_measured_rows_wins(self, tmp_path): + thin = make_result_csv(tmp_path / "node_0" / "r.csv", count=1) + rich = make_result_csv(tmp_path / "node_1" / "r.csv", count=8) + assert result_csv.select_best([thin, rich]) == rich + assert result_csv.select_best([rich, thin]) == rich + + def test_an_empty_file_loses_to_a_measured_one(self, tmp_path): + empty = write_csv(tmp_path / "node_0" / "r.csv", DUMMY_HEADER, "m,,tok/s") + measured = make_result_csv(tmp_path / "node_1" / "r.csv", count=1) + assert result_csv.select_best([empty, measured]) == measured + + def test_ties_break_on_the_newer_file(self, tmp_path): + older = make_result_csv(tmp_path / "a" / "r.csv", count=2) + newer = make_result_csv(tmp_path / "b" / "r.csv", count=2) + os.utime(older, (1_000_000, 1_000_000)) + os.utime(newer, (2_000_000, 2_000_000)) + assert result_csv.select_best([older, newer]) == newer + assert result_csv.select_best([newer, older]) == newer + + def test_a_full_tie_keeps_the_order_it_was_given(self, tmp_path): + first = make_result_csv(tmp_path / "a" / "r.csv", count=2) + second = make_result_csv(tmp_path / "b" / "r.csv", count=2) + os.utime(first, (1_000_000, 1_000_000)) + os.utime(second, (1_000_000, 1_000_000)) + assert result_csv.select_best([first, second]) == first + assert result_csv.select_best([second, first]) == second + + def test_no_candidates_is_not_an_error(self): + assert result_csv.select_best([]) is None + + +class TestDiscovery: + """What a depth-1 search finds, and what it refuses to.""" + + def test_a_results_csv_in_the_run_directory_is_found(self, tmp_path): + run_dir = tmp_path / "run_directory" + found = make_result_csv(run_dir / "perf_primus-megatron-Megatron-LM.csv", PRIMUS_HEADER) + assert result_csv.discover([run_dir, tmp_path]).winner == found + + def test_the_parent_of_the_run_directory_is_searched_too(self, tmp_path): + """Primus writes to $(pwd)/../ -- benchmark_report.sh:109.""" + run_dir = tmp_path / "run_directory" + run_dir.mkdir() + found = make_result_csv(tmp_path / "perf_primus-megatron-Megatron-LM.csv", PRIMUS_HEADER) + assert result_csv.discover([run_dir, tmp_path]).winner == found + + def test_the_search_does_not_walk_the_tree(self, tmp_path): + make_result_csv(tmp_path / "deep" / "nested" / "r.csv") + assert result_csv.discover([tmp_path]).winner is None + + def test_a_decoy_csv_is_rejected_and_explained(self, tmp_path): + write_csv(tmp_path / "gpu_info_power.csv", "gpu,power", "0,300") + discovery = result_csv.discover([tmp_path]) + assert discovery.winner is None + assert discovery.seen == 1 + rejected_paths = [str(path) for path, _ in discovery.rejected] + assert str(tmp_path / "gpu_info_power.csv") in rejected_paths + + def test_the_same_file_reached_twice_is_counted_once(self, tmp_path): + make_result_csv(tmp_path / "r.csv") + discovery = result_csv.discover([tmp_path, tmp_path, tmp_path / "missing"]) + assert discovery.seen == 1 + assert len(discovery.searched) == 1 + + def test_a_file_from_an_earlier_model_is_not_this_run(self, tmp_path): + """The workspace root is shared, so a stale CSV must not be adopted.""" + stale = make_result_csv(tmp_path / "r.csv", count=4) + os.utime(stale, (1_000_000, 1_000_000)) + run_started = time.time() + discovery = result_csv.discover([tmp_path], min_mtime=run_started) + assert discovery.winner is None + assert discovery.rejected[0][1] == "written before this run started" + + def test_a_file_written_during_the_run_is_kept(self, tmp_path): + run_started = time.time() + fresh = make_result_csv(tmp_path / "r.csv", count=4) + assert result_csv.discover([tmp_path], min_mtime=run_started).winner == fresh + + def test_an_excluded_path_is_skipped(self, tmp_path): + skip_me = make_result_csv(tmp_path / "r.csv") + assert result_csv.discover([tmp_path], excluded=[skip_me]).winner is None + + +class TestLogReportingModelsStaySilent: + """Some models report only through stdout, e.g. scripts/huggingface_gpt2/run.sh:86. + + Of the 39 cards that declare no multiple_results, four write no results CSV at all; + most of the rest write one to /run_logs on shared storage, which no depth-1 search + reaches. Either way the search must come back empty without adding noise. + """ + + def test_a_directory_with_no_csv_yields_nothing_and_no_noise(self, tmp_path): + (tmp_path / "run.log").write_text("performance: 14164 samples_per_second\n") + discovery = result_csv.discover([tmp_path]) + assert discovery.winner is None + assert discovery.seen == 0 + assert discovery.rejected == [] + + def test_a_missing_directory_is_not_an_error(self, tmp_path): + discovery = result_csv.discover([tmp_path / "never_created"]) + assert discovery.winner is None + assert discovery.searched == [] + + +class TestDescription: + """The diagnostic says where it looked and why each file was refused.""" + + def test_it_names_the_directories_and_the_reasons(self, tmp_path): + write_csv(tmp_path / "gpu_info.csv", "gpu,power", "0,300") + lines = result_csv.describe(result_csv.discover([tmp_path])) + assert any(str(tmp_path) in line for line in lines) + assert any("CSV files seen: 1" in line for line in lines) + assert any("header lacks" in line for line in lines) + + def test_it_stays_short_when_there_is_a_lot_to_say(self, tmp_path): + for index in range(12): + write_csv(tmp_path / f"decoy_{index}.csv", "gpu,power", "0,300") + lines = result_csv.describe(result_csv.discover([tmp_path]), limit=3) + assert len(lines) == 6 + assert "and 9 more" in lines[-1] + + +class TestSettleResultsCsv: + """Which file the Docker path reports from, and what it says about the choice.""" + + def settle(self, tmp_path, monkeypatch, declared=None, min_mtime=None): + from madengine.execution.container_runner import _settle_results_csv + + monkeypatch.chdir(tmp_path) + said = [] + model_info = {"name": "dummy"} + if declared is not None: + model_info["multiple_results"] = declared + path, discovery = _settle_results_csv( + model_info, "run_directory", said.append, min_mtime=min_mtime + ) + return path, discovery, "\n".join(said) + + def test_a_declared_file_that_exists_wins(self, tmp_path, monkeypatch): + """Even against a candidate with more rows: the card said where to look.""" + (tmp_path / "run_directory").mkdir() + make_result_csv(tmp_path / "run_directory" / "declared.csv", count=1) + make_result_csv(tmp_path / "run_directory" / "richer.csv", count=9) + path, discovery, said = self.settle(tmp_path, monkeypatch, declared="declared.csv") + assert Path(path).name == "declared.csv" + assert discovery is None + assert said == "" + + def test_a_typo_recovers_and_says_both_things(self, tmp_path, monkeypatch): + (tmp_path / "run_directory").mkdir() + make_result_csv(tmp_path / "run_directory" / "perf_dummy.csv", count=2) + path, _, said = self.settle(tmp_path, monkeypatch, declared="perf_dumy.csv") + assert Path(path).name == "perf_dummy.csv" + assert "declares multiple_results='perf_dumy.csv' but no such file" in said + assert "found by its header" in said + assert "the declared file was not there" in said + + def test_an_undeclared_file_is_found_and_named(self, tmp_path, monkeypatch): + (tmp_path / "run_directory").mkdir() + make_result_csv(tmp_path / "run_directory" / "perf_dummy.csv", count=2) + path, _, said = self.settle(tmp_path, monkeypatch) + assert Path(path).name == "perf_dummy.csv" + assert "declares no multiple_results" in said + assert "Warning: model" not in said + + def test_the_workspace_root_is_searched_after_the_run_directory(self, tmp_path, monkeypatch): + (tmp_path / "run_directory").mkdir() + make_result_csv(tmp_path / "perf_primus-megatron-Megatron-LM.csv", PRIMUS_HEADER, count=4) + path, _, said = self.settle(tmp_path, monkeypatch) + assert Path(path).name == "perf_primus-megatron-Megatron-LM.csv" + + def test_nothing_found_reports_nothing_and_does_not_raise(self, tmp_path, monkeypatch): + (tmp_path / "run_directory").mkdir() + path, discovery, said = self.settle(tmp_path, monkeypatch) + assert path is None + assert discovery.winner is None + assert said == "" + + def test_a_log_only_model_stays_silent(self, tmp_path, monkeypatch): + """A card with no CSV writer must not gain a warning it never had.""" + (tmp_path / "run_directory").mkdir() + (tmp_path / "run.log").write_text("performance: 14164 samples_per_second\n") + path, _, said = self.settle(tmp_path, monkeypatch) + assert path is None + assert said == "" + + def test_a_previous_model_result_is_not_adopted(self, tmp_path, monkeypatch): + (tmp_path / "run_directory").mkdir() + stale = make_result_csv(tmp_path / "perf_other_model.csv", count=3) + os.utime(stale, (1_000_000, 1_000_000)) + path, _, _ = self.settle(tmp_path, monkeypatch, min_mtime=time.time()) + assert path is None + + def test_madengine_own_perf_csv_is_never_adopted(self, tmp_path, monkeypatch): + (tmp_path / "run_directory").mkdir() + make_result_csv(tmp_path / "perf.csv", count=5) + path, _, _ = self.settle(tmp_path, monkeypatch) + assert path is None diff --git a/tests/unit/test_result_csv_templates.py b/tests/unit/test_result_csv_templates.py new file mode 100644 index 00000000..9d89f179 --- /dev/null +++ b/tests/unit/test_result_csv_templates.py @@ -0,0 +1,115 @@ +"""The job templates must copy a results CSV nobody declared. + +The per-node copy used to sit inside ``{% if multiple_results %}``, so a card without the +field left nothing behind on the node and there was nothing for the collector to find -- +discovery on the login node would have had no input. These tests pin the sweep into all +four blocks: the multi-node task script and the single-node tail of the SLURM job, and +both container blocks of the Kubernetes job. + +Copyright (c) Advanced Micro Devices, Inc. All rights reserved. +""" + +from pathlib import Path + +import madengine.deployment as deployment_package +from madengine.deployment.base import create_jinja_env +from madengine.deployment.slurm import SlurmDeployment + + +#: The shape test the shell does, in the one form it is written in both templates. +SHAPE_TEST = 'if($i=="model")m=1;if($i=="performance")p=1;if($i=="metric")t=1' +SWEEP_LOOP = "for _result_csv in" + + +def render_slurm(tmp_path, model_overrides=None, nodes=2): + """Render job.sh.j2 the way prepare() does, with the model entry adjusted.""" + from tests.unit.test_slurm_job_template import MODEL_ENTRY, _build_deployment + + deployment = _build_deployment(tmp_path, {"nodes": nodes}, {"nnodes": nodes}) + model_entry = dict(MODEL_ENTRY) + model_entry.update(model_overrides or {}) + context = deployment._prepare_template_context(model_entry) + return deployment.jinja_env.get_template("job.sh.j2").render(**context) + + +def render_kubernetes(**context): + """Render job.yaml.j2 on its own; undeclared variables render empty, as in Jinja.""" + templates = Path(deployment_package.__file__).parent / "templates" / "kubernetes" + context.setdefault("env_vars", {}) + return create_jinja_env(templates).get_template("job.yaml.j2").render(**context) + + +class TestSlurmJobScript: + """Both copy sites in the SLURM job script: the task script and the single-node tail.""" + + def test_the_multi_node_task_script_carries_it(self, tmp_path): + script = render_slurm(tmp_path, {"multiple_results": ""}, nodes=2) + task_script = script.split("TASK_SCRIPT_EOF")[1] + assert SWEEP_LOOP in task_script + assert SHAPE_TEST in task_script + assert 'cp "$_result_csv" "$NODE_COLLECTION_DIR"/' in task_script + + def test_the_single_node_tail_carries_it(self, tmp_path): + script = render_slurm(tmp_path, {"multiple_results": ""}, nodes=1) + assert "TASK_SCRIPT_EOF" not in script + assert SWEEP_LOOP in script + assert SHAPE_TEST in script + assert 'mkdir -p "$NODE_COLLECTION_DIR"' in script + + def test_a_declared_file_is_still_copied_by_name(self, tmp_path): + script = render_slurm(tmp_path, {"multiple_results": "perf_dummy.csv"}, nodes=1) + assert '"$WORKSPACE/run_directory/perf_dummy.csv"' in script + assert SWEEP_LOOP in script + + def test_madengine_own_outputs_are_skipped_by_the_sweep(self, tmp_path): + script = render_slurm(tmp_path, {"multiple_results": ""}) # noqa: E501 + assert "perf.csv|perf_super*|perf_entry*" in script + + def test_the_sweep_looks_in_the_workspace_and_the_run_directory(self, tmp_path): + script = render_slurm(tmp_path, {"multiple_results": ""}) + assert '"$WORKSPACE"/*.csv "$WORKSPACE"/run_directory/*.csv' in script + + +class TestKubernetesJob: + """Both container blocks: the launcher arm and the direct-script arm.""" + + def test_the_launcher_arm_carries_it(self): + manifest = render_kubernetes(launcher_command="bash /tmp/run_launcher.sh") + assert SWEEP_LOOP in manifest + assert SHAPE_TEST in manifest + + def test_the_direct_script_arm_carries_it(self): + manifest = render_kubernetes() + assert SWEEP_LOOP in manifest + assert SHAPE_TEST in manifest + + def test_it_copies_into_the_results_volume(self): + manifest = render_kubernetes() + assert 'cp "$_result_csv" /results/${HOSTNAME}/' in manifest + + def test_a_declared_file_is_still_copied_by_name(self): + manifest = render_kubernetes(multiple_results="perf_dummy.csv") + assert "/workspace/perf_dummy.csv" in manifest + assert SWEEP_LOOP in manifest + + +class TestRankingIsShared: + """The SLURM collector ranks candidates with the same code the Docker path uses.""" + + def test_the_deployment_delegates_to_the_shared_ranking(self, tmp_path): + from unittest.mock import MagicMock + + from tests.unit.test_result_csv_discovery import make_result_csv + + deployment = object.__new__(SlurmDeployment) + deployment.console = MagicMock() + thin = make_result_csv(tmp_path / "node_0" / "r.csv", count=1) + rich = make_result_csv(tmp_path / "node_1" / "r.csv", count=8) + assert deployment._select_best_multiple_results_csv([thin, rich]) == rich + + def test_no_candidates_yields_nothing(self): + from unittest.mock import MagicMock + + deployment = object.__new__(SlurmDeployment) + deployment.console = MagicMock() + assert deployment._select_best_multiple_results_csv([]) is None