diff --git a/src/madengine/execution/bare_metal_runner.py b/src/madengine/execution/bare_metal_runner.py new file mode 100644 index 00000000..34de716c --- /dev/null +++ b/src/madengine/execution/bare_metal_runner.py @@ -0,0 +1,654 @@ +#!/usr/bin/env python3 +"""Bare-metal (conda) execution backend for madengine. + +Runs a model's scripts directly on the host inside a conda environment, without +Docker. The conda env (created in the build phase by :class:`CondaEnvManager`) +provides dependency isolation; this runner wraps each script invocation in +``conda run -n `` and reuses madengine's pre/post-script, performance +extraction, and ``perf.csv`` reporting. + +This is the non-Docker sibling of ``ContainerRunner``. It is single-node/local +only. It parallels ``ContainerRunner._run_self_managed`` (which already runs +scripts on the host for self-managed launchers) but adds conda wrapping, GPU +environment setup, and full reporting. + +Copyright (c) Advanced Micro Devices, Inc. All rights reserved. +""" + +import os +import shlex +import time +import typing +from contextlib import redirect_stderr, redirect_stdout + +from rich.console import Console as RichConsole + +from madengine.core.console import Console +from madengine.core.context import Context +from madengine.core.dataprovider import Data +from madengine.execution.conda_env import CondaEnvManager, resolve_conda_env_name +from madengine.execution.container_runner_helpers import ( + make_run_log_file_path, + resolve_run_timeout, +) +from madengine.execution.run_reporting import ( + determine_status, + extract_performance_from_log, + write_perf_records, +) +from madengine.reporting.update_perf_csv import PERF_CSV_HEADER, flatten_tags +from madengine.utils.config_parser import ConfigParser +from madengine.utils.gpu_config import resolve_runtime_gpus +from madengine.utils.ops import PythonicTee, file_print +from madengine.utils.path_utils import scripts_base_dir_from +from madengine.utils.run_details import get_build_number, get_pipeline + + +class BareMetalRunner: + """Run models on bare metal inside a conda environment (no Docker).""" + + def __init__( + self, + context: Context = None, + data: Data = None, + console: Console = None, + live_output: bool = False, + additional_context: typing.Dict = None, + ): + """Initialize the bare-metal runner. + + Args: + context: The madengine context (with runtime GPU detection). + data: The data provider instance. + console: Optional console instance. + live_output: Whether to stream output live. + additional_context: Additional configuration context. + """ + self.context = context + self.data = data + self.console = console or Console(live_output=live_output) + self.live_output = live_output + self.rich_console = RichConsole() + self.credentials = None + self.perf_csv_path = "perf.csv" + self.additional_context = additional_context or {} + self.bm_config = self.additional_context.get("bare_metal", {}) or {} + gpu_vendor = self.context.ctx.get("gpu_vendor") if self.context else None + gpu_arch = ( + (self.context.ctx.get("docker_env_vars") or {}).get( + "MAD_SYSTEM_GPU_ARCHITECTURE" + ) + if self.context + else None + ) + self.conda = CondaEnvManager( + console=self.console, + bm_config=self.bm_config, + gpu_vendor=gpu_vendor, + gpu_arch=gpu_arch, + ) + + def set_perf_csv_path(self, path: str) -> None: + """Set the perf.csv output path.""" + self.perf_csv_path = path + + def set_credentials(self, credentials: typing.Dict) -> None: + """Set credentials for model execution.""" + self.credentials = credentials + + def ensure_perf_csv_exists(self) -> None: + """Ensure perf.csv exists with the standard header.""" + if not os.path.exists(self.perf_csv_path): + file_print(PERF_CSV_HEADER, filename=self.perf_csv_path, mode="w") + print(f"Created performance CSV file: {self.perf_csv_path}") + + def _gpu_env(self, resolved_gpu_count: int) -> typing.Dict[str, str]: + """Build GPU-visibility env vars for the resolved GPU count. + + On bare metal there is no Docker ``--device`` flag; instead we expose a + GPU subset via ``HIP_VISIBLE_DEVICES`` (AMD) / ``CUDA_VISIBLE_DEVICES`` + (NVIDIA). A count of -1 (all) leaves visibility unrestricted. + + Args: + resolved_gpu_count: Number of GPUs requested. + + Returns: + Dict of GPU env vars (possibly empty). + """ + env: typing.Dict[str, str] = {} + if resolved_gpu_count is None or resolved_gpu_count < 0: + return env + vendor = "" + if self.context: + vendor = str(self.context.ctx.get("gpu_vendor", "")).upper() + device_list = ",".join(str(i) for i in range(resolved_gpu_count)) + if "AMD" in vendor: + env["HIP_VISIBLE_DEVICES"] = device_list + env["ROCR_VISIBLE_DEVICES"] = device_list + elif "NVIDIA" in vendor: + env["CUDA_VISIBLE_DEVICES"] = device_list + return env + + def _build_run_env( + self, model_info: typing.Dict, resolved_gpu_count: int + ) -> typing.Dict[str, str]: + """Assemble the environment for the model script. + + Layers (later overrides earlier): host env, context docker_env_vars, + GPU-visibility vars, model card env_vars, additional_context env_vars, + MAD_MODEL_NAME / build number. + + Args: + model_info: Model definition dict. + resolved_gpu_count: Number of GPUs requested. + + Returns: + Environment dict for subprocess execution. + """ + env = os.environ.copy() + + if self.context and "docker_env_vars" in self.context.ctx: + for key, value in self.context.ctx["docker_env_vars"].items(): + env[key] = str(value) + + env.update(self._gpu_env(resolved_gpu_count)) + env["MAD_RUNTIME_NGPUS"] = str( + resolved_gpu_count if resolved_gpu_count is not None else "" + ) + + if model_info.get("env_vars"): + for key, value in model_info["env_vars"].items(): + env[key] = str(value) + print(f" ENV: {key}=") + + if self.additional_context.get("env_vars"): + for key, value in self.additional_context["env_vars"].items(): + env[key] = str(value) + print(f" ENV: {key}=") + + env["MAD_MODEL_NAME"] = model_info["name"] + env["JENKINS_BUILD_NUMBER"] = get_build_number() + multiple_results = model_info.get("multiple_results") + if multiple_results: + env["MAD_OUTPUT_CSV"] = multiple_results + + return env + + def _resolve_script(self, model_info: typing.Dict) -> typing.Tuple[str, str, str]: + """Resolve script path, working directory, and interpreter. + + Mirrors ContainerRunner._run_self_managed: a ``.sh``/``.py`` path is used + directly; a directory falls back to ``run.sh``. + + Args: + model_info: Model definition dict. + + Returns: + Tuple ``(script_path, working_dir, interpreter)`` where interpreter is + "python3" or "bash". + + Raises: + FileNotFoundError: If the resolved script does not exist. + """ + scripts_arg = model_info["scripts"] + if scripts_arg.endswith((".sh", ".slurm", ".py")): + script_path = scripts_arg + else: + script_path = os.path.join(scripts_arg, "run.sh") + + if not os.path.isabs(script_path): + script_path = os.path.join(os.getcwd(), script_path) + + if not os.path.exists(script_path): + raise FileNotFoundError(f"Script not found: {script_path}") + + working_dir = os.path.dirname(script_path) or os.getcwd() + interpreter = "python3" if script_path.endswith(".py") else "bash" + return script_path, working_dir, interpreter + + def _sh_in_dir( + self, + cmd: str, + env: typing.Dict[str, str], + working_dir: str, + timeout: typing.Optional[int], + can_fail: bool, + ) -> str: + """Run *cmd* in *working_dir* via Console.sh, capturing output to the log. + + Console.sh captures the child's stdout/stderr and (in live mode) streams + it; the returned text is re-printed by the caller so it is teed into the + run-log file (subprocess fds are not affected by redirect_stdout, so this + capture-then-print path is what actually lands output in the log). + + Args: + cmd: Command to run. + env: Environment dict for the child process. + working_dir: Directory to cd into before running. + timeout: Timeout in seconds (None disables). + can_fail: If False, a nonzero exit raises RuntimeError. + + Returns: + The captured command output. + """ + full_cmd = f"cd {shlex.quote(working_dir)} && {cmd}" + return self.console.sh( + full_cmd, + canFail=can_fail, + timeout=timeout if timeout and timeout > 0 else None, + env=env, + ) + + def _run_scripts( + self, + scripts: typing.List[typing.Dict], + env_name: str, + env: typing.Dict[str, str], + working_dir: str, + timeout: typing.Optional[int], + ) -> None: + """Run pre/post scripts inside the conda env. + + Each entry is ``{"path": ..., "args": ...}``. Scripts run from + *working_dir* so relative ``scripts/common`` references resolve. A + nonzero exit raises (pre/post scripts are setup steps that must succeed). + + Args: + scripts: List of script descriptors. + env_name: Conda env name. + env: Environment dict. + working_dir: Working directory for execution. + timeout: Per-script timeout in seconds. + """ + prefix = self.conda.conda_run_prefix(env_name) + for script in scripts: + script_path = script["path"].strip() + script_args = script.get("args", "").strip() if "args" in script else "" + args_q = ( + " ".join(shlex.quote(a) for a in shlex.split(script_args)) + if script_args + else "" + ) + cmd = f"{prefix} bash {shlex.quote(script_path)} {args_q}".rstrip() + print(f"๐Ÿ”ง Pre/Post script: {cmd}") + output = self._sh_in_dir(cmd, env, working_dir, timeout, can_fail=False) + if not self.live_output: + print(output) + + def _create_run_details( + self, + model_info: typing.Dict, + build_info: typing.Dict, + run_results: typing.Dict, + resolved_gpu_count: int, + ) -> typing.Dict: + """Build a perf.csv run-details dict for a bare-metal run. + + Args: + model_info: Model definition dict. + build_info: Build info from the manifest entry. + run_results: Accumulated run results (status, performance, etc.). + resolved_gpu_count: Number of GPUs used. + + Returns: + Run details dict compatible with update_perf_csv. + """ + gpu_arch = "" + if self.context: + gpu_arch = (self.context.ctx.get("docker_env_vars") or {}).get( + "MAD_SYSTEM_GPU_ARCHITECTURE", "" + ) + n_gpus = str(resolved_gpu_count if resolved_gpu_count is not None else "") + + run_details = { + "model": model_info["name"], + "n_gpus": n_gpus, + "nnodes": "1", + "gpus_per_node": n_gpus, + "training_precision": model_info.get("training_precision", ""), + "pipeline": get_pipeline(), + "args": model_info.get("args", ""), + "tags": model_info.get("tags", ""), + "docker_file": "", + "base_docker": "", + "docker_sha": "", + "docker_image": "", + "git_commit": run_results.get("git_commit", ""), + "machine_name": run_results.get("machine_name", ""), + "deployment_type": "bare_metal", + "launcher": "conda", + "gpu_architecture": gpu_arch, + "performance": run_results.get("performance", ""), + "metric": run_results.get("metric", ""), + "relative_change": "", + "status": run_results.get("status", "FAILURE"), + "build_duration": build_info.get("build_duration", ""), + "test_duration": run_results.get("test_duration", ""), + "dataname": run_results.get("dataname", ""), + "data_provider_type": run_results.get("data_provider_type", ""), + "data_size": run_results.get("data_size", ""), + "data_download_duration": run_results.get("data_download_duration", ""), + "build_number": get_build_number(), + "additional_docker_run_options": "", + } + flatten_tags(run_details) + + try: + scripts_base_dir = scripts_base_dir_from(model_info.get("scripts", "")) + config_parser = ConfigParser(scripts_base_dir=scripts_base_dir) + run_details["configs"] = config_parser.parse_and_load( + model_info.get("args", ""), model_info.get("scripts", "") + ) + except Exception as e: + print(f"โš ๏ธ Warning: Could not parse config file: {e}") + run_details["configs"] = None + + return run_details + + def run_model( + self, + model_info: typing.Dict, + build_info: typing.Dict = None, + timeout: int = 7200, + skip_model_run: bool = False, + phase_suffix: str = "", + ) -> typing.Dict: + """Run a single model on bare metal inside its conda env. + + Args: + model_info: Model definition dict. + build_info: Build info from the manifest entry. + timeout: Execution timeout in seconds. + skip_model_run: If True, run pre-scripts but skip the model script. + phase_suffix: Suffix for the run-log file (e.g. ".run"). + + Returns: + Run results dict (status, performance, metric, test_duration, ...). + """ + build_info = build_info or {} + env_name = resolve_conda_env_name(model_info, self.bm_config) + timeout = resolve_run_timeout(model_info, timeout) + # Bare-metal logs are named like the Docker path but with an explicit + # bare-metal marker instead of an image reference. + log_file_path = make_run_log_file_path( + model_info, f"bare_metal_{env_name}", phase_suffix + ) + print(f"Run log will be written to: {log_file_path}") + + machine_name = self.console.sh("hostname") + + run_results = { + "model": model_info["name"], + "docker_image": "", + "status": "FAILURE", + "performance": "", + "metric": "", + "test_duration": 0, + "machine_name": machine_name, + "log_file": log_file_path, + } + + self.rich_console.print( + f"[bold green]๐Ÿƒ Running model:[/bold green] " + f"[bold cyan]{model_info['name']}[/bold cyan] " + f"[dim]on bare metal in conda env[/dim] [yellow]{env_name}[/yellow]" + ) + + resolved_gpu_count = resolve_runtime_gpus(model_info, self.additional_context) + env = self._build_run_env(model_info, resolved_gpu_count) + + # Collect pre/post scripts from context (populated from scripts/common). + pre_scripts = ( + list(self.context.ctx.get("pre_scripts", [])) if self.context else [] + ) + post_scripts = ( + list(self.context.ctx.get("post_scripts", [])) if self.context else [] + ) + encapsulate = ( + self.context.ctx.get("encapsulate_script", "") if self.context else "" + ) + + try: + script_path, working_dir, interpreter = self._resolve_script(model_info) + except FileNotFoundError as e: + run_results["status"] = "FAILURE" + run_results["status_detail"] = str(e) + self.rich_console.print(f"[red]โœ— {e}[/red]") + self._record(model_info, build_info, run_results, resolved_gpu_count) + return run_results + + model_args = ( + self.context.ctx.get("model_args", model_info.get("args", "")) + if self.context + else model_info.get("args", "") + ) + args_q = ( + " ".join(shlex.quote(a) for a in shlex.split(model_args)) + if model_args + else "" + ) + prefix = self.conda.conda_run_prefix(env_name) + encap = f"{encapsulate} " if encapsulate else "" + model_cmd = f"{prefix} {encap}{interpreter} {shlex.quote(script_path)} {args_q}".rstrip() + + test_start_time = time.time() + try: + with open(log_file_path, mode="w", buffering=1) as outlog: + with redirect_stdout( + PythonicTee(outlog, self.live_output) + ), redirect_stderr(PythonicTee(outlog, self.live_output)): + print(f"โฐ Setting timeout to {timeout} seconds.") + print(f"๐Ÿ“‚ Working directory: {working_dir}") + print(f"๐Ÿ Conda env: {env_name}") + + if pre_scripts: + self._run_scripts( + pre_scripts, env_name, env, working_dir, timeout + ) + + if skip_model_run: + run_results["status"] = "SKIPPED" + print("Skipping model run (--skip-model-run).") + else: + print(f"๐Ÿš€ Executing: {model_cmd}") + print("=" * 80) + # canFail=True: a nonzero model exit is not fatal here; + # status is decided from perf metrics + log error scan + # (matches the Docker path's status semantics). + model_output = self._sh_in_dir( + model_cmd, env, working_dir, timeout, can_fail=True + ) + if not self.live_output: + print(model_output) + print("=" * 80) + + if post_scripts: + self._run_scripts( + post_scripts, env_name, env, working_dir, timeout + ) + + run_results["test_duration"] = time.time() - test_start_time + print(f"test_duration: {run_results['test_duration']:.2f}s") + + if not skip_model_run: + performance, metric = extract_performance_from_log(log_file_path) + run_results["performance"] = performance or "" + run_results["metric"] = metric or "" + run_results["status"] = determine_status( + log_file_path, performance, model_info, self.additional_context + ) + if run_results["status"] == "SUCCESS": + self.rich_console.print("[green]Status: SUCCESS[/green]") + else: + self.rich_console.print("[red]Status: FAILURE[/red]") + + except Exception as e: + run_results["status"] = "FAILURE" + run_results["status_detail"] = str(e) + run_results["test_duration"] = time.time() - test_start_time + self.rich_console.print(f"[red]โœ— Bare-metal run failed: {e}[/red]") + + self._record(model_info, build_info, run_results, resolved_gpu_count) + return run_results + + def _record( + self, + model_info: typing.Dict, + build_info: typing.Dict, + run_results: typing.Dict, + resolved_gpu_count: int, + ) -> None: + """Write perf records for a run (skipped for SKIPPED / deferred perf).""" + if run_results.get("status") == "SKIPPED": + return + if self.additional_context.get("skip_perf_collection", False): + return + self.ensure_perf_csv_exists() + try: + run_details = self._create_run_details( + model_info, build_info, run_results, resolved_gpu_count + ) + write_perf_records( + run_details, + model_info, + self.perf_csv_path, + run_results.get("status", "FAILURE"), + ) + print(f"Updated perf.csv with result for {model_info['name']}") + except Exception as e: + self.rich_console.print( + f"[yellow]Warning: Could not update perf.csv: {e}[/yellow]" + ) + + # Opt-in teardown: remove the conda env after perf is recorded so a + # removal failure never loses performance data. Best-effort (remove() + # uses canFail=True) โ€” leave the node clean on shared bare-metal hosts. + if self.bm_config.get("cleanup_env", False): + env_name = resolve_conda_env_name(model_info, self.bm_config) + try: + self.conda.remove(env_name) + print(f"๐Ÿงน Removed conda env '{env_name}' (cleanup_env)") + except Exception as e: + self.rich_console.print( + f"[yellow]Warning: Could not remove conda env " + f"'{env_name}': {e}[/yellow]" + ) + + def run_models_from_manifest( + self, + manifest_file: str, + registry: str = None, + timeout: int = 7200, + keep_alive: bool = False, + keep_model_dir: bool = False, + skip_model_run: bool = False, + phase_suffix: str = "", + ) -> typing.Dict: + """Run all models from a build manifest on bare metal. + + Signature mirrors ContainerRunner.run_models_from_manifest so the + orchestrator can call either interchangeably. ``registry``, + ``keep_alive`` and ``keep_model_dir`` are Docker-only and ignored here. + + Args: + manifest_file: Path to build_manifest.json. + registry: Ignored (Docker-only). + timeout: Execution timeout per model in seconds. + keep_alive: Ignored (Docker-only). + keep_model_dir: Ignored (Docker-only). + skip_model_run: Whether to skip the model script invocation. + phase_suffix: Suffix for log files (e.g. ".run"). + + Returns: + Execution summary: successful_runs, failed_runs, total_runs. + """ + import json + + self.rich_console.print( + f"[bold blue]๐Ÿ“ฆ Loading manifest:[/bold blue] {manifest_file}" + ) + with open(manifest_file, "r") as f: + manifest = json.load(f) + + built_images = manifest.get("built_images", {}) + built_models = manifest.get("built_models", {}) + + if "context" in manifest and isinstance(manifest["context"], dict): + self.additional_context = { + **(self.additional_context or {}), + **manifest["context"], + } + self.bm_config = self.additional_context.get("bare_metal", {}) or {} + self.conda.bm_config = self.bm_config + + if not built_models: + self.rich_console.print("[yellow]โš ๏ธ No models found in manifest[/yellow]") + return {"successful_runs": [], "failed_runs": [], "total_runs": 0} + + keys = built_images.keys() if built_images else built_models.keys() + + successful_runs: typing.List[typing.Dict] = [] + failed_runs: typing.List[typing.Dict] = [] + + for key in keys: + model_info = built_models.get(key, {}) + if not model_info: + self.rich_console.print( + f"[yellow]โš ๏ธ No model info for {key}, skipping[/yellow]" + ) + continue + build_info = built_images.get(key, {}) if built_images else {} + try: + run_results = self.run_model( + model_info=model_info, + build_info=build_info, + timeout=timeout, + skip_model_run=skip_model_run, + phase_suffix=phase_suffix, + ) + status = run_results.get("status", "FAILURE") + if status in ("SUCCESS", "SKIPPED"): + successful_runs.append( + { + "model": model_info["name"], + "image": "bare_metal", + "status": status, + "performance": run_results.get("performance"), + "duration": run_results.get("test_duration"), + } + ) + else: + failed_runs.append( + { + "model": model_info["name"], + "image": "bare_metal", + "status": status, + "error": "Bare-metal execution failed - check logs", + } + ) + self.rich_console.print( + f"[red]โŒ Run failed for {model_info['name']}: {status}[/red]" + ) + except Exception as e: + self.rich_console.print( + f"[red]โŒ Failed to run {model_info.get('name', key)}: {e}[/red]" + ) + failed_runs.append( + { + "model": model_info.get("name", key), + "image": "bare_metal", + "error": str(e), + } + ) + + self.rich_console.print("\n[bold]๐Ÿ“Š Execution Summary:[/bold]") + self.rich_console.print( + f" [green]โœ“ Successful:[/green] {len(successful_runs)}" + ) + self.rich_console.print(f" [red]โœ— Failed:[/red] {len(failed_runs)}") + + return { + "successful_runs": successful_runs, + "failed_runs": failed_runs, + "total_runs": len(successful_runs) + len(failed_runs), + } diff --git a/src/madengine/execution/conda_env.py b/src/madengine/execution/conda_env.py new file mode 100644 index 00000000..ea120ca2 --- /dev/null +++ b/src/madengine/execution/conda_env.py @@ -0,0 +1,554 @@ +#!/usr/bin/env python3 +"""Conda environment lifecycle management for bare-metal execution. + +This module manages conda/mamba environments used by the bare-metal execution +backend, which runs models directly on the host (no Docker). The conda env plays +the role the Docker image plays in the container path: it isolates model +dependencies from the host and from other models. + +Copyright (c) Advanced Micro Devices, Inc. All rights reserved. +""" + +import hashlib +import os +import platform +import shlex +import shutil +import stat +import typing +import urllib.request + +from madengine.core.console import Console + +# Base URL for downloading pinned micromamba static binaries when no +# conda/mamba is present on the node. micromamba is a drop-in for the +# conda/mamba CLI surface CondaEnvManager relies on. +_MICROMAMBA_BASE_URL = "https://micro.mamba.pm/api/micromamba" + +# Base URL for TheRock's per-architecture ROCm/torch pip wheel indexes. The +# resolved index for a gfx arch is "//" and serves both the +# rocm[...] userspace packages and matching torch/torchvision wheels. Kept as a +# single constant because TheRock's preview URLs may shift. +_ROCM_NIGHTLIES_BASE_URL = "https://rocm.nightlies.amd.com/v2" + + +def resolve_conda_env_name( + model_info: typing.Dict, bm_config: typing.Optional[typing.Dict] = None +) -> str: + """Resolve the conda env name for a model. + + Priority: bare_metal config ``conda_env`` > model card ``conda_env`` > + derived from the model name (``/`` replaced with ``_``). + + Args: + model_info: Model definition dict. + bm_config: The ``bare_metal`` block from additional_context. + + Returns: + The conda environment name. + """ + bm_config = bm_config or {} + name = bm_config.get("conda_env") or model_info.get("conda_env") + if name: + return str(name) + return "mad_" + str(model_info.get("name", "model")).replace("/", "_") + + +def _vendor_suffix(vendor: typing.Optional[str]) -> str: + """Map a GPU vendor string to the file suffix used for vendor variants. + + Mirrors the Dockerfile convention (``*.amd.Dockerfile`` / + ``*.nvidia.Dockerfile``): AMD -> ``amd``, NVIDIA -> ``nvidia``. + + Args: + vendor: GPU vendor string (case-insensitive), e.g. ``"AMD"``. + + Returns: + The lowercase suffix token, or ``""`` if the vendor is unknown. + """ + v = str(vendor or "").strip().upper() + if v == "AMD": + return "amd" + if v == "NVIDIA": + return "nvidia" + return "" + + +def resolve_rocm_index_url(index_url: str, gfx_arch: str) -> str: + """Resolve the ROCm pip wheel index URL for a GPU architecture. + + An explicit *index_url* (anything other than ``"auto"``) is returned + unchanged. ``"auto"`` (or empty) resolves to TheRock's per-arch index + ``//`` using the detected gfx architecture. + + Args: + index_url: Configured index URL, or ``"auto"`` to derive from the arch. + gfx_arch: Detected GPU architecture (e.g. ``"gfx942"``). + + Returns: + The resolved pip index URL. + + Raises: + RuntimeError: If ``"auto"`` is requested but no gfx arch is available. + """ + if index_url and index_url != "auto": + return index_url + if not gfx_arch: + raise RuntimeError( + "rocm.index_url=auto requires a detected GPU architecture " + "(MAD_SYSTEM_GPU_ARCHITECTURE); set rocm.index_url explicitly instead." + ) + return f"{_ROCM_NIGHTLIES_BASE_URL}/{gfx_arch}/" + + +def resolve_environment_file(base_path: str, vendor: typing.Optional[str]) -> str: + """Resolve a vendor-specific variant of a file, else the base file. + + Given ``scripts/dummy/environment.yml`` and ``vendor="AMD"``, prefer + ``scripts/dummy/environment.amd.yml`` when it exists, otherwise return the + base path unchanged. Mirrors the Dockerfile suffix convention used by + ``DockerBuilder._get_dockerfiles_for_model`` (suffix match only, no + ``# CONTEXT`` build-arg filtering โ€” there is no build-arg matrix here). + + Args: + base_path: The base file path (e.g. an ``environment.yml``). + vendor: GPU vendor string used to pick the variant. + + Returns: + The vendor-specific path if it exists, else *base_path*. + """ + if not base_path: + return base_path + suffix = _vendor_suffix(vendor) + if not suffix: + return base_path + root, ext = os.path.splitext(base_path) + candidate = f"{root}.{suffix}{ext}" + if os.path.isfile(candidate): + return candidate + return base_path + + +def bootstrap_micromamba() -> str: + """Download a pinned micromamba static binary when no conda/mamba exists. + + Fetches the micromamba binary for the host architecture into + ``~/.cache/madengine/micromamba/micromamba`` and returns that path. + Idempotent: if the cached binary already exists, the download is skipped. + + micromamba is a drop-in for the ``conda``/``mamba`` CLI surface used by + :class:`CondaEnvManager` (``env list``, ``create``, ``env update``, + ``run -n --no-capture-output``, ``env remove``), so no other code + needs to branch on which binary is in use. + + Returns: + Path to the executable micromamba binary. + + Raises: + RuntimeError: If the host architecture is unsupported or the download + fails. + """ + cache_dir = os.path.join( + os.path.expanduser("~"), ".cache", "madengine", "micromamba" + ) + target = os.path.join(cache_dir, "micromamba") + if os.path.isfile(target) and os.access(target, os.X_OK): + return target + + machine = platform.machine().lower() + if machine in ("x86_64", "amd64"): + platform_tag = "linux-64" + elif machine in ("aarch64", "arm64"): + platform_tag = "linux-aarch64" + else: + raise RuntimeError( + f"Unsupported architecture for micromamba bootstrap: {machine}" + ) + + url = f"{_MICROMAMBA_BASE_URL}/{platform_tag}/latest" + os.makedirs(cache_dir, exist_ok=True) + print(f"โค“ Bootstrapping micromamba for {platform_tag} from {url}") + try: + # The endpoint serves a tar.bz2 archive with the binary at + # bin/micromamba; stream it and extract just that member. + import tarfile + + tmp_archive = target + ".tar.bz2" + urllib.request.urlretrieve(url, tmp_archive) # noqa: S310 (trusted URL) + with tarfile.open(tmp_archive, "r:bz2") as tar: + member = tar.getmember("bin/micromamba") + member.name = "micromamba" + tar.extract(member, path=cache_dir) + os.remove(tmp_archive) + except Exception as exc: + raise RuntimeError(f"Failed to bootstrap micromamba from {url}: {exc}") from exc + + mode = os.stat(target).st_mode + os.chmod(target, mode | stat.S_IXUSR | stat.S_IXGRP | stat.S_IXOTH) + print(f"โœ“ micromamba ready: {target}") + return target + + +class CondaEnvManager: + """Manage conda/mamba environments for bare-metal execution.""" + + def __init__( + self, + console: typing.Optional[Console] = None, + bm_config: typing.Optional[typing.Dict] = None, + gpu_vendor: typing.Optional[str] = None, + gpu_arch: typing.Optional[str] = None, + ) -> None: + """Initialize the conda env manager. + + Args: + console: Console for shell execution (created if not provided). + bm_config: The ``bare_metal`` block from additional_context. + gpu_vendor: Detected GPU vendor (e.g. ``"AMD"``/``"NVIDIA"``) used to + pick vendor-specific dependency files. Falls back to + ``bare_metal.gpu_vendor`` when not given. + gpu_arch: Detected GPU architecture (e.g. ``"gfx942"``) used to + resolve ``rocm.index_url=auto``. Falls back to + ``bare_metal.gpu_arch`` when not given. + """ + self.console = console or Console() + self.bm_config = bm_config or {} + self.gpu_vendor = gpu_vendor or self.bm_config.get("gpu_vendor") + self.gpu_arch = gpu_arch or self.bm_config.get("gpu_arch") + self._conda_bin: typing.Optional[str] = None + + def detect_conda_bin(self) -> str: + """Locate the conda (or mamba/micromamba) executable. + + Respects ``bare_metal.conda_bin`` when set, else searches PATH for + ``micromamba``, ``mamba``, then ``conda``. When none is found, bootstraps + a pinned micromamba binary (see :func:`bootstrap_micromamba`). Caches the + result. + + Returns: + Path to the conda/mamba/micromamba executable. + + Raises: + RuntimeError: If a configured ``conda_bin`` is invalid, or micromamba + bootstrap fails. + """ + if self._conda_bin: + return self._conda_bin + + configured = self.bm_config.get("conda_bin") + if configured: + if os.path.isfile(configured) or shutil.which(configured): + self._conda_bin = configured + return self._conda_bin + raise RuntimeError( + f"Configured conda_bin '{configured}' not found or not executable." + ) + + for candidate in ("micromamba", "mamba", "conda"): + found = shutil.which(candidate) + if found: + self._conda_bin = found + return self._conda_bin + + # Nothing on PATH and no conda_bin configured: bootstrap micromamba so + # bare-metal runs work with zero manual conda install. + self._conda_bin = bootstrap_micromamba() + return self._conda_bin + + def env_exists(self, env_name: str) -> bool: + """Return True if a conda env with *env_name* already exists. + + Matches by exact env name in ``conda env list`` output (the last path + component of each env directory, and any explicitly named env). + + Args: + env_name: The conda environment name. + + Returns: + True if the environment exists. + """ + conda_bin = self.detect_conda_bin() + try: + output = self.console.sh(f"{shlex.quote(conda_bin)} env list", canFail=True) + except Exception: + return False + + for line in output.splitlines(): + line = line.strip() + if not line or line.startswith("#"): + continue + # Format: " [*] " OR just a path. + parts = line.split() + if not parts: + continue + first = parts[0] + if first == env_name: + return True + # Match on the basename of the env path (last column). + last = parts[-1] + if os.path.basename(last.rstrip("/")) == env_name: + return True + return False + + def conda_run_prefix(self, env_name: str) -> str: + """Return the ``conda run`` command prefix for executing inside *env_name*. + + Uses ``--no-capture-output`` so the child process streams stdout/stderr + directly (matching live-output behavior of the Docker path). + + Args: + env_name: The conda environment name. + + Returns: + Command prefix string, e.g. ``/path/conda run -n env --no-capture-output``. + """ + conda_bin = self.detect_conda_bin() + # `--no-capture-output` is a conda/mamba flag that makes `run` stream the + # child's stdout/stderr instead of buffering. micromamba does not accept + # it (it streams by default) and forwards the unrecognized token into its + # internal `exec`, producing "exec: --: invalid option". So only add the + # flag for conda/mamba. + is_micromamba = "micromamba" in os.path.basename(conda_bin).lower() + capture_flag = "" if is_micromamba else " --no-capture-output" + return f"{shlex.quote(conda_bin)} run -n {shlex.quote(env_name)}{capture_flag}" + + def _dep_hash_path(self, env_name: str) -> str: + """Return the path of the dependency-hash stamp file for *env_name*.""" + cache_dir = os.path.join( + os.path.expanduser("~"), ".cache", "madengine", "envhash" + ) + return os.path.join(cache_dir, env_name) + + def _compute_dep_hash(self, files: typing.List[str]) -> str: + """Hash the content of file-driven dependency inputs. + + Only existing files contribute. Returns ``""`` when no such files are + given, which callers treat as "no file-driven deps" (pure reuse). + + Args: + files: Candidate dependency file paths (environment_file, + requirements_file); missing/empty entries are skipped. + + Returns: + A hex digest, or ``""`` if no existing files were provided. + """ + digest = hashlib.sha256() + found_any = False + for path in files: + if not path or not os.path.isfile(path): + continue + found_any = True + digest.update(path.encode("utf-8")) + with open(path, "rb") as f: + digest.update(f.read()) + return digest.hexdigest() if found_any else "" + + def _write_dep_hash(self, env_name: str, dep_hash: str) -> None: + """Persist the dependency hash stamp for *env_name* (best-effort).""" + if not dep_hash: + return + path = self._dep_hash_path(env_name) + try: + os.makedirs(os.path.dirname(path), exist_ok=True) + with open(path, "w") as f: + f.write(dep_hash) + except OSError as exc: + print(f"โš ๏ธ Could not write dep-hash stamp for '{env_name}': {exc}") + + def _read_dep_hash(self, env_name: str) -> str: + """Read the stored dependency hash stamp for *env_name* (``""`` if none).""" + path = self._dep_hash_path(env_name) + try: + with open(path, "r") as f: + return f.read().strip() + except OSError: + return "" + + def _install_rocm_wheels( + self, + env_name: str, + rocm_config: typing.Dict, + timeout: typing.Optional[int], + ) -> None: + """Install ROCm userspace (and optionally torch) pip wheels into the env. + + Resolves the per-arch index URL (``index_url=auto`` -> gfx-arch index) + and pip-installs the configured packages, then torch/torchvision from the + same index when ``torch`` is set. + + Args: + env_name: The conda environment name. + rocm_config: The resolved ``rocm`` block. + timeout: Per-command timeout in seconds. + + Raises: + RuntimeError: If the index URL cannot be resolved. + """ + index_url = resolve_rocm_index_url( + rocm_config.get("index_url", "auto"), self.gpu_arch + ) + packages = rocm_config.get("packages") or ["rocm[libraries,devel]"] + prefix = self.conda_run_prefix(env_name) + url_q = shlex.quote(index_url) + pkgs_q = " ".join(shlex.quote(p) for p in packages) + self.console.sh( + f"{prefix} pip install --index-url {url_q} {pkgs_q}", + timeout=timeout, + ) + print(f"โœ“ Installed ROCm wheels in '{env_name}' from {index_url}") + if rocm_config.get("torch", False): + self.console.sh( + f"{prefix} pip install --index-url {url_q} torch torchvision", + timeout=timeout, + ) + print(f"โœ“ Installed torch/torchvision in '{env_name}' from {index_url}") + + def create_or_update( + self, model_info: typing.Dict, timeout: typing.Optional[int] = 3600 + ) -> str: + """Create or update the conda env for a model and install dependencies. + + Pipeline (each step optional and additive): + 1. Create/reuse the env: ``conda env create/update -f`` when + ``environment_file`` is given, else + ``conda create -n python=``. + 2. If ``rocm.enabled`` -> pip-install ROCm userspace (+ torch) from the + gfx-arch index. + 3. If ``requirements_file`` -> ``pip install -r`` inside the env. + 4. If ``setup_script`` -> run it inside the env. + + Idempotency: + - When the env exists and ``reuse_env`` (default True) is set, steps 1 + and 3 are skipped *unless* the content of ``environment_file`` / + ``requirements_file`` changed since last run (tracked via a content- + hash stamp), which forces an env update + requirements reinstall. + - The ROCm-wheel install (step 2) is skip-on-reuse regardless of file + hashes (it is not file-driven and reinstalling is slow); force a + refresh with ``reuse_env: False``. + - ``setup_script`` (step 4) always runs. + + Args: + model_info: Model definition dict (may carry conda_env, + environment_file, requirements_file, python_version, rocm, + setup_script). + timeout: Per-command timeout in seconds (None disables it). + + Returns: + The resolved conda environment name. + + Raises: + RuntimeError: If env creation or dependency setup fails. + """ + env_name = resolve_conda_env_name(model_info, self.bm_config) + conda_bin = self.detect_conda_bin() + conda_q = shlex.quote(conda_bin) + + environment_file = ( + self.bm_config.get("environment_file") + or model_info.get("environment_file") + or "" + ) + if environment_file: + environment_file = resolve_environment_file( + environment_file, self.gpu_vendor + ) + requirements_file = ( + self.bm_config.get("requirements_file") + or model_info.get("requirements_file") + or "" + ) + if requirements_file: + requirements_file = resolve_environment_file( + requirements_file, self.gpu_vendor + ) + python_version = ( + self.bm_config.get("python_version") + or model_info.get("python_version") + or "" + ) + setup_script = ( + self.bm_config.get("setup_script") or model_info.get("setup_script") or "" + ) + if setup_script: + setup_script = resolve_environment_file(setup_script, self.gpu_vendor) + rocm_config = self.bm_config.get("rocm") or model_info.get("rocm") or {} + reuse_env = self.bm_config.get("reuse_env", True) + + already = self.env_exists(env_name) + + # File-driven deps (environment_file/requirements_file) invalidate reuse + # when their content changes; ROCm wheels do not (see docstring). + dep_hash = self._compute_dep_hash([environment_file, requirements_file]) + deps_changed = bool(dep_hash) and dep_hash != self._read_dep_hash(env_name) + file_step_needed = (not already) or (not reuse_env) or deps_changed + + if already and reuse_env and not deps_changed: + print(f"โ„น๏ธ Conda env '{env_name}' already exists; reusing (reuse_env).") + else: + if already and reuse_env and deps_changed: + print( + f"โ„น๏ธ Dependency files changed for '{env_name}'; " + f"updating env despite reuse_env." + ) + if environment_file: + if not os.path.isfile(environment_file): + raise RuntimeError( + f"environment_file not found: {environment_file}" + ) + # `env update` is idempotent and also creates when missing. + verb = "update" if already else "create" + self.console.sh( + f"{conda_q} env {verb} -n {shlex.quote(env_name)} " + f"-f {shlex.quote(environment_file)}", + timeout=timeout, + ) + else: + py = f"python={python_version}" if python_version else "python" + self.console.sh( + f"{conda_q} create -y -n {shlex.quote(env_name)} {shlex.quote(py)}", + timeout=timeout, + ) + print(f"โœ“ Conda env ready: {env_name}") + + # ROCm wheels: skip-on-reuse (install only on fresh create or reuse_env=False). + if rocm_config.get("enabled") and (not already or not reuse_env): + self._install_rocm_wheels(env_name, rocm_config, timeout) + + if requirements_file and file_step_needed: + if not os.path.isfile(requirements_file): + raise RuntimeError(f"requirements_file not found: {requirements_file}") + prefix = self.conda_run_prefix(env_name) + self.console.sh( + f"{prefix} pip install -r {shlex.quote(requirements_file)}", + timeout=timeout, + ) + print(f"โœ“ Installed requirements in '{env_name}': {requirements_file}") + + if file_step_needed: + self._write_dep_hash(env_name, dep_hash) + + if setup_script: + if not os.path.isfile(setup_script): + raise RuntimeError(f"setup_script not found: {setup_script}") + prefix = self.conda_run_prefix(env_name) + self.console.sh( + f"{prefix} bash {shlex.quote(setup_script)}", + timeout=timeout, + ) + print(f"โœ“ Ran setup_script in env '{env_name}': {setup_script}") + + return env_name + + def remove(self, env_name: str, timeout: typing.Optional[int] = 600) -> None: + """Remove a conda environment (best-effort). + + Args: + env_name: The conda environment name. + timeout: Command timeout in seconds. + """ + conda_bin = self.detect_conda_bin() + self.console.sh( + f"{shlex.quote(conda_bin)} env remove -y -n {shlex.quote(env_name)}", + canFail=True, + timeout=timeout, + ) diff --git a/src/madengine/execution/run_reporting.py b/src/madengine/execution/run_reporting.py new file mode 100644 index 00000000..ac0b965f --- /dev/null +++ b/src/madengine/execution/run_reporting.py @@ -0,0 +1,277 @@ +#!/usr/bin/env python3 +"""Shared run-reporting helpers (performance extraction, status, perf.csv). + +Pure-Python helpers used by execution backends to turn a run log into a +performance/metric pair, a SUCCESS/FAILURE status, and rows in ``perf.csv`` / +``perf_super``. These are Docker-independent so both the container and +bare-metal backends can share the same reporting semantics. + +The performance-extraction and status logic mirror the inline behavior in +``ContainerRunner.run_container`` so results are consistent across backends. + +Copyright (c) Advanced Micro Devices, Inc. All rights reserved. +""" + +import json +import os +import re +import typing + +from madengine.deployment.base import PERFORMANCE_LOG_PATTERN +from madengine.execution.container_runner_helpers import ( + log_text_has_error_pattern, + resolve_log_error_scan_config, +) +from madengine.reporting.update_perf_csv import update_perf_csv +from madengine.reporting.update_perf_super import ( + update_perf_super_csv, + update_perf_super_json, +) +from madengine.utils.path_utils import scripts_base_dir_from + +# HuggingFace Trainer emits "'train_samples_per_second': 4.23"; used as a +# fallback when the canonical "performance: " line is absent. +_HF_PERF_PATTERN = r"train_samples_per_second[\'\"\s:=]+([0-9][0-9.eE+-]*)" + +# Benign log substrings that must not be treated as errors (mirrors +# ContainerRunner.run_container). +_DEFAULT_BENIGN_SUBSTRINGS: typing.Tuple[str, ...] = ( + "Failed to establish connection to the metrics exporter agent", + "RpcError: Running out of retries to initialize the metrics agent", + "Metrics will not be exported", + "FutureWarning", + "Opened result file:", + "SQLite3 generation ::", + "rocpd_op:", + "rpd_tracer:", +) + +_DEFAULT_BENIGN_REGEXES: typing.Tuple[str, ...] = ( + r"^E[0-9]{8}.*generateRocpd\.cpp", + r"^W[0-9]{8}.*simple_timer\.cpp", + r"^W[0-9]{8}.*generateRocpd\.cpp", + r"^E[0-9]{8}.*tool\.cpp", + r"\[rocprofv3\]", +) + + +def extract_performance_from_log( + log_file_path: str, +) -> typing.Tuple[typing.Optional[str], typing.Optional[str]]: + """Extract ``(performance, metric)`` from a run log file. + + Tries the canonical ``performance: [][,] `` pattern + first, then the HuggingFace ``train_samples_per_second`` fallback. + + Args: + log_file_path: Path to the run log file. + + Returns: + A ``(performance, metric)`` tuple; either element may be None. + """ + if not log_file_path or not os.path.exists(log_file_path): + return None, None + + with open(log_file_path, "r", encoding="utf-8", errors="ignore") as f: + log_content = f.read() + + match = re.search(PERFORMANCE_LOG_PATTERN, log_content) + if match: + return match.group(1).strip(), match.group(2).strip() + + hf_match = re.search(_HF_PERF_PATTERN, log_content) + if hf_match: + return hf_match.group(1).strip(), "samples_per_second" + + return None, None + + +def log_has_errors( + log_file_path: str, + model_info: typing.Dict, + additional_context: typing.Optional[typing.Dict] = None, +) -> bool: + """Scan a run log for error patterns (respecting benign exclusions). + + Args: + log_file_path: Path to the run log file. + model_info: Model definition dict (may override scan config). + additional_context: Additional context (may override scan config). + + Returns: + True if an error pattern was found and scanning is enabled. + """ + scan_enabled, error_patterns, extra_benign = resolve_log_error_scan_config( + model_info, additional_context + ) + if not scan_enabled or not log_file_path or not os.path.exists(log_file_path): + return False + + benign_substrings = list(_DEFAULT_BENIGN_SUBSTRINGS) + list(extra_benign) + + with open(log_file_path, "r", encoding="utf-8", errors="ignore") as f: + log_scan_text = f.read() + + for pattern in error_patterns: + if log_text_has_error_pattern( + log_scan_text, + pattern, + benign_substrings, + _DEFAULT_BENIGN_REGEXES, + ): + print(f"Found error pattern '{pattern}' in logs") + return True + return False + + +def determine_status( + log_file_path: str, + performance: typing.Optional[str], + model_info: typing.Dict, + additional_context: typing.Optional[typing.Dict] = None, +) -> str: + """Determine SUCCESS/FAILURE from log errors and presence of performance. + + Mirrors ContainerRunner.run_container status logic: a run is SUCCESS when it + has a performance metric and no error patterns; worker nodes + (``MAD_COLLECT_METRICS=false``) and deferred-collection runs + (``skip_perf_collection``) succeed without a local metric. + + Args: + log_file_path: Path to the run log file. + performance: Extracted performance value (or None). + model_info: Model definition dict. + additional_context: Additional context dict. + + Returns: + "SUCCESS" or "FAILURE". + """ + additional_context = additional_context or {} + + has_errors = log_has_errors(log_file_path, model_info, additional_context) + has_performance = bool( + performance and str(performance).strip() and str(performance).strip() != "N/A" + ) + is_worker_node = os.environ.get("MAD_COLLECT_METRICS", "true").lower() == "false" + + if has_errors: + return "FAILURE" + if has_performance: + return "SUCCESS" + if is_worker_node: + return "SUCCESS" + if additional_context.get("skip_perf_collection", False): + return "SUCCESS" + return "FAILURE" + + +def resolve_multiple_results_path( + multiple_results: str, model_dir: str +) -> typing.Optional[str]: + """Resolve a multiple_results CSV path: try cwd then model_dir. + + Args: + multiple_results: Relative CSV filename from the model card. + model_dir: Directory the model ran in. + + Returns: + First existing path, or None. + """ + if not multiple_results: + return None + if os.path.isfile(multiple_results): + return multiple_results + candidate = os.path.join(model_dir, multiple_results) + if os.path.isfile(candidate): + return candidate + return None + + +def write_perf_records( + run_details: typing.Dict, + model_info: typing.Dict, + perf_csv_path: str, + status: str, + model_dir: str = ".", +) -> None: + """Write perf.csv and perf_super records for a completed run. + + Handles the multiple_results (CSV) case and the single-result case, matching + ContainerRunner.run_container's reporting. Failures write an exception row. + + Args: + run_details: Run details dict (from a create_run_details_dict-style call). + model_info: Model definition dict. + perf_csv_path: Path to perf.csv. + status: Run status ("SUCCESS"/"FAILURE"). + model_dir: Directory the model ran in (for multiple_results resolution). + """ + scripts_path = model_info.get("scripts", "") + scripts_base_dir = scripts_base_dir_from(scripts_path) + + multiple_results = (model_info.get("multiple_results") or "").strip() + resolved_multiple = ( + resolve_multiple_results_path(multiple_results, model_dir) + if multiple_results + else None + ) + + if resolved_multiple and status == "SUCCESS": + common_info = run_details.copy() + for key in ("model", "performance", "metric", "status"): + common_info.pop(key, None) + with open("common_info.json", "w") as f: + json.dump(common_info, f) + + update_perf_csv( + multiple_results=resolved_multiple, + perf_csv=perf_csv_path, + model_name=run_details["model"], + common_info="common_info.json", + ) + try: + num_entries = update_perf_super_json( + multiple_results=resolved_multiple, + perf_super_json="perf_super.json", + model_name=run_details["model"], + common_info="common_info.json", + 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, + ) + except Exception as e: + print(f"โš ๏ธ Warning: Could not update perf_super files: {e}") + return + + # Single-result path. + with open("perf_entry.json", "w") as f: + json.dump(run_details, f) + + if status == "SUCCESS": + update_perf_csv(single_result="perf_entry.json", perf_csv=perf_csv_path) + else: + update_perf_csv(exception_result="perf_entry.json", perf_csv=perf_csv_path) + + try: + if status == "SUCCESS": + num_entries = update_perf_super_json( + single_result="perf_entry.json", + perf_super_json="perf_super.json", + scripts_base_dir=scripts_base_dir, + ) + else: + num_entries = update_perf_super_json( + exception_result="perf_entry.json", + perf_super_json="perf_super.json", + 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, + ) + except Exception as e: + print(f"โš ๏ธ Warning: Could not update perf_super files: {e}") diff --git a/src/madengine/orchestration/build_orchestrator.py b/src/madengine/orchestration/build_orchestrator.py index 246701cf..e124bedd 100644 --- a/src/madengine/orchestration/build_orchestrator.py +++ b/src/madengine/orchestration/build_orchestrator.py @@ -17,10 +17,10 @@ from rich.console import Console as RichConsole from rich.panel import Panel -from madengine.core.console import Console -from madengine.core.context import Context from madengine.core.additional_context_defaults import apply_build_context_defaults from madengine.core.auth import load_credentials +from madengine.core.console import Console +from madengine.core.context import Context from madengine.core.errors import ( BuildError, ConfigurationError, @@ -28,11 +28,11 @@ create_error_context, ) from madengine.deployment.common import is_self_managed_launcher -from madengine.utils.discover_models import DiscoverModels from madengine.execution.docker_builder import DockerBuilder from madengine.execution.dockerfile_utils import ( dockerfile_requires_explicit_mad_arch_build_arg, ) +from madengine.utils.discover_models import DiscoverModels class BuildOrchestrator: @@ -47,7 +47,12 @@ class BuildOrchestrator: - Save deployment_config from --additional-context """ - def __init__(self, args, additional_context: Optional[Dict] = None, detect_local_gpu_arch: bool = False): + def __init__( + self, + args, + additional_context: Optional[Dict] = None, + detect_local_gpu_arch: bool = False, + ): """ Initialize build orchestrator. @@ -65,7 +70,7 @@ def __init__(self, args, additional_context: Optional[Dict] = None, detect_local # Merge additional_context from args and parameter merged_context = {} - + # Load from file first if provided if hasattr(args, "additional_context_file") and args.additional_context_file: try: @@ -73,7 +78,7 @@ def __init__(self, args, additional_context: Optional[Dict] = None, detect_local merged_context = json.load(f) except (FileNotFoundError, json.JSONDecodeError) as e: print(f"Warning: Could not load additional_context_file: {e}") - + # Then merge string additional_context (overrides file) if hasattr(args, "additional_context") and args.additional_context: try: @@ -81,6 +86,7 @@ def __init__(self, args, additional_context: Optional[Dict] = None, detect_local # Use ast.literal_eval for Python dict syntax (single quotes) # This matches what Context class expects import ast + context_from_string = ast.literal_eval(args.additional_context) merged_context.update(context_from_string) elif isinstance(args.additional_context, dict): @@ -98,31 +104,42 @@ def __init__(self, args, additional_context: Optional[Dict] = None, detect_local self.additional_context = merged_context self._original_user_slurm_keys = set(merged_context.get("slurm", {}).keys()) - + # Apply ConfigLoader to infer deploy type, validate, and apply defaults if self.additional_context: try: from madengine.deployment.config_loader import ConfigLoader + # This will: # 1. Infer deploy type from k8s/slurm presence # 2. Validate for conflicts (e.g., both k8s and slurm) # 3. Apply appropriate defaults # 4. Add 'deploy' field for internal use - self.additional_context = ConfigLoader.load_config(self.additional_context) + self.additional_context = ConfigLoader.load_config( + self.additional_context + ) except ValueError as e: # Re-raise as ConfigurationError so the CLI layer handles the exit code raise ConfigurationError(str(e)) except Exception as e: # Other errors during config loading - warn but continue - self.rich_console.print(f"[yellow]Warning: Could not apply config defaults: {e}[/yellow]") + self.rich_console.print( + f"[yellow]Warning: Could not apply config defaults: {e}[/yellow]" + ) self.rich_console.print("[bold blue]Build additional context[/bold blue]\n") - self.rich_console.print(Panel( - json.dumps(self.additional_context, indent=2) if self.additional_context else "(empty)", - title="[bold]Context[/bold] (from --additional-context / --additional-context-file)", - border_style="dim", - padding=(0, 1), - )) + self.rich_console.print( + Panel( + ( + json.dumps(self.additional_context, indent=2) + if self.additional_context + else "(empty)" + ), + title="[bold]Context[/bold] (from --additional-context / --additional-context-file)", + border_style="dim", + padding=(0, 1), + ) + ) self.rich_console.print() # Initialize context in build-only mode (no GPU detection by default). @@ -143,7 +160,7 @@ def __init__(self, args, additional_context: Optional[Dict] = None, detect_local def _copy_scripts(self): """[DEPRECATED] Copy common scripts to model directories. - + This method is no longer called during build phase as it's not needed. Build phase only creates Docker images - script execution happens in run phase. Scripts are copied by run_orchestrator._copy_scripts() for local execution. @@ -233,6 +250,13 @@ def execute( ], ) + # Handle bare-metal (conda) build: create conda envs, no Docker build. + # Use membership, not truthiness: an empty dict ({'bare_metal': {}}) means + # "bare-metal with defaults" and must still route here. This mirrors the + # run-phase check in RunOrchestrator._infer_deployment_target. + if "bare_metal" in self.additional_context: + return self._execute_bare_metal_build(manifest_output=manifest_output) + # Handle pre-built image mode if use_image: # If use_image is "auto", resolve from model card @@ -271,8 +295,11 @@ def execute( # env_vars.DOCKER_IMAGE_NAME, treat it as an implicit --use-image so # users do not have to repeat the image on the CLI for slurm_multi. slurm_multi_models = [ - m for m in _discovered_models - if is_self_managed_launcher((m.get("distributed") or {}).get("launcher", "")) + m + for m in _discovered_models + if is_self_managed_launcher( + (m.get("distributed") or {}).get("launcher", "") + ) ] if slurm_multi_models and not registry: card_images = { @@ -356,14 +383,18 @@ def execute( ) self._warn_if_mad_arch_unresolved_for_dockerfiles(models, builder) - resolved_arch = self.context.ctx.get("docker_build_arg", {}).get("MAD_SYSTEM_GPU_ARCHITECTURE") + resolved_arch = self.context.ctx.get("docker_build_arg", {}).get( + "MAD_SYSTEM_GPU_ARCHITECTURE" + ) if resolved_arch: self.rich_console.print( f"[green]โœ“ MAD_SYSTEM_GPU_ARCHITECTURE resolved: {resolved_arch}[/green]\n" ) # Step 3: Build Docker images - self.rich_console.print("[bold cyan]๐Ÿ—๏ธ Building Docker images...[/bold cyan]") + self.rich_console.print( + "[bold cyan]๐Ÿ—๏ธ Building Docker images...[/bold cyan]" + ) # Determine phase suffix for log files # Build phase always uses .build suffix to avoid conflicts with run logs @@ -411,8 +442,12 @@ def execute( self.rich_console.print(f" [red]โ€ข {model_name}: {error_msg}[/red]") # Step 4: ALWAYS generate manifest (even with partial failures) - self.rich_console.print("\n[bold cyan]๐Ÿ“„ Generating build manifest...[/bold cyan]") - builder.export_build_manifest(manifest_output, registry, batch_build_metadata) + self.rich_console.print( + "\n[bold cyan]๐Ÿ“„ Generating build manifest...[/bold cyan]" + ) + builder.export_build_manifest( + manifest_output, registry, batch_build_metadata + ) # Step 5: Save build summary to manifest self._save_build_summary(manifest_output, build_summary) @@ -420,7 +455,9 @@ def execute( # Step 6: Save deployment_config to manifest self._save_deployment_config(manifest_output) - self.rich_console.print(f"[green]โœ“ Build complete: {manifest_output}[/green]") + self.rich_console.print( + f"[green]โœ“ Build complete: {manifest_output}[/green]" + ) self.rich_console.print(f"[dim]{'=' * 60}[/dim]\n") # Step 7: Check if we should fail (only if ALL builds failed) @@ -472,20 +509,22 @@ def _execute_with_prebuilt_image( ) -> str: """ Generate manifest for a pre-built Docker image (skip Docker build). - + This is useful when using external images like: - lmsysorg/sglang:v0.5.2rc1-rocm700-mi30x - nvcr.io/nvidia/pytorch:24.01-py3 - + Args: use_image: Pre-built Docker image name manifest_output: Output file for build manifest - + Returns: Path to generated build_manifest.json """ self.rich_console.print(f"\n[dim]{'=' * 60}[/dim]") - self.rich_console.print("[bold blue]๐Ÿ”จ BUILD PHASE (Pre-built Image Mode)[/bold blue]") + self.rich_console.print( + "[bold blue]๐Ÿ”จ BUILD PHASE (Pre-built Image Mode)[/bold blue]" + ) self.rich_console.print(f"[cyan]Using pre-built image: {use_image}[/cyan]") self.rich_console.print(f"[dim]{'=' * 60}[/dim]\n") @@ -511,8 +550,10 @@ def _execute_with_prebuilt_image( self.rich_console.print(f"[green]โœ“ Found {len(models)} models[/green]\n") # Step 2: Generate manifest with pre-built image - self.rich_console.print("[bold cyan]๐Ÿ“„ Generating manifest for pre-built image...[/bold cyan]") - + self.rich_console.print( + "[bold cyan]๐Ÿ“„ Generating manifest for pre-built image...[/bold cyan]" + ) + manifest = { # built_images and built_models MUST share the same key set so # ContainerRunner.run_models_from_manifest() can join them via @@ -521,7 +562,7 @@ def _execute_with_prebuilt_image( # same pre-built use_image) so multi-model --use-image runs work. "built_images": {}, "built_models": {}, - "context": self.context.ctx if hasattr(self.context, 'ctx') else {}, + "context": self.context.ctx if hasattr(self.context, "ctx") else {}, "credentials_required": [], "summary": { "successful_builds": [], @@ -582,23 +623,25 @@ def _execute_with_prebuilt_image( # Save deployment config self._save_deployment_config(manifest_output) - + # Merge model's distributed and slurm config into deployment_config # This ensures launcher and slurm settings are in deployment_config even if not in additional-context if models: with open(manifest_output, "r") as f: saved_manifest = json.load(f) - + if "deployment_config" not in saved_manifest: saved_manifest["deployment_config"] = {} - + # Merge model's distributed config from the first model. # If multiple models have differing distributed configs, warn โ€” only the first wins here. # Use json.dumps for the hash key so nested dicts (e.g. sglang_disagg / vllm_disagg) # don't trigger TypeError: unhashable type: 'dict' from `tuple(sorted(items()))`. if len(models) > 1: distinct_distributed = { - json.dumps(m.get("distributed") or {}, sort_keys=True, default=str) + json.dumps( + m.get("distributed") or {}, sort_keys=True, default=str + ) for m in models } if len(distinct_distributed) > 1: @@ -610,12 +653,26 @@ def _execute_with_prebuilt_image( if model_distributed: if "distributed" not in saved_manifest["deployment_config"]: saved_manifest["deployment_config"]["distributed"] = {} - + # Copy launcher and other critical fields from model config - for key in ["launcher", "nnodes", "nproc_per_node", "backend", "port", "sglang_disagg", "vllm_disagg"]: - if key in model_distributed and key not in saved_manifest["deployment_config"]["distributed"]: - saved_manifest["deployment_config"]["distributed"][key] = model_distributed[key] - + for key in [ + "launcher", + "nnodes", + "nproc_per_node", + "backend", + "port", + "sglang_disagg", + "vllm_disagg", + ]: + if ( + key in model_distributed + and key + not in saved_manifest["deployment_config"]["distributed"] + ): + saved_manifest["deployment_config"]["distributed"][key] = ( + model_distributed[key] + ) + # Merge model's slurm config into deployment_config.slurm from the first model. # This enables run phase to auto-detect SLURM deployment without --additional-context. # Warn when multiple models have differing slurm configs (only the first wins here). @@ -634,20 +691,36 @@ def _execute_with_prebuilt_image( if model_slurm: if "slurm" not in saved_manifest["deployment_config"]: saved_manifest["deployment_config"]["slurm"] = {} - + # Copy slurm settings from model config (model card fills in # values not explicitly set by --additional-context). # Use _original_user_slurm_keys (captured before ConfigLoader # applies defaults) so model card values override defaults # but user's explicit CLI values still win. - for key in ["partition", "nodes", "gpus_per_node", "time", "exclusive", "reservation", "output_dir", "nodelist"]: - if key in model_slurm and key not in self._original_user_slurm_keys: - saved_manifest["deployment_config"]["slurm"][key] = model_slurm[key] - + for key in [ + "partition", + "nodes", + "gpus_per_node", + "time", + "exclusive", + "reservation", + "output_dir", + "nodelist", + ]: + if ( + key in model_slurm + and key not in self._original_user_slurm_keys + ): + saved_manifest["deployment_config"]["slurm"][key] = ( + model_slurm[key] + ) + with open(manifest_output, "w") as f: json.dump(saved_manifest, f, indent=2) - self.rich_console.print(f"[green]โœ“ Generated manifest: {manifest_output}[/green]") + self.rich_console.print( + f"[green]โœ“ Generated manifest: {manifest_output}[/green]" + ) self.rich_console.print(f" Pre-built image: {use_image}") self.rich_console.print(f" Models: {len(models)}") self.rich_console.print(f"[dim]{'=' * 60}[/dim]\n") @@ -668,23 +741,25 @@ def _execute_with_prebuilt_image( def _resolve_image_from_model_card(self) -> str: """ Resolve Docker image name from model card's DOCKER_IMAGE_NAME env var. - + This method discovers models and extracts the DOCKER_IMAGE_NAME from env_vars. If multiple models have different images, uses the first and prints a warning. - + Returns: Docker image name from model card - + Raises: ConfigurationError: If no DOCKER_IMAGE_NAME found in any model """ - self.rich_console.print("[bold cyan]๐Ÿ” Auto-detecting image from model card...[/bold cyan]") - + self.rich_console.print( + "[bold cyan]๐Ÿ” Auto-detecting image from model card...[/bold cyan]" + ) + # Discover models to get their env_vars discover_models = DiscoverModels(args=self.args) models = discover_models.run() - + if not models: raise ConfigurationError( "No models discovered for image auto-detection", @@ -698,17 +773,17 @@ def _resolve_image_from_model_card(self) -> str: "Verify --tags parameter is correct", ], ) - + # Collect DOCKER_IMAGE_NAME from all models images_found = {} for model in models: model_name = model.get("name", "unknown") env_vars = model.get("env_vars", {}) docker_image = env_vars.get("DOCKER_IMAGE_NAME") - + if docker_image: images_found[model_name] = docker_image - + if not images_found: model_names = [m.get("name", "unknown") for m in models] raise ConfigurationError( @@ -724,11 +799,11 @@ def _resolve_image_from_model_card(self) -> str: 'Example: "env_vars": {"DOCKER_IMAGE_NAME": "myimage:tag"}', ], ) - + # Use first model's image first_model = list(images_found.keys())[0] resolved_image = images_found[first_model] - + # Warn if multiple models have different images unique_images = set(images_found.values()) if len(unique_images) > 1: @@ -741,8 +816,10 @@ def _resolve_image_from_model_card(self) -> str: f"[yellow] Using image from '{first_model}': {resolved_image}[/yellow]\n" ) else: - self.rich_console.print(f"[green]โœ“ Auto-detected image: {resolved_image}[/green]\n") - + self.rich_console.print( + f"[green]โœ“ Auto-detected image: {resolved_image}[/green]\n" + ) + return resolved_image def _execute_build_on_compute( @@ -754,29 +831,33 @@ def _execute_build_on_compute( ) -> str: """ Execute Docker build on a SLURM compute node and push to registry. - + Build workflow: 1. Build on 1 compute node only 2. Push image to registry 3. Store registry image name in manifest 4. Run phase will pull image in parallel on all nodes - + Args: registry: Registry to push images to (REQUIRED) clean_cache: Whether to use --no-cache for Docker builds manifest_output: Output file for build manifest batch_build_metadata: Optional batch build metadata - + Returns: Path to generated build_manifest.json """ - import subprocess - import os import glob - + import os + import subprocess + self.rich_console.print(f"\n[dim]{'=' * 60}[/dim]") - self.rich_console.print("[bold blue]๐Ÿ”จ BUILD PHASE (Compute Node Mode)[/bold blue]") - self.rich_console.print("[cyan]Building on 1 compute node, pushing to registry...[/cyan]") + self.rich_console.print( + "[bold blue]๐Ÿ”จ BUILD PHASE (Compute Node Mode)[/bold blue]" + ) + self.rich_console.print( + "[cyan]Building on 1 compute node, pushing to registry...[/cyan]" + ) self.rich_console.print(f"[dim]{'=' * 60}[/dim]\n") # registry is required for the build-on-compute flow (it must be pushed somewhere @@ -819,12 +900,12 @@ def _execute_build_on_compute( additional_info={"registry": r}, ), suggestions=[ - 'Dockerhub: --registry docker.io/(/)', - 'GHCR: --registry ghcr.io/(/)', - 'Quay: --registry quay.io/(/)', - 'NGC: --registry nvcr.io/(/)', - 'Self-hosted: --registry (:)(/)', - 'Local: --registry localhost:5000(/)', + "Dockerhub: --registry docker.io/(/)", + "GHCR: --registry ghcr.io/(/)", + "Quay: --registry quay.io/(/)", + "NGC: --registry nvcr.io/(/)", + "Self-hosted: --registry (:)(/)", + "Local: --registry localhost:5000(/)", ], ) normalized = registry.strip().rstrip("/") @@ -859,7 +940,7 @@ def _execute_build_on_compute( self.rich_console.print("[bold cyan]๐Ÿ” Discovering models...[/bold cyan]") discover_models = DiscoverModels(args=self.args) models = discover_models.run() - + if not models: raise DiscoveryError( "No models discovered for build-on-compute", @@ -872,7 +953,7 @@ def _execute_build_on_compute( "Verify --tags parameter is correct", ], ) - + # SLURM config is derived from the first model card + --additional-context overrides. # All models are built in a single sbatch job, so one SLURM config applies to all. first_model = models[0] @@ -881,11 +962,17 @@ def _execute_build_on_compute( slurm_config = {**model_slurm_config, **context_slurm_config} self.rich_console.print(f"[green]โœ“ Found {len(models)} model(s)[/green]\n") - self.rich_console.print("[bold cyan]๐Ÿ“‹ SLURM Configuration (merged):[/bold cyan]") + self.rich_console.print( + "[bold cyan]๐Ÿ“‹ SLURM Configuration (merged):[/bold cyan]" + ) if model_slurm_config: - self.rich_console.print(f" [dim]From model card:[/dim] {list(model_slurm_config.keys())}") + self.rich_console.print( + f" [dim]From model card:[/dim] {list(model_slurm_config.keys())}" + ) if context_slurm_config: - self.rich_console.print(f" [dim]From --additional-context (overrides):[/dim] {list(context_slurm_config.keys())}") + self.rich_console.print( + f" [dim]From --additional-context (overrides):[/dim] {list(context_slurm_config.keys())}" + ) # Validate required fields partition = slurm_config.get("partition") @@ -904,21 +991,21 @@ def _execute_build_on_compute( reservation = slurm_config.get("reservation", "") time_limit = slurm_config.get("time", "02:00:00") - + self.rich_console.print(f" Partition: {partition}") self.rich_console.print(f" Time limit: {time_limit}") if reservation: self.rich_console.print(f" Reservation: {reservation}") self.rich_console.print("") - + # Validate registry credentials self.rich_console.print("[bold cyan]๐Ÿ” Registry Configuration:[/bold cyan]") self.rich_console.print(f" Registry: {registry}") - + # Check for credentials - either from environment or credential.json dockerhub_user = os.environ.get("MAD_DOCKERHUB_USER", "") dockerhub_password = os.environ.get("MAD_DOCKERHUB_PASSWORD", "") - + # Try to load from credential.json if env vars not set credential_file = Path("credential.json") if not dockerhub_user and credential_file.exists(): @@ -929,16 +1016,22 @@ def _execute_build_on_compute( dockerhub_user = dockerhub_creds.get("username", "") dockerhub_password = dockerhub_creds.get("password", "") if dockerhub_user: - self.rich_console.print(f" Credentials: Found in credential.json") + self.rich_console.print( + f" Credentials: Found in credential.json" + ) except (json.JSONDecodeError, IOError) as e: - self.rich_console.print(f" [yellow]Warning: Could not read credential.json: {e}[/yellow]") + self.rich_console.print( + f" [yellow]Warning: Could not read credential.json: {e}[/yellow]" + ) elif dockerhub_user: - self.rich_console.print(f" Credentials: Found in environment (MAD_DOCKERHUB_USER)") - + self.rich_console.print( + f" Credentials: Found in environment (MAD_DOCKERHUB_USER)" + ) + # Determine if registry requires authentication public_registries = ["docker.io", "ghcr.io", "gcr.io", "quay.io", "nvcr.io"] registry_lower = registry.lower() if registry else "" - + # For docker.io pushes, authentication is always required # Per-registry guidance for the missing-credentials error message. # Today only Docker Hub credentials (MAD_DOCKERHUB_USER/PASSWORD or credential.json) @@ -958,7 +1051,7 @@ def _execute_build_on_compute( ], "gcr.io": [ "Google Container Registry: use a service-account JSON key as password", - "Set MAD_DOCKERHUB_USER=_json_key, MAD_DOCKERHUB_PASSWORD=\"$(cat key.json)\"", + 'Set MAD_DOCKERHUB_USER=_json_key, MAD_DOCKERHUB_PASSWORD="$(cat key.json)"', ], "quay.io": [ "Quay.io: use a robot account or encrypted password", @@ -970,7 +1063,11 @@ def _execute_build_on_compute( ], } _matched_hints = next( - (hints for reg_key, hints in _registry_hints.items() if reg_key in registry_lower), + ( + hints + for reg_key, hints in _registry_hints.items() + if reg_key in registry_lower + ), _registry_hints["docker.io"], ) if any(pub_reg in registry_lower for pub_reg in public_registries): @@ -987,10 +1084,12 @@ def _execute_build_on_compute( self.rich_console.print(f" Auth: Will login to registry before push") else: # Private/internal registry - may not need auth - self.rich_console.print(f" Auth: Private registry (auth may not be required)") - + self.rich_console.print( + f" Auth: Private registry (auth may not be required)" + ) + self.rich_console.print("") - + # Check if we're inside an existing allocation inside_allocation = os.environ.get("SLURM_JOB_ID") is not None existing_job_id = os.environ.get("SLURM_JOB_ID", "") @@ -1031,21 +1130,27 @@ def _execute_build_on_compute( # Docker repository names must be lowercase; cover the edge case where the # Dockerfile filename is just `Dockerfile` (no suffix to strip). dockerfile_basename = ( - Path(dockerfile_path).name.replace(".Dockerfile", "").replace(".ubuntu.amd", "") + Path(dockerfile_path) + .name.replace(".Dockerfile", "") + .replace(".ubuntu.amd", "") ) local_image_name = f"ci-{model_name}_{dockerfile_basename}".lower() if len(registry_parts) >= 2: registry_image_name = f"{registry}:{model_name}" else: registry_image_name = f"{registry}/{model_name}:latest" - self.rich_console.print(f" {model_name}: {dockerfile_path} -> {registry_image_name}") - per_model_data.append({ - "model": model, - "model_name": model_name, - "dockerfile_path": dockerfile_path, - "local_image_name": local_image_name, - "registry_image_name": registry_image_name, - }) + self.rich_console.print( + f" {model_name}: {dockerfile_path} -> {registry_image_name}" + ) + per_model_data.append( + { + "model": model, + "model_name": model_name, + "dockerfile_path": dockerfile_path, + "local_image_name": local_image_name, + "registry_image_name": registry_image_name, + } + ) self.rich_console.print("") # Shell-quote values that end up inside bash commands (defense-in-depth, @@ -1150,29 +1255,31 @@ def _execute_build_on_compute( exit 0 """ - + build_script_path = Path("madengine_build_job.sh") build_script_path.write_text(build_script_content) build_script_path.chmod(0o755) - + if inside_allocation: - self.rich_console.print(f"[cyan]Running build via srun (inside allocation {existing_job_id})...[/cyan]") + self.rich_console.print( + f"[cyan]Running build via srun (inside allocation {existing_job_id})...[/cyan]" + ) cmd = ["srun", "-N1", "--ntasks=1", "bash", str(build_script_path)] else: self.rich_console.print("[cyan]Submitting build job via sbatch...[/cyan]") cmd = ["sbatch", "--wait", str(build_script_path)] - + self.rich_console.print(f" Build script: {build_script_path}") self.rich_console.print(f" Command: {' '.join(cmd)}") self.rich_console.print("") - + try: result = subprocess.run( cmd, capture_output=False, text=True, ) - + if result.returncode != 0: raise BuildError( f"Build on compute node failed with exit code {result.returncode}", @@ -1187,10 +1294,12 @@ def _execute_build_on_compute( "Verify registry credentials are configured", ], ) - + # Generate manifest keyed by model_name โ€” same shape as _execute_with_prebuilt_image # so ContainerRunner.run_models_from_manifest() can join via built_models.get(key). - self.rich_console.print(f"\n[bold cyan]๐Ÿ“„ Generating manifest...[/bold cyan]") + self.rich_console.print( + f"\n[bold cyan]๐Ÿ“„ Generating manifest...[/bold cyan]" + ) built_images: Dict = {} built_models: Dict = {} @@ -1235,7 +1344,9 @@ def _execute_build_on_compute( "successful_builds": list(built_models.keys()), "failed_builds": [], "total_build_time": 0, - "successful_pushes": [pmd["registry_image_name"] for pmd in per_model_data], + "successful_pushes": [ + pmd["registry_image_name"] for pmd in per_model_data + ], "failed_pushes": [], }, } @@ -1245,12 +1356,14 @@ def _execute_build_on_compute( self.rich_console.print(f"[green]โœ“ Build completed on compute node[/green]") for pmd in per_model_data: - self.rich_console.print(f"[green]โœ“ Image pushed: {pmd['registry_image_name']}[/green]") + self.rich_console.print( + f"[green]โœ“ Image pushed: {pmd['registry_image_name']}[/green]" + ) self.rich_console.print(f"[green]โœ“ Manifest: {manifest_output}[/green]") self.rich_console.print(f"[dim]{'=' * 60}[/dim]\n") - + return manifest_output - + except subprocess.TimeoutExpired: raise BuildError( "Build on compute node timed out", @@ -1269,6 +1382,196 @@ def _execute_build_on_compute( component="BuildOrchestrator", ), ) from e + + def _execute_bare_metal_build( + self, manifest_output: str = "build_manifest.json" + ) -> str: + """Create conda environments for discovered models (no Docker build). + + The conda env replaces the Docker image as the dependency-isolation unit. + For each model this creates/updates the env (via CondaEnvManager) and + writes a synthetic manifest whose entries mark ``bare_metal: True`` so the + run phase routes to BareMetalRunner. The manifest shape mirrors + ``_execute_with_prebuilt_image`` (keyed by model name) so the run-phase + join ``built_models.get(image_name)`` still works. + + Args: + manifest_output: Output file for the build manifest. + + Returns: + Path to the generated manifest. + + Raises: + DiscoveryError: If no models are discovered. + BuildError: If conda env creation fails for all models. + """ + from madengine.execution.conda_env import ( + CondaEnvManager, + resolve_conda_env_name, + ) + from madengine.utils.gpu_validator import ( + GPUInstallationError, + validate_gpu_installation, + ) + + self.rich_console.print(f"\n[dim]{'=' * 60}[/dim]") + self.rich_console.print( + "[bold blue]๐Ÿ”จ BUILD PHASE (Bare-Metal / Conda)[/bold blue]" + ) + self.rich_console.print(f"[dim]{'=' * 60}[/dim]\n") + + bm_config = self.additional_context.get("bare_metal", {}) or {} + + # Preflight: bare-metal runs the workload directly on the host, so the + # GPU driver/runtime must be healthy before we spend time solving conda + # envs. (Docker builds defer this to the container; there is no container + # here.) Fail fast with the validator's actionable message. + self.context.init_gpu_context() + gpu_vendor = self.context.ctx.get("gpu_vendor") + gpu_arch = (self.context.ctx.get("docker_env_vars") or {}).get( + "MAD_SYSTEM_GPU_ARCHITECTURE" + ) + rocm_path = getattr(self.context, "_rocm_path", None) + self.rich_console.print( + "[bold cyan]๐Ÿ” Validating GPU installation...[/bold cyan]" + ) + try: + validate_gpu_installation( + rocm_path=rocm_path, + raise_on_error=True, + ) + except GPUInstallationError as gpu_err: + raise BuildError( + str(gpu_err), + context=create_error_context( + operation="bare_metal_build", + component="BuildOrchestrator", + ), + suggestions=[ + "Check that GPU drivers and the ROCm/CUDA runtime are installed", + "Verify ROCM_PATH points at a valid ROCm installation (AMD)", + ], + ) from gpu_err + self.rich_console.print("[green]โœ“ GPU installation validated[/green]\n") + + self.rich_console.print("[bold cyan]๐Ÿ” Discovering models...[/bold cyan]") + discover_models = DiscoverModels(args=self.args) + models = discover_models.run() + if not models: + raise DiscoveryError( + "No models discovered for bare-metal build", + context=create_error_context( + operation="bare_metal_build", + component="BuildOrchestrator", + ), + suggestions=[ + "Check if models.json exists", + "Verify --tags parameter is correct", + ], + ) + self.rich_console.print(f"[green]โœ“ Found {len(models)} model(s)[/green]\n") + + conda = CondaEnvManager( + console=self.console, + bm_config=bm_config, + gpu_vendor=gpu_vendor, + gpu_arch=gpu_arch, + ) + + manifest = { + "built_images": {}, + "built_models": {}, + "context": self.context.ctx if hasattr(self.context, "ctx") else {}, + "bare_metal": True, + "credentials_required": [], + "summary": { + "successful_builds": [], + "failed_builds": [], + "total_build_time": 0, + "successful_pushes": [], + "failed_pushes": [], + }, + } + + for model in models: + model_name = model.get("name", "unknown") + try: + env_name = conda.create_or_update(model) + except Exception as e: + self.rich_console.print( + f"[red]โœ— Conda env setup failed for {model_name}: {e}[/red]" + ) + manifest["summary"]["failed_builds"].append( + {"model": model_name, "error": str(e)} + ) + continue + + manifest["built_images"][model_name] = { + "docker_image": "", + "dockerfile": "N/A (bare-metal mode)", + "build_status": "SUCCESS", + "build_time": 0, + "bare_metal": True, + "conda_env": env_name, + } + + data_field = model.get("data", []) + if isinstance(data_field, list): + data_str = ",".join(data_field) if data_field else "" + else: + data_str = data_field if data_field else "" + + manifest["built_models"][model_name] = { + "name": model_name, + "tags": model.get("tags", []), + "scripts": model.get("scripts", ""), + "n_gpus": model.get("n_gpus", "1"), + "owner": model.get("owner", ""), + "training_precision": model.get("training_precision", ""), + "args": model.get("args", ""), + "timeout": model.get("timeout", None), + "data": data_str, + "cred": model.get("cred", ""), + "multiple_results": model.get("multiple_results", ""), + "skip_gpu_arch": model.get("skip_gpu_arch", []), + "env_vars": model.get("env_vars", {}), + "conda_env": env_name, + "environment_file": model.get("environment_file", ""), + "requirements_file": model.get("requirements_file", ""), + "python_version": model.get("python_version", ""), + "setup_script": model.get("setup_script", ""), + "rocm": model.get("rocm", {}), + "bare_metal": True, + } + manifest["summary"]["successful_builds"].append(model_name) + self.rich_console.print( + f"[green]โœ“ {model_name} -> conda env '{env_name}'[/green]" + ) + + with open(manifest_output, "w") as f: + json.dump(manifest, f, indent=2) + + self._save_deployment_config(manifest_output) + + if not manifest["summary"]["successful_builds"]: + raise BuildError( + "All bare-metal env setups failed - no models ready to run", + context=create_error_context( + operation="bare_metal_build", + component="BuildOrchestrator", + ), + suggestions=[ + "Check conda/mamba is installed and on PATH", + "Verify environment_file / setup_script paths in models.json", + ], + ) + + self.rich_console.print( + f"[green]โœ“ Bare-metal build complete: {manifest_output}[/green]" + ) + self.rich_console.print(f"[dim]{'=' * 60}[/dim]\n") + return manifest_output + def _save_build_summary(self, manifest_file: str, build_summary: Dict): """Save build summary to manifest for display purposes.""" try: @@ -1282,12 +1585,16 @@ def _save_build_summary(self, manifest_file: str, build_summary: Dict): json.dump(manifest, f, indent=2) except Exception as e: - self.rich_console.print(f"[yellow]Warning: Could not save build summary: {e}[/yellow]") + self.rich_console.print( + f"[yellow]Warning: Could not save build summary: {e}[/yellow]" + ) def _save_deployment_config(self, manifest_file: str): """Save deployment_config from --additional-context to manifest.""" if not self.additional_context: - self.rich_console.print("[dim]No additional_context provided, skipping deployment config[/dim]") + self.rich_console.print( + "[dim]No additional_context provided, skipping deployment config[/dim]" + ) return try: @@ -1298,23 +1605,31 @@ def _save_deployment_config(self, manifest_file: str): # Auto-detect target from config presence if not explicitly set target = self.additional_context.get("deploy") if not target: - # Auto-detect based on config presence - if self.additional_context.get("slurm"): + # Auto-detect based on config presence. Membership, not truthiness: + # an empty {'bare_metal': {}} still means bare-metal (see execute()). + if "bare_metal" in self.additional_context: + target = "bare_metal" + elif self.additional_context.get("slurm"): target = "slurm" - elif self.additional_context.get("k8s") or self.additional_context.get("kubernetes"): + elif self.additional_context.get("k8s") or self.additional_context.get( + "kubernetes" + ): target = "k8s" else: target = "local" - + # Get env_vars and filter out MIOPEN_USER_DB_PATH # This variable must be set per-process in multi-GPU training to avoid database conflicts env_vars = self.additional_context.get("env_vars", {}).copy() if "MIOPEN_USER_DB_PATH" in env_vars: del env_vars["MIOPEN_USER_DB_PATH"] - print("โ„น๏ธ Filtered MIOPEN_USER_DB_PATH from env_vars (will be set per-process in training)") - + print( + "โ„น๏ธ Filtered MIOPEN_USER_DB_PATH from env_vars (will be set per-process in training)" + ) + deployment_config = { "target": target, + "bare_metal": self.additional_context.get("bare_metal"), "slurm": self.additional_context.get("slurm"), "k8s": self.additional_context.get("k8s"), "kubernetes": self.additional_context.get("kubernetes"), @@ -1329,17 +1644,25 @@ def _save_deployment_config(self, manifest_file: str): k: v for k, v in deployment_config.items() if v is not None } - if deployment_config and deployment_config != {"target": "local", "env_vars": {}}: + if deployment_config and deployment_config != { + "target": "local", + "env_vars": {}, + }: manifest["deployment_config"] = deployment_config with open(manifest_file, "w") as f: json.dump(manifest, f, indent=2) - self.rich_console.print(f"[green]โœ“ Saved deployment config to {manifest_file}[/green]") + self.rich_console.print( + f"[green]โœ“ Saved deployment config to {manifest_file}[/green]" + ) else: - self.rich_console.print("[dim]No deployment config to save (local execution)[/dim]") + self.rich_console.print( + "[dim]No deployment config to save (local execution)[/dim]" + ) except Exception as e: # Non-fatal - just warn - self.rich_console.print(f"[yellow]Warning: Could not save deployment config: {e}[/yellow]") - + self.rich_console.print( + f"[yellow]Warning: Could not save deployment config: {e}[/yellow]" + ) diff --git a/src/madengine/orchestration/run_orchestrator.py b/src/madengine/orchestration/run_orchestrator.py index 296af1ee..a9882663 100644 --- a/src/madengine/orchestration/run_orchestrator.py +++ b/src/madengine/orchestration/run_orchestrator.py @@ -21,8 +21,8 @@ from rich.console import Console as RichConsole from rich.panel import Panel -from madengine.core.console import Console from madengine.core.auth import load_credentials +from madengine.core.console import Console from madengine.core.context import Context from madengine.core.dataprovider import Data from madengine.core.errors import ( @@ -31,11 +31,13 @@ ExecutionError, create_error_context, ) -from madengine.utils.session_tracker import SessionTracker from madengine.orchestration.image_filtering import ( filter_images_by_gpu_compatibility as _filter_by_gpu_compat, +) +from madengine.orchestration.image_filtering import ( filter_images_by_skip_gpu_arch as _filter_by_skip_gpu_arch, ) +from madengine.utils.session_tracker import SessionTracker class RunOrchestrator: @@ -69,33 +71,44 @@ def __init__(self, args, additional_context: Optional[Dict] = None): # Use ast.literal_eval for Python dict syntax (single quotes) # This matches what Context class expects import ast + parsed = ast.literal_eval(args.additional_context) merged_context = parsed if isinstance(parsed, dict) else {} elif isinstance(args.additional_context, dict): merged_context = args.additional_context except (ValueError, SyntaxError) as e: - self.rich_console.print(f"[yellow]Warning: Could not parse additional_context: {e}[/yellow]") + self.rich_console.print( + f"[yellow]Warning: Could not parse additional_context: {e}[/yellow]" + ) if args.additional_context: - self.rich_console.print(f"[dim]Raw (first 200 chars): {str(args.additional_context)[:200]}[/dim]") + self.rich_console.print( + f"[dim]Raw (first 200 chars): {str(args.additional_context)[:200]}[/dim]" + ) pass if additional_context: merged_context.update(additional_context) self.additional_context = merged_context - keys_str = ", ".join(sorted(self.additional_context.keys())) if self.additional_context else "(none)" - self.rich_console.print(f"[dim]Run additional context (CLI):[/dim] [cyan]{keys_str}[/cyan]") + keys_str = ( + ", ".join(sorted(self.additional_context.keys())) + if self.additional_context + else "(none)" + ) + self.rich_console.print( + f"[dim]Run additional context (CLI):[/dim] [cyan]{keys_str}[/cyan]" + ) # Track if we copied MODEL_DIR contents (for cleanup) self._copied_from_model_dir = False - + # Track if we ran build phase in this workflow (for log combination) self._did_build_phase = False - + # Initialize session tracker for filtering current run results perf_csv_path = getattr(args, "output", "perf.csv") self.session_tracker = SessionTracker(perf_csv_path) - + # Initialize context in runtime mode (with GPU detection for local) # This will be lazy-initialized only when needed self.context = None @@ -105,14 +118,14 @@ def _init_runtime_context(self): """Initialize runtime context (with GPU detection).""" # Always reinitialize context in runtime mode for run phase # This ensures GPU detection and proper runtime context even after build phase - + # Context expects additional_context as a string representation of Python dict # Use repr() instead of json.dumps() because Context uses ast.literal_eval() if self.additional_context: context_string = repr(self.additional_context) else: context_string = None - + self.context = Context( additional_context=context_string, build_only_mode=False, @@ -172,7 +185,7 @@ def execute( mad_container_image = None if self.additional_context: mad_container_image = self.additional_context.get("MAD_CONTAINER_IMAGE") - + if mad_container_image: # Local image mode: Skip build, create synthetic manifest if not tags: @@ -187,14 +200,16 @@ def execute( "Example: --tags model_name --additional-context \"{'MAD_CONTAINER_IMAGE': 'rocm/tensorflow:latest'}\"", ], ) - + # Generate synthetic manifest using the provided image manifest_file = self._create_manifest_from_local_image( image_name=mad_container_image, tags=tags, - manifest_output=getattr(self.args, "manifest_output", "build_manifest.json"), + manifest_output=getattr( + self.args, "manifest_output", "build_manifest.json" + ), ) - + # Step 1: Ensure we have a manifest (build if needed) elif not manifest_file or not os.path.exists(manifest_file): if not tags: @@ -210,7 +225,9 @@ def execute( ], ) - self.rich_console.print("[cyan]No manifest found, building first...[/cyan]\n") + self.rich_console.print( + "[cyan]No manifest found, building first...[/cyan]\n" + ) manifest_file = self._build_phase(tags, registry) self._did_build_phase = True # Mark that we built in this workflow @@ -221,73 +238,104 @@ def execute( # (with optional runtime override) with open(manifest_file) as f: manifest = json.load(f) - + deployment_config = manifest.get("deployment_config", {}) - + # Update additional_context with deployment_config for deployment layer if not self.additional_context: self.additional_context = {} - + # Merge deployment_config into additional_context (for deployment layer to use) - for key in ["slurm", "k8s", "kubernetes", "distributed", "vllm", "env_vars", "debug"]: + for key in [ + "slurm", + "k8s", + "kubernetes", + "bare_metal", + "distributed", + "vllm", + "env_vars", + "debug", + ]: if key in deployment_config and key not in self.additional_context: self.additional_context[key] = deployment_config[key] - + # Display manifest entries: context (from build) and deployment_config (run/deploy) self.rich_console.print("[bold blue]Build manifest breakdown[/bold blue]\n") manifest_context = manifest.get("context", {}) - self.rich_console.print(Panel( - json.dumps(manifest_context, indent=2) if manifest_context else "(empty)", - title="[bold]Manifest context[/bold] (from build additional context)", - border_style="dim", - padding=(0, 1), - )) - self.rich_console.print(Panel( - json.dumps(deployment_config, indent=2) if deployment_config else "(empty)", - title="[bold]Manifest deployment_config[/bold]", - border_style="dim", - padding=(0, 1), - )) + self.rich_console.print( + Panel( + ( + json.dumps(manifest_context, indent=2) + if manifest_context + else "(empty)" + ), + title="[bold]Manifest context[/bold] (from build additional context)", + border_style="dim", + padding=(0, 1), + ) + ) + self.rich_console.print( + Panel( + ( + json.dumps(deployment_config, indent=2) + if deployment_config + else "(empty)" + ), + title="[bold]Manifest deployment_config[/bold]", + border_style="dim", + padding=(0, 1), + ) + ) self.rich_console.print() # Infer deployment target from config structure (Convention over Configuration) # No explicit "deploy" field needed - presence of k8s/slurm indicates deployment type target = self._infer_deployment_target(self.additional_context) - + # Legacy support: check manifest for explicit target if not target or target == "local": target = deployment_config.get("target", "local") - - self.rich_console.print(f"[bold cyan]Deployment target: {target}[/bold cyan]\n") + + self.rich_console.print( + f"[bold cyan]Deployment target: {target}[/bold cyan]\n" + ) # Step 4: Execute based on target try: - if target == "local" or target == "docker": + if target == "bare_metal": + results = self._execute_bare_metal(manifest_file, timeout) + elif target == "local" or target == "docker": results = self._execute_local(manifest_file, timeout) else: results = self._execute_distributed(target, manifest_file) - + # Combine build and run logs for full workflow if self._did_build_phase and (target == "local" or target == "docker"): self._combine_build_and_run_logs(manifest_file) - + # Add session information to results for filtering results["session_start_row"] = session_start_row - results["session_row_count"] = self.session_tracker.get_session_row_count() - + results["session_row_count"] = ( + self.session_tracker.get_session_row_count() + ) + # Always cleanup madengine package files after execution - self.rich_console.print("\n[dim]๐Ÿงน Cleaning up madengine package files...[/dim]") + self.rich_console.print( + "\n[dim]๐Ÿงน Cleaning up madengine package files...[/dim]" + ) self._cleanup_model_dir_copies() - + # NOTE: Do NOT cleanup session marker here! # It's needed by display functions in CLI layer # Cleanup happens in CLI after display (via perf_csv_path) - + return results - + except Exception as e: # Always cleanup madengine package files even on error - self.rich_console.print("\n[dim]๐Ÿงน Cleaning up madengine package files...[/dim]") + self.rich_console.print( + "\n[dim]๐Ÿงน Cleaning up madengine package files...[/dim]" + ) self._cleanup_model_dir_copies() raise @@ -333,59 +381,68 @@ def _build_phase(self, tags: list, registry: Optional[str] = None) -> str: return manifest_file def _create_manifest_from_local_image( - self, - image_name: str, - tags: list, - manifest_output: str = "build_manifest.json" + self, image_name: str, tags: list, manifest_output: str = "build_manifest.json" ) -> str: """ Create a synthetic manifest for a user-provided local image. - + This enables MAD_CONTAINER_IMAGE functionality where users can skip the build phase and directly run models using a pre-existing Docker image. - + Args: image_name: Docker image name/tag (e.g., 'rocm/tensorflow:latest') tags: Model tags to discover manifest_output: Output path for the manifest file - + Returns: Path to the generated manifest file - + Raises: DiscoveryError: If no models are found RuntimeError: If image validation fails """ - from madengine.utils.discover_models import DiscoverModels from madengine.core.errors import DiscoveryError - - self.rich_console.print(f"[yellow]๐Ÿ  Local Image Mode: Using {image_name}[/yellow]") - self.rich_console.print(f"[dim]Skipping build phase, creating synthetic manifest...[/dim]\n") - + from madengine.utils.discover_models import DiscoverModels + + self.rich_console.print( + f"[yellow]๐Ÿ  Local Image Mode: Using {image_name}[/yellow]" + ) + self.rich_console.print( + f"[dim]Skipping build phase, creating synthetic manifest...[/dim]\n" + ) + # Validate that the image exists locally or can be pulled. # image_name is interpolated into shell commands run with shell=True, # so shell-escape it to avoid command injection / breakage on special chars. quoted_image_name = shlex.quote(image_name) try: - self.console.sh(f"docker image inspect {quoted_image_name} > /dev/null 2>&1") - self.rich_console.print(f"[green]โœ“ Image {image_name} found locally[/green]") + self.console.sh( + f"docker image inspect {quoted_image_name} > /dev/null 2>&1" + ) + self.rich_console.print( + f"[green]โœ“ Image {image_name} found locally[/green]" + ) except (subprocess.CalledProcessError, RuntimeError) as e: - self.rich_console.print(f"[yellow]โš ๏ธ Image {image_name} not found locally, attempting to pull...[/yellow]") + self.rich_console.print( + f"[yellow]โš ๏ธ Image {image_name} not found locally, attempting to pull...[/yellow]" + ) try: self.console.sh(f"docker pull {quoted_image_name}") - self.rich_console.print(f"[green]โœ“ Successfully pulled {image_name}[/green]") + self.rich_console.print( + f"[green]โœ“ Successfully pulled {image_name}[/green]" + ) except Exception as e: raise RuntimeError( f"Failed to find or pull image {image_name}. " f"Ensure the image exists locally or can be pulled from a registry. " f"Error: {e}" ) - + # Discover models by tags (without building) self.args.tags = tags discover_models = DiscoverModels(args=self.args) models = discover_models.run() - + if not models: raise DiscoveryError( "No models discovered for local image mode", @@ -399,17 +456,21 @@ def _create_manifest_from_local_image( "Ensure model definitions have matching tags", ], ) - - self.rich_console.print(f"[green]โœ“ Discovered {len(models)} model(s) for tags: {tags}[/green]\n") - + + self.rich_console.print( + f"[green]โœ“ Discovered {len(models)} model(s) for tags: {tags}[/green]\n" + ) + # Initialize build-only context for manifest generation # (we need context structure, but skip GPU detection since we're not building) - context_string = repr(self.additional_context) if self.additional_context else None + context_string = ( + repr(self.additional_context) if self.additional_context else None + ) build_context = Context( additional_context=context_string, build_only_mode=True, ) - + # Create manifest structure manifest = { "built_images": {}, @@ -419,13 +480,13 @@ def _create_manifest_from_local_image( "local_image_name": image_name, "deployment_config": self.additional_context.get("deployment_config", {}), } - + # For each model, create a synthetic entry using the provided image for model in models: model_name = model["name"] # Create a synthetic image identifier (not an actual built image) synthetic_image_id = f"local-{model_name.replace('/', '_')}" - + manifest["built_images"][synthetic_image_id] = { "docker_image": image_name, # Use user-provided image "dockerfile": "N/A (local image mode)", @@ -434,22 +495,26 @@ def _create_manifest_from_local_image( "local_image": True, "registry_image": None, } - + # Convert data list to comma-separated string (required by dataprovider) data_field = model.get("data", []) if isinstance(data_field, list): data_str = ",".join(data_field) if data_field else "" else: data_str = data_field if data_field else "" - + # Build model info dict with all fields that ContainerRunner expects # Use exact field names from models.json format manifest["built_models"][synthetic_image_id] = { "name": model_name, "tags": model.get("tags", []), "dockerfile": "N/A (local image mode)", - "scripts": model.get("scripts", ""), # models.json uses "scripts" (plural) - "n_gpus": model.get("n_gpus", "1"), # models.json uses "n_gpus" (string format) + "scripts": model.get( + "scripts", "" + ), # models.json uses "scripts" (plural) + "n_gpus": model.get( + "n_gpus", "1" + ), # models.json uses "n_gpus" (string format) "owner": model.get("owner", ""), "training_precision": model.get("training_precision", ""), "args": model.get("args", ""), # Required field for docker run @@ -458,16 +523,22 @@ def _create_manifest_from_local_image( "cred": model.get("cred", ""), "deprecated": model.get("deprecated", False), "skip_gpu_arch": model.get("skip_gpu_arch", []), - "additional_docker_run_options": model.get("additional_docker_run_options", ""), + "additional_docker_run_options": model.get( + "additional_docker_run_options", "" + ), } - + # Write manifest to file with open(manifest_output, "w") as f: json.dump(manifest, f, indent=2) - - self.rich_console.print(f"[green]โœ“ Generated synthetic manifest: {manifest_output}[/green]") - self.rich_console.print(f"[yellow]โš ๏ธ Warning: User-provided image {image_name}. Model support not guaranteed.[/yellow]\n") - + + self.rich_console.print( + f"[green]โœ“ Generated synthetic manifest: {manifest_output}[/green]" + ) + self.rich_console.print( + f"[yellow]โš ๏ธ Warning: User-provided image {image_name}. Model support not guaranteed.[/yellow]\n" + ) + return manifest_output def _load_and_merge_manifest(self, manifest_file: str) -> str: @@ -486,22 +557,32 @@ def _load_and_merge_manifest(self, manifest_file: str) -> str: if "deployment_config" in manifest: stored_config = manifest["deployment_config"] # Runtime --additional-context overrides stored config - for key in ["deploy", "slurm", "k8s", "kubernetes", "distributed", "vllm", "env_vars", "debug"]: + for key in [ + "deploy", + "slurm", + "k8s", + "kubernetes", + "bare_metal", + "distributed", + "vllm", + "env_vars", + "debug", + ]: if key in self.additional_context: stored_config[key] = self.additional_context[key] manifest["deployment_config"] = stored_config - + # Merge context (tools, pre_scripts, post_scripts, encapsulate_script) if "context" not in manifest: manifest["context"] = {} - + merge_keys = ["tools", "pre_scripts", "post_scripts", "encapsulate_script"] context_updated = False for key in merge_keys: if key in self.additional_context: manifest["context"][key] = self.additional_context[key] context_updated = True - + if context_updated or "deployment_config" in manifest: # Write back merged config with open(manifest_file, "w") as f: @@ -517,16 +598,18 @@ def _execute_local(self, manifest_file: str, timeout: int) -> Dict: # Load manifest first to check if we have Docker images with open(manifest_file, "r") as f: manifest = json.load(f) - + has_docker_images = bool(manifest.get("built_images", {})) - + if has_docker_images: # Using Docker containers - containers have GPU support built-in - self.rich_console.print("[dim cyan]Using Docker containers with built-in GPU support[/dim cyan]\n") - + self.rich_console.print( + "[dim cyan]Using Docker containers with built-in GPU support[/dim cyan]\n" + ) + # Initialize runtime context (runs full GPU detection on compute nodes) self._init_runtime_context() - + # Show node info self._show_node_info() @@ -545,7 +628,9 @@ def _execute_local(self, manifest_file: str, timeout: int) -> Dict: if "docker_mounts" in manifest_context: if "docker_mounts" not in self.context.ctx: self.context.ctx["docker_mounts"] = {} - for container_path, host_path in manifest_context["docker_mounts"].items(): + for container_path, host_path in manifest_context[ + "docker_mounts" + ].items(): if container_path not in self.context.ctx["docker_mounts"]: self.context.ctx["docker_mounts"][container_path] = host_path if "docker_build_arg" in manifest_context: @@ -554,9 +639,15 @@ def _execute_local(self, manifest_file: str, timeout: int) -> Dict: for key, value in manifest_context["docker_build_arg"].items(): if key not in self.context.ctx["docker_build_arg"]: self.context.ctx["docker_build_arg"][key] = value - if "docker_gpus" in manifest_context and "docker_gpus" not in self.context.ctx: + if ( + "docker_gpus" in manifest_context + and "docker_gpus" not in self.context.ctx + ): self.context.ctx["docker_gpus"] = manifest_context["docker_gpus"] - if "gpu_vendor" in manifest_context and "gpu_vendor" not in self.context.ctx: + if ( + "gpu_vendor" in manifest_context + and "gpu_vendor" not in self.context.ctx + ): self.context.ctx["gpu_vendor"] = manifest_context["gpu_vendor"] if "guest_os" in manifest_context and "guest_os" not in self.context.ctx: self.context.ctx["guest_os"] = manifest_context["guest_os"] @@ -567,12 +658,17 @@ def _execute_local(self, manifest_file: str, timeout: int) -> Dict: if "post_scripts" in manifest_context: self.context.ctx["post_scripts"] = manifest_context["post_scripts"] if "encapsulate_script" in manifest_context: - self.context.ctx["encapsulate_script"] = manifest_context["encapsulate_script"] + self.context.ctx["encapsulate_script"] = manifest_context[ + "encapsulate_script" + ] # Restore docker_env_vars from build context (e.g. MAD_SECRETS_HFTOKEN for Primus HF-backed configs). # Keep runtime-detected values as priority (consistent with docker_mounts / docker_build_arg): # values already populated by Context (e.g. MAD_SECRETS_* read from os.environ) must not be # overwritten by manifest entries that may still contain unexpanded "${VAR}" placeholders. - if "docker_env_vars" in manifest_context and manifest_context["docker_env_vars"]: + if ( + "docker_env_vars" in manifest_context + and manifest_context["docker_env_vars"] + ): if "docker_env_vars" not in self.context.ctx: self.context.ctx["docker_env_vars"] = {} for k, v in manifest_context["docker_env_vars"].items(): @@ -590,9 +686,13 @@ def _execute_local(self, manifest_file: str, timeout: int) -> Dict: if "pre_scripts" in self.additional_context: self.context.ctx["pre_scripts"] = self.additional_context["pre_scripts"] if "post_scripts" in self.additional_context: - self.context.ctx["post_scripts"] = self.additional_context["post_scripts"] + self.context.ctx["post_scripts"] = self.additional_context[ + "post_scripts" + ] if "encapsulate_script" in self.additional_context: - self.context.ctx["encapsulate_script"] = self.additional_context["encapsulate_script"] + self.context.ctx["encapsulate_script"] = self.additional_context[ + "encapsulate_script" + ] # Filter images by GPU vendor and architecture # Filter images by GPU compatibility @@ -605,10 +705,14 @@ def _execute_local(self, manifest_file: str, timeout: int) -> Dict: if has_docker_images: # Docker images: filter by GPU vendor at runtime to avoid cross-vendor execution - self.rich_console.print("[dim cyan]Filtering Docker images by runtime GPU compatibility...[/dim cyan]") + self.rich_console.print( + "[dim cyan]Filtering Docker images by runtime GPU compatibility...[/dim cyan]" + ) else: # Bare-metal execution: filter by runtime GPU - self.rich_console.print("[dim cyan]Filtering bare-metal images by runtime GPU compatibility...[/dim cyan]") + self.rich_console.print( + "[dim cyan]Filtering bare-metal images by runtime GPU compatibility...[/dim cyan]" + ) compatible_images = self._filter_images_by_gpu_compatibility( manifest["built_images"], runtime_gpu_vendor, runtime_gpu_arch @@ -630,30 +734,37 @@ def _execute_local(self, manifest_file: str, timeout: int) -> Dict: manifest["built_images"] = compatible_images print(f"Filtered to {len(compatible_images)} compatible images\n") - + # Filter by skip_gpu_arch from model definitions (applies to both Docker and bare-metal) runtime_gpu_arch = self.context.get_system_gpu_architecture() if "built_models" in manifest and compatible_images: - self.rich_console.print("[cyan]Checking skip_gpu_arch model restrictions...[/cyan]") + self.rich_console.print( + "[cyan]Checking skip_gpu_arch model restrictions...[/cyan]" + ) compatible_images = self._filter_images_by_skip_gpu_arch( compatible_images, manifest["built_models"], runtime_gpu_arch ) manifest["built_images"] = compatible_images - print(f"After skip_gpu_arch filtering: {len(compatible_images)} images to run\n") - + print( + f"After skip_gpu_arch filtering: {len(compatible_images)} images to run\n" + ) + # NOTE: Dockerfile context filtering is already done during build phase # Re-filtering during run phase causes issues because: # 1. The build phase already filtered dockerfiles based on build-time context # 2. All built images should be runnable on the runtime node # 3. Legacy behavior: filtering happens once (either build or run, not both) - + # Write filtered manifest back to file so runner sees the filtered list with open(manifest_file, "w") as f: json.dump(manifest, f, indent=2) except Exception as e: import traceback - self.rich_console.print(f"[yellow]Warning: GPU/Context filtering failed: {e}[/yellow]") + + self.rich_console.print( + f"[yellow]Warning: GPU/Context filtering failed: {e}[/yellow]" + ) self.rich_console.print(f"[red]Traceback: {traceback.format_exc()}[/red]") self.rich_console.print("[yellow]Proceeding with all images[/yellow]\n") @@ -693,6 +804,72 @@ def _execute_local(self, manifest_file: str, timeout: int) -> Dict: return results + def _execute_bare_metal(self, manifest_file: str, timeout: int) -> Dict: + """Execute models on bare metal (conda env, no Docker). + + Like ``_execute_local`` but skips Docker-image GPU filtering (images are + synthetic bare-metal entries) and uses ``BareMetalRunner`` instead of + ``ContainerRunner``. GPU detection still runs natively via the runtime + context; scripts/common is still populated for pre/post scripts. + """ + self.rich_console.print("[cyan]Executing on bare metal (conda)...[/cyan]\n") + + with open(manifest_file, "r") as f: + manifest = json.load(f) + + # Initialize runtime context (native GPU detection, no Docker needed). + self._init_runtime_context() + self._show_node_info() + + from madengine.execution.bare_metal_runner import BareMetalRunner + + credentials = load_credentials() + + # Restore host-level context from manifest (pre/post scripts, tools, + # encapsulate, guest_os) so pre/post scripts and tools apply on bare metal. + if "context" in manifest: + manifest_context = manifest["context"] + for key in ( + "tools", + "pre_scripts", + "post_scripts", + "encapsulate_script", + "guest_os", + ): + if key in manifest_context and key not in self.context.ctx: + self.context.ctx[key] = manifest_context[key] + + # Runtime --additional-context overrides manifest. + if self.additional_context: + for key in ("tools", "pre_scripts", "post_scripts", "encapsulate_script"): + if key in self.additional_context: + self.context.ctx[key] = self.additional_context[key] + + # Populate scripts/common so pre/post scripts resolve. + self._copy_scripts() + + runner = BareMetalRunner( + self.context, + self.data, + self.console, + live_output=getattr(self.args, "live_output", False), + additional_context=self.additional_context, + ) + runner.set_credentials(credentials) + if hasattr(self.args, "output") and self.args.output: + runner.set_perf_csv_path(self.args.output) + + results = runner.run_models_from_manifest( + manifest_file=manifest_file, + timeout=timeout, + skip_model_run=getattr(self.args, "skip_model_run", False), + phase_suffix=".run", + ) + + self.rich_console.print("\n[green]โœ“ Bare-metal execution complete[/green]") + self.rich_console.print(f"[dim]{'=' * 60}[/dim]\n") + return results + def _execute_distributed(self, target: str, manifest_file: str) -> Dict: """Execute on distributed infrastructure.""" self.rich_console.print(f"[cyan]Deploying to {target}...[/cyan]\n") @@ -712,13 +889,15 @@ def _execute_distributed(self, target: str, manifest_file: str) -> Dict: ) # Import from deployment layer - from madengine.deployment.factory import DeploymentFactory from madengine.deployment.base import DeploymentConfig + from madengine.deployment.factory import DeploymentFactory # Add runtime flags to additional_context for deployment layer if "live_output" not in self.additional_context: - self.additional_context["live_output"] = getattr(self.args, "live_output", False) - + self.additional_context["live_output"] = getattr( + self.args, "live_output", False + ) + # Pass session_start_row for result filtering in collect_results session_start_row = self.session_tracker.session_start_row if "session_start_row" not in self.additional_context: @@ -773,37 +952,39 @@ def _show_node_info(self): elif "HOST_AZURE" in host_os: print(self.console.sh("tdnf info rocm-libs", canFail=True)) else: - self.rich_console.print("[yellow]Warning: Unable to detect host OS[/yellow]") + self.rich_console.print( + "[yellow]Warning: Unable to detect host OS[/yellow]" + ) def _cleanup_model_dir_copies(self): """Clean up only madengine package files from scripts/common directory. - + This cleanup removes ONLY the files that were copied from madengine package: - scripts/common/tools.json - scripts/common/test_echo.sh - scripts/common/pre_scripts/ - scripts/common/post_scripts/ - scripts/common/tools/ - + This preserves the user's actual scripts/ and docker/ directories in MAD project. """ import shutil import subprocess - + # Only clean up scripts/common/ subdirectories that came from madengine package common_dir = Path("scripts/common") if not common_dir.exists(): return - + # List of items to clean up (from madengine package) items_to_cleanup = [ "tools.json", "test_echo.sh", "pre_scripts", "post_scripts", - "tools" + "tools", ] - + for item_name in items_to_cleanup: item_path = common_dir / item_name if item_path.exists(): @@ -814,14 +995,20 @@ def _cleanup_model_dir_copies(self): subprocess.run( ["chmod", "-R", "+w", str(item_path)], capture_output=True, - timeout=10 + timeout=10, ) - except (subprocess.TimeoutExpired, subprocess.CalledProcessError, OSError) as e: + except ( + subprocess.TimeoutExpired, + subprocess.CalledProcessError, + OSError, + ) as e: print(f"Warning: chmod failed for {item_path}: {e}") shutil.rmtree(item_path) else: item_path.unlink() - self.rich_console.print(f"[dim] Cleaned up: scripts/common/{item_name}[/dim]") + self.rich_console.print( + f"[dim] Cleaned up: scripts/common/{item_name}[/dim]" + ) except Exception as e: # Try with sudo for permission issues try: @@ -829,9 +1016,11 @@ def _cleanup_model_dir_copies(self): ["sudo", "rm", "-rf", str(item_path)], check=True, capture_output=True, - timeout=10 + timeout=10, + ) + self.rich_console.print( + f"[dim] Cleaned up: scripts/common/{item_name} (elevated)[/dim]" ) - self.rich_console.print(f"[dim] Cleaned up: scripts/common/{item_name} (elevated)[/dim]") except Exception as e2: self.rich_console.print( f"[yellow]โš ๏ธ Warning: Could not clean up {item_path}: {e2}[/yellow]" @@ -839,84 +1028,88 @@ def _cleanup_model_dir_copies(self): def _combine_build_and_run_logs(self, manifest_file: str): """Combine build.live.log and run.live.log into live.log for full workflow. - + For full workflow (build + run), this creates a unified log file by: 1. Reading the manifest to find models that were actually executed in this session 2. Finding corresponding *.build.live.log and *.run.live.log files for those models 3. Concatenating them into *.live.log 4. Keeping the original build and run logs for reference - + Args: manifest_file: Path to the manifest file containing executed models """ import json - + # Load manifest to get list of build log files try: with open(manifest_file, "r") as f: manifest = json.load(f) - + built_images = manifest.get("built_images", {}) if not built_images: return # No models to process except Exception as e: - self.rich_console.print(f"[yellow]โš ๏ธ Warning: Could not load manifest for log combining: {e}[/yellow]") + self.rich_console.print( + f"[yellow]โš ๏ธ Warning: Could not load manifest for log combining: {e}[/yellow]" + ) return - + self.rich_console.print("\n[dim]๐Ÿ“ Combining build and run logs...[/dim]") combined_count = 0 - + # Process each built image for image_name, image_info in built_images.items(): # Get build log file name from manifest build_log = image_info.get("log_file") if not build_log or not os.path.exists(build_log): continue # Skip if build log doesn't exist - + # Derive the base name and corresponding run log base_name = build_log.replace(".build.live.log", "") run_log = f"{base_name}.run.live.log" combined_log = f"{base_name}.live.log" - + # Check if run log exists if not os.path.exists(run_log): continue # Skip if run log doesn't exist - + try: # Combine build and run logs - with open(combined_log, 'w') as outfile: + with open(combined_log, "w") as outfile: # Add build log - with open(build_log, 'r') as infile: + with open(build_log, "r") as infile: outfile.write(infile.read()) - + # Add separator outfile.write("\n" + "=" * 80 + "\n") outfile.write("RUN PHASE LOG\n") outfile.write("=" * 80 + "\n\n") - + # Add run log - with open(run_log, 'r') as infile: + with open(run_log, "r") as infile: outfile.write(infile.read()) - + combined_count += 1 self.rich_console.print(f"[dim] Combined: {combined_log}[/dim]") - + except Exception as e: self.rich_console.print( f"[yellow]โš ๏ธ Warning: Could not combine logs for {base_name}: {e}[/yellow]" ) - + if combined_count > 0: - self.rich_console.print(f"[dim]โœ“ Combined {combined_count} log file(s)[/dim]") + self.rich_console.print( + f"[dim]โœ“ Combined {combined_count} log file(s)[/dim]" + ) def _copy_scripts(self): """Copy common scripts to model directories. - + Handles scenarios: 1. MAD Project: scripts/ already exists in current directory - just add madengine common files 2. External MODEL_DIR: Copy from external path to current directory 3. madengine Testing: Copy from src/madengine/scripts/common - + NOTE: Does NOT delete existing scripts/ or docker/ directories in current working directory. """ import shutil @@ -924,19 +1117,27 @@ def _copy_scripts(self): # Define ignore function for cache files (used for all copy operations) def ignore_cache_files(directory, files): """Ignore Python cache files and directories.""" - return [f for f in files if f.endswith('.pyc') or f == '__pycache__' or f.endswith('.pyo')] - + return [ + f + for f in files + if f.endswith(".pyc") or f == "__pycache__" or f.endswith(".pyo") + ] + # Step 1: Check if MODEL_DIR points to external directory and copy if needed # MODEL_DIR default is "." (current directory), so only copy if it's different model_dir_env = os.environ.get("MODEL_DIR", ".") model_dir_abs = os.path.abspath(model_dir_env) current_dir_abs = os.path.abspath(".") - + # Only copy if MODEL_DIR points to a different directory (not current dir) if model_dir_abs != current_dir_abs and os.path.exists(model_dir_env): - self.rich_console.print(f"[yellow]๐Ÿ“ External MODEL_DIR detected: {model_dir_env}[/yellow]") - self.rich_console.print("[yellow]Copying MODEL_DIR contents for run phase...[/yellow]") - + self.rich_console.print( + f"[yellow]๐Ÿ“ External MODEL_DIR detected: {model_dir_env}[/yellow]" + ) + self.rich_console.print( + "[yellow]Copying MODEL_DIR contents for run phase...[/yellow]" + ) + # Copy docker/ and scripts/ from MODEL_DIR (without deleting existing ones first) for subdir in ["docker", "scripts"]: src_path = Path(model_dir_env) / subdir @@ -945,18 +1146,29 @@ def ignore_cache_files(directory, files): # Use copytree with dirs_exist_ok=True to merge instead of replace if dest_path.exists(): # Only warn, don't delete existing directories - self.rich_console.print(f"[dim] Note: Merging {subdir}/ from MODEL_DIR with existing directory[/dim]") - shutil.copytree(src_path, dest_path, dirs_exist_ok=True, ignore=ignore_cache_files) - - self.rich_console.print("[green]โœ“ MODEL_DIR structure copied (docker/, scripts/)[/green]") + self.rich_console.print( + f"[dim] Note: Merging {subdir}/ from MODEL_DIR with existing directory[/dim]" + ) + shutil.copytree( + src_path, + dest_path, + dirs_exist_ok=True, + ignore=ignore_cache_files, + ) + + self.rich_console.print( + "[green]โœ“ MODEL_DIR structure copied (docker/, scripts/)[/green]" + ) elif not os.path.exists(model_dir_env): - self.rich_console.print(f"[yellow]โš ๏ธ Warning: MODEL_DIR '{model_dir_env}' does not exist, using current directory[/yellow]") + self.rich_console.print( + f"[yellow]โš ๏ธ Warning: MODEL_DIR '{model_dir_env}' does not exist, using current directory[/yellow]" + ) # Step 2: Copy madengine's common scripts (pre_scripts, post_scripts, tools) # This provides the execution framework scripts # Find madengine installation path (works for both development and installed package) madengine_common = None - + # Option 1: Development mode - check if running from source dev_path = Path("src/madengine/scripts/common") if dev_path.exists(): @@ -966,23 +1178,34 @@ def ignore_cache_files(directory, files): # Option 2: Installed package - find via module location try: import madengine + madengine_module_path = Path(madengine.__file__).parent installed_path = madengine_module_path / "scripts" / "common" if installed_path.exists(): madengine_common = installed_path - print(f"Found madengine scripts in installed package: {madengine_common}") + print( + f"Found madengine scripts in installed package: {madengine_common}" + ) except Exception as e: print(f"Could not locate madengine scripts: {e}") - + if madengine_common and madengine_common.exists(): - print(f"Copying madengine common scripts from {madengine_common} to scripts/common") - + print( + f"Copying madengine common scripts from {madengine_common} to scripts/common" + ) + dest_common = Path("scripts/common") # Ensure the destination directory exists before copying dest_common.mkdir(parents=True, exist_ok=True) - + # Copy pre_scripts, post_scripts, tools if they exist - for item in ["pre_scripts", "post_scripts", "tools", "tools.json", "test_echo.sh"]: + for item in [ + "pre_scripts", + "post_scripts", + "tools", + "tools.json", + "test_echo.sh", + ]: src_item = madengine_common / item if src_item.exists(): dest_item = dest_common / item @@ -991,19 +1214,21 @@ def ignore_cache_files(directory, files): shutil.rmtree(dest_item) else: dest_item.unlink() - + if src_item.is_dir(): shutil.copytree(src_item, dest_item, ignore=ignore_cache_files) else: shutil.copy2(src_item, dest_item) print(f" Copied {item}") else: - self.rich_console.print("[yellow]โš ๏ธ Could not find madengine scripts directory[/yellow]") + self.rich_console.print( + "[yellow]โš ๏ธ Could not find madengine scripts directory[/yellow]" + ) # Step 3: REMOVED - Distribution to model directories is incorrect # scripts/common should remain at /scripts/common/ for proper relative path access # Model scripts reference it via ../scripts/common/ from their directory (e.g., scripts/dummy/) - # + # # This ensures compatibility with legacy workflow where: # - scripts/common/ stays at working directory root # - Model scripts use ../scripts/common/ relative paths @@ -1024,7 +1249,9 @@ def _filter_images_by_gpu_compatibility( ) compatible_images[model_name] = image_info continue - built_with_vendor = {k: v for k, v in built_images.items() if v.get("gpu_vendor")} + built_with_vendor = { + k: v for k, v in built_images.items() if v.get("gpu_vendor") + } compat, skipped = _filter_by_gpu_compat( built_with_vendor, runtime_gpu_vendor, runtime_gpu_arch ) @@ -1032,7 +1259,7 @@ def _filter_images_by_gpu_compatibility( for model_name, reason in skipped: self.rich_console.print(f"[dim] Skipping {model_name}: {reason}[/dim]") return compatible_images - + def _filter_images_by_gpu_architecture( self, built_images: Dict, runtime_gpu_arch: str ) -> Dict: @@ -1063,19 +1290,22 @@ def _filter_images_by_skip_gpu_arch( self._write_skipped_status(model_name, image_info, gpu_arch) return compatible_images - def _write_skipped_status(self, model_name: str, image_info: Dict, gpu_arch: str) -> None: + def _write_skipped_status( + self, model_name: str, image_info: Dict, gpu_arch: str + ) -> None: """Write SKIPPED status to perf CSV for models that were skipped. - + Args: model_name: Name of the model that was skipped image_info: Image information dictionary gpu_arch: GPU architecture that caused the skip """ try: - from madengine.reporting.update_perf_csv import update_perf_csv import json import tempfile - + + from madengine.reporting.update_perf_csv import update_perf_csv + # Create a perf entry for the skipped model perf_entry = { "model": model_name, @@ -1083,45 +1313,51 @@ def _write_skipped_status(self, model_name: str, image_info: Dict, gpu_arch: str "reason": f"Model not supported on {gpu_arch} architecture", "gpu_architecture": gpu_arch, } - + # Write to temporary JSON file - with tempfile.NamedTemporaryFile(mode='w', suffix='.json', delete=False) as f: + with tempfile.NamedTemporaryFile( + mode="w", suffix=".json", delete=False + ) as f: json.dump(perf_entry, f) temp_file = f.name - + # Get output CSV path from args - output_csv = getattr(self.args, 'output', 'perf.csv') - + output_csv = getattr(self.args, "output", "perf.csv") + # Update perf CSV with skipped entry update_perf_csv(exception_result=temp_file, perf_csv=output_csv) - + # Clean up temp file import os + os.unlink(temp_file) - + except Exception as e: - self.rich_console.print(f"[dim] Warning: Could not write SKIPPED status to CSV: {e}[/dim]") + self.rich_console.print( + f"[dim] Warning: Could not write SKIPPED status to CSV: {e}[/dim]" + ) def _infer_deployment_target(self, config: Dict) -> str: """ Infer deployment target from configuration structure. - + Convention over Configuration: + - Presence of "bare_metal" field โ†’ bare-metal (conda, no Docker) - Presence of "k8s" or "kubernetes" field โ†’ k8s deployment - Presence of "slurm" field โ†’ slurm deployment - - Neither present โ†’ local execution - + - None present โ†’ local execution (Docker) + Args: config: Configuration dictionary - + Returns: - Deployment target: "k8s", "slurm", or "local" + Deployment target: "bare_metal", "k8s", "slurm", or "local" """ - if "k8s" in config or "kubernetes" in config: + if "bare_metal" in config: + return "bare_metal" + elif "k8s" in config or "kubernetes" in config: return "k8s" elif "slurm" in config: return "slurm" else: return "local" - - diff --git a/src/madengine/utils/discover_models.py b/src/madengine/utils/discover_models.py index fe795e7b..e6d1135c 100644 --- a/src/madengine/utils/discover_models.py +++ b/src/madengine/utils/discover_models.py @@ -5,11 +5,12 @@ # built-in modules import argparse -import os -import json import importlib.util +import json +import os import typing -from dataclasses import dataclass, field, asdict +from dataclasses import asdict, dataclass, field + from rich.console import Console as RichConsole @@ -32,6 +33,11 @@ class CustomModel: args: str = "" multiple_results: str = "" skip_gpu_arch: str = "" + # Bare-metal (conda) execution fields (all optional). + conda_env: str = "" + environment_file: str = "" + python_version: str = "" + setup_script: str = "" def to_dict(self) -> dict: return asdict(self) @@ -73,51 +79,57 @@ def _setup_model_dir_if_needed(self) -> None: This copies docker/, scripts/, and config files (models.json, credential.json, data.json) from MODEL_DIR to the current working directory to support the model discovery process. This operation is safe for build-only (CPU) nodes as it only involves file operations. - + MODEL_DIR defaults to "." (current directory) if not set. Only copies if MODEL_DIR points to a different directory than current working directory. """ model_dir_env = os.environ.get("MODEL_DIR", ".") - + # Get absolute paths to compare model_dir_abs = os.path.abspath(model_dir_env) cwd_abs = os.path.abspath(".") - + # Only copy if MODEL_DIR points to a different directory (not current dir) if model_dir_abs != cwd_abs: import shlex import subprocess from pathlib import Path - self.rich_console.print(f"[bold cyan]๐Ÿ“ MODEL_DIR environment variable detected:[/bold cyan] [yellow]{model_dir_env}[/yellow]") + self.rich_console.print( + f"[bold cyan]๐Ÿ“ MODEL_DIR environment variable detected:[/bold cyan] [yellow]{model_dir_env}[/yellow]" + ) print(f"Copying required files to current working directory: {cwd_abs}") try: # Check if source directory exists if not os.path.exists(model_dir_env): - self.rich_console.print(f"[yellow]โš ๏ธ Warning: MODEL_DIR path does not exist: {model_dir_env}[/yellow]") + self.rich_console.print( + f"[yellow]โš ๏ธ Warning: MODEL_DIR path does not exist: {model_dir_env}[/yellow]" + ) return # Copy specific directories and files only (not everything with /*) # This prevents copying unwanted subdirectories from MODEL_DIR items_to_copy = [] - + # Directories to copy for subdir in ["docker", "scripts"]: src_path = Path(model_dir_env) / subdir if src_path.exists(): items_to_copy.append((src_path, subdir, "directory")) - + # Files to copy for file in ["models.json", "credential.json", "data.json"]: src_file = Path(model_dir_env) / file if src_file.exists(): items_to_copy.append((src_file, file, "file")) - + if not items_to_copy: - self.rich_console.print(f"[yellow]โš ๏ธ No required files/directories found in MODEL_DIR[/yellow]") + self.rich_console.print( + f"[yellow]โš ๏ธ No required files/directories found in MODEL_DIR[/yellow]" + ) return - + # Copy each item copied_count = 0 for src_path, item_name, item_type in items_to_copy: @@ -127,7 +139,7 @@ def _setup_model_dir_if_needed(self) -> None: cmd, shell=True, capture_output=True, text=True, check=True ) copied_count += 1 - + if result.stdout: # Show summary for directories, full output for files if item_type == "directory": @@ -135,21 +147,29 @@ def _setup_model_dir_if_needed(self) -> None: if len(lines) < 10: print(result.stdout) else: - print(f" โœ“ Copied {item_name}/ ({len(lines)} files)") + print( + f" โœ“ Copied {item_name}/ ({len(lines)} files)" + ) else: print(f" โœ“ Copied {item_name}") except subprocess.CalledProcessError as e: - self.rich_console.print(f"[yellow]โš ๏ธ Warning: Failed to copy {item_name}: {e}[/yellow]") + self.rich_console.print( + f"[yellow]โš ๏ธ Warning: Failed to copy {item_name}: {e}[/yellow]" + ) if e.stderr: print(f" Error details: {e.stderr}") # Continue with other items even if one fails - + if copied_count > 0: - self.rich_console.print(f"[green]โœ… Successfully copied {copied_count} item(s) from MODEL_DIR[/green]") - + self.rich_console.print( + f"[green]โœ… Successfully copied {copied_count} item(s) from MODEL_DIR[/green]" + ) + print(f"Model dir: {model_dir_env} โ†’ current dir: {cwd_abs}") except Exception as e: - self.rich_console.print(f"[yellow]โš ๏ธ Warning: Unexpected error copying MODEL_DIR: {e}[/yellow]") + self.rich_console.print( + f"[yellow]โš ๏ธ Warning: Unexpected error copying MODEL_DIR: {e}[/yellow]" + ) # Continue execution even if copy fails def discover_models(self) -> None: @@ -179,7 +199,9 @@ def discover_models(self) -> None: files = os.listdir(root) if "models.json" in files and "get_models_json.py" in files: - self.rich_console.print(f"[red]โŒ Both models.json and get_models_json.py found in {root}.[/red]") + self.rich_console.print( + f"[red]โŒ Both models.json and get_models_json.py found in {root}.[/red]" + ) raise ValueError( f"Both models.json and get_models_json.py found in {root}." ) @@ -311,11 +333,27 @@ def select_models(self) -> None: custom_model.update_model() dirname = custom_model.name.split("/")[0] custom_model.dockerfile = os.path.normpath( - os.path.join("scripts", dirname, custom_model.dockerfile) + os.path.join( + "scripts", dirname, custom_model.dockerfile + ) ) custom_model.scripts = os.path.normpath( os.path.join("scripts", dirname, custom_model.scripts) ) + if custom_model.environment_file: + custom_model.environment_file = os.path.normpath( + os.path.join( + "scripts", + dirname, + custom_model.environment_file, + ) + ) + if custom_model.setup_script: + custom_model.setup_script = os.path.normpath( + os.path.join( + "scripts", dirname, custom_model.setup_script + ) + ) model_dict = custom_model.to_dict() model_dict["args"] = model_dict["args"] + extra_args tag_models.append(model_dict) @@ -339,17 +377,35 @@ def select_models(self) -> None: custom_model.update_model() dirname = custom_model.name.split("/")[0] custom_model.dockerfile = os.path.normpath( - os.path.join("scripts", dirname, custom_model.dockerfile) + os.path.join( + "scripts", dirname, custom_model.dockerfile + ) ) custom_model.scripts = os.path.normpath( os.path.join("scripts", dirname, custom_model.scripts) ) + if custom_model.environment_file: + custom_model.environment_file = os.path.normpath( + os.path.join( + "scripts", + dirname, + custom_model.environment_file, + ) + ) + if custom_model.setup_script: + custom_model.setup_script = os.path.normpath( + os.path.join( + "scripts", dirname, custom_model.setup_script + ) + ) model_dict = custom_model.to_dict() model_dict["args"] = model_dict["args"] + extra_args tag_models.append(model_dict) if not tag_models: - self.rich_console.print(f"[red]โŒ No models found corresponding to the given tag: {tag}[/red]") + self.rich_console.print( + f"[red]โŒ No models found corresponding to the given tag: {tag}[/red]" + ) raise ValueError( f"No models found corresponding to the given tag: {tag}" ) @@ -359,11 +415,15 @@ def select_models(self) -> None: def print_models(self) -> None: if self.selected_models: # print selected models using parsed tags and adding backslash-separated extra args - self.rich_console.print(f"[bold green]๐Ÿ“‹ Selected Models ({len(self.selected_models)} models):[/bold green]") + self.rich_console.print( + f"[bold green]๐Ÿ“‹ Selected Models ({len(self.selected_models)} models):[/bold green]" + ) print(json.dumps(self.selected_models, indent=4)) else: # print list of all model names - self.rich_console.print(f"[bold cyan]๐Ÿ“Š Available Models ({len(self.model_list)} total):[/bold cyan]") + self.rich_console.print( + f"[bold cyan]๐Ÿ“Š Available Models ({len(self.model_list)} total):[/bold cyan]" + ) for model_name in self.model_list: print(f" {model_name}") diff --git a/tests/e2e/test_run_workflows.py b/tests/e2e/test_run_workflows.py index 32f6141b..6d6e9c27 100644 --- a/tests/e2e/test_run_workflows.py +++ b/tests/e2e/test_run_workflows.py @@ -3,36 +3,38 @@ Copyright (c) Advanced Micro Devices, Inc. All rights reserved. """ +import csv +import json + # built-in modules import os -import csv # third-party modules import pytest -import json + +from madengine.core.context import Context # project modules -from tests.fixtures.utils import BASE_DIR, MODEL_DIR -from tests.fixtures.utils import global_data -from tests.fixtures.utils import clean_test_temp_files -from tests.fixtures.utils import get_gpu_nodeid_map -from tests.fixtures.utils import get_num_gpus -from tests.fixtures.utils import get_num_cpus -from tests.fixtures.utils import requires_gpu -from tests.fixtures.utils import generate_additional_context_for_machine from tests.fixtures.utils import ( + BASE_DIR, DEFAULT_CLEAN_FILES, - build_run_command, + MODEL_DIR, assert_model_in_perf_csv, + build_run_command, + clean_test_temp_files, + generate_additional_context_for_machine, + get_gpu_nodeid_map, + get_num_cpus, + get_num_gpus, + global_data, + requires_gpu, ) -from madengine.core.context import Context - - # ============================================================================ # Context Handling Tests # ============================================================================ + class TestContexts: @pytest.mark.parametrize( @@ -331,9 +333,7 @@ def test_docker_mounts_mount_host_paths_in_docker_container( ) @requires_gpu("docker gpus requires GPU hardware") - @pytest.mark.skipif( - get_num_gpus() < 8, reason="test requires atleast 8 gpus" - ) + @pytest.mark.skipif(get_num_gpus() < 8, reason="test requires atleast 8 gpus") @pytest.mark.parametrize( "clean_test_temp_files", [["perf.csv", "perf.html", "results_dummy_gpubind.csv"]], @@ -364,24 +364,24 @@ def test_docker_gpus(self, global_data, clean_test_temp_files): gpu_node_ids.append(row["performance"]) else: pytest.fail("model in perf_test.csv did not run successfully.") - + # Debug information print(f"GPU node IDs from performance: {gpu_node_ids}") print(f"GPU nodeid map: {gpu_nodeid_map}") mapped_gpus = [gpu_nodeid_map.get(node_id) for node_id in gpu_node_ids] print(f"Mapped GPUs: {mapped_gpus}") - + # Filter out None values and sort valid_mapped_gpus = [gpu for gpu in mapped_gpus if gpu is not None] sorted_gpus = sorted(valid_mapped_gpus) print(f"Sorted valid GPUs: {sorted_gpus}") - + if sorted_gpus != [0, 2, 3, 4, 5, 7]: - pytest.fail(f"docker_gpus did not bind expected gpus in docker container. Expected: [0, 2, 3, 4, 5, 7], Got: {sorted_gpus}, Raw node IDs: {gpu_node_ids}, Mapping: {gpu_nodeid_map}") + pytest.fail( + f"docker_gpus did not bind expected gpus in docker container. Expected: [0, 2, 3, 4, 5, 7], Got: {sorted_gpus}, Raw node IDs: {gpu_node_ids}, Mapping: {gpu_nodeid_map}" + ) - @pytest.mark.skipif( - get_num_cpus() < 64, reason="test requires atleast 64 cpus" - ) + @pytest.mark.skipif(get_num_cpus() < 64, reason="test requires atleast 64 cpus") @pytest.mark.parametrize( "clean_test_temp_files", [["perf.csv", "perf.html", "results_dummy_cpubind.csv"]], @@ -425,25 +425,25 @@ def test_gpu_product_name_matches_arch(self): """ context = Context() - product_name = context.ctx['docker_env_vars']["MAD_SYSTEM_GPU_PRODUCT_NAME"] + product_name = context.ctx["docker_env_vars"]["MAD_SYSTEM_GPU_PRODUCT_NAME"] - #fail the test if GPU product name is empty + # fail the test if GPU product name is empty if not product_name or not product_name.strip(): pytest.fail("GPU product name is empty or just whitespaces") product_name = product_name.upper() - #if product name has AMD or NVIDIA in it then it's a safe bet - #that it was parsed properly + # if product name has AMD or NVIDIA in it then it's a safe bet + # that it was parsed properly if not ("AMD" in product_name or "NVIDIA" in product_name): pytest.fail(f"Incorrect product name={product_name!r}") - # ============================================================================ # Tag Filtering Tests # ============================================================================ + class TestTagsFunctionality: @pytest.mark.parametrize( @@ -463,7 +463,7 @@ def test_can_select_model_subset_with_commandline_tag_argument( + "MODEL_DIR=" + MODEL_DIR + " " - + f"python3 -m madengine.cli.app run --tags dummy_group_1 --live-output --additional-context '{json.dumps(context)}'" + + f"python3 -m madengine.cli.app run --tags dummy_group_1 --live-output --additional-context '{json.dumps(context)}'" ) # Check for model execution (handles ANSI codes in output) @@ -520,7 +520,7 @@ def test_model_names_are_automatically_tags( + "MODEL_DIR=" + MODEL_DIR + " " - + f"python3 -m madengine.cli.app run --tags dummy --live-output --additional-context '{json.dumps(context)}'" + + f"python3 -m madengine.cli.app run --tags dummy --live-output --additional-context '{json.dumps(context)}'" ) # Check for model execution (handles ANSI codes in output) @@ -528,3 +528,54 @@ def test_model_names_are_automatically_tags( pytest.fail("dummy tag not selected with commandline --tags argument") +class TestBareMetalExecution: + """E2E tests for the bare-metal (conda) execution backend. + + These run without a real conda/mamba install by pointing + ``bare_metal.conda_bin`` at a checked-in fake-conda shim that execs the + wrapped command directly. The conda_bin path must be absolute because the + runner cd's into the model's working directory before invoking it. + """ + + FAKE_CONDA = os.path.join( + BASE_DIR, + MODEL_DIR, + "scripts", + "dummy_bare_metal", + "fake_conda.sh", + ) + RUN_LOG = "dummy_bare_metal_bare_metal_mad_dummy_bare_metal.run.live.log" + + @pytest.mark.parametrize( + "clean_test_temp_files", + [DEFAULT_CLEAN_FILES + [RUN_LOG]], + indirect=True, + ) + def test_bare_metal_run_succeeds(self, global_data, clean_test_temp_files): + """dummy_bare_metal runs via the conda backend and reports SUCCESS. + + Verifies the run routed to the bare-metal path (deployment_type, + launcher) and that performance was extracted from the model log. + """ + context = { + "bare_metal": { + "conda_bin": os.path.abspath(self.FAKE_CONDA), + "reuse_env": True, + } + } + global_data["console"].sh( + build_run_command("dummy_bare_metal", additional_context=context) + ) + + perf_csv = os.path.join(BASE_DIR, "perf.csv") + assert_model_in_perf_csv(perf_csv, "dummy_bare_metal", status="SUCCESS") + + with open(perf_csv, "r") as csv_file: + for row in csv.DictReader(csv_file): + if row["model"] == "dummy_bare_metal": + assert row["deployment_type"] == "bare_metal" + assert row["launcher"] == "conda" + assert row["performance"], "expected a performance value" + break + else: + pytest.fail("dummy_bare_metal not found in perf.csv") diff --git a/tests/fixtures/dummy/models.json b/tests/fixtures/dummy/models.json index 8c8a29c1..bff526a7 100644 --- a/tests/fixtures/dummy/models.json +++ b/tests/fixtures/dummy/models.json @@ -28,6 +28,21 @@ "args": "", "timeout": 360 }, + { + "name": "dummy_bare_metal", + "dockerfile": "docker/dummy", + "scripts": "scripts/dummy_bare_metal/run.sh", + "n_gpus": "1", + "owner": "mad.support@amd.com", + "training_precision": "", + "tags": [ + "dummies", + "dummy_bare_metal" + ], + "args": "", + "conda_env": "mad_dummy_bare_metal", + "python_version": "3.11" + }, { "name": "dummy3", "dockerfile": "docker/dummy", diff --git a/tests/fixtures/dummy/scripts/dummy_bare_metal/fake_conda.sh b/tests/fixtures/dummy/scripts/dummy_bare_metal/fake_conda.sh new file mode 100755 index 00000000..679550bb --- /dev/null +++ b/tests/fixtures/dummy/scripts/dummy_bare_metal/fake_conda.sh @@ -0,0 +1,39 @@ +#!/bin/bash +# +# Copyright (c) Advanced Micro Devices, Inc. +# All rights reserved. +# +# Minimal fake conda for e2e-testing the bare-metal runner without a real +# conda/mamba install. Supports the subcommands CondaEnvManager invokes: +# conda env list | conda env create | conda env update | conda env remove +# conda create | conda run -n --no-capture-output +# `run` execs the wrapped command with CONDA_DEFAULT_ENV set to . + +case "$1" in + env) + case "$2" in + list) + echo "# conda environments" + echo "base * /tmp/fake_conda/base" + ;; + create|update) echo "fake conda env $2 ok" ;; + remove) echo "fake conda env remove ok" ;; + esac + ;; + create) echo "fake conda create ok" ;; + run) + shift + env_name="" + while [[ "$1" == -* || "$1" == "-n" ]]; do + if [[ "$1" == "-n" ]]; then + env_name="$2" + shift 2 + else + shift + fi + done + export CONDA_DEFAULT_ENV="$env_name" + exec "$@" + ;; + *) echo "fake conda: unknown $*" ;; +esac diff --git a/tests/fixtures/dummy/scripts/dummy_bare_metal/run.sh b/tests/fixtures/dummy/scripts/dummy_bare_metal/run.sh new file mode 100755 index 00000000..cf5b9bda --- /dev/null +++ b/tests/fixtures/dummy/scripts/dummy_bare_metal/run.sh @@ -0,0 +1,8 @@ +#!/bin/bash +# +# Copyright (c) Advanced Micro Devices, Inc. +# All rights reserved. +# + +echo "running dummy_bare_metal in conda env: ${CONDA_DEFAULT_ENV:-unknown}" +echo "performance: $RANDOM samples_per_second" diff --git a/tests/unit/test_bare_metal.py b/tests/unit/test_bare_metal.py new file mode 100644 index 00000000..b0f71ae2 --- /dev/null +++ b/tests/unit/test_bare_metal.py @@ -0,0 +1,449 @@ +"""Unit tests for the bare-metal (conda) execution backend.""" + +import os +from unittest import mock + +import pytest + +from madengine.execution.conda_env import ( + CondaEnvManager, + bootstrap_micromamba, + resolve_conda_env_name, + resolve_environment_file, + resolve_rocm_index_url, +) +from madengine.execution.run_reporting import ( + determine_status, + extract_performance_from_log, +) +from madengine.orchestration.run_orchestrator import RunOrchestrator + +# ---- Deployment target inference ---- + + +class TestBareMetalInference: + """_infer_deployment_target routes bare_metal before slurm/k8s/local.""" + + def _make_orch(self): + args = mock.MagicMock() + args.additional_context = "{}" + args.output = "perf.csv" + return RunOrchestrator(args) + + def test_bare_metal_key_selects_bare_metal(self): + orch = self._make_orch() + assert orch._infer_deployment_target({"bare_metal": {}}) == "bare_metal" + + def test_bare_metal_wins_over_slurm(self): + orch = self._make_orch() + assert ( + orch._infer_deployment_target({"bare_metal": {}, "slurm": {}}) + == "bare_metal" + ) + + def test_no_key_is_local(self): + orch = self._make_orch() + assert orch._infer_deployment_target({}) == "local" + + +# ---- Conda env name resolution ---- + + +class TestResolveCondaEnvName: + def test_config_overrides_model(self): + assert resolve_conda_env_name({"conda_env": "m"}, {"conda_env": "c"}) == "c" + + def test_model_used_when_no_config(self): + assert resolve_conda_env_name({"conda_env": "m"}, {}) == "m" + + def test_derived_from_name(self): + assert resolve_conda_env_name({"name": "foo/bar"}, {}) == "mad_foo_bar" + + +# ---- CondaEnvManager command construction ---- + + +class TestCondaEnvManager: + def _mgr(self, bm_config=None): + console = mock.MagicMock() + mgr = CondaEnvManager(console=console, bm_config=bm_config or {}) + mgr._conda_bin = "/opt/conda/bin/conda" # bypass PATH detection + return mgr, console + + def test_conda_run_prefix(self): + mgr, _ = self._mgr() + prefix = mgr.conda_run_prefix("myenv") + assert "run -n myenv" in prefix + assert "--no-capture-output" in prefix + + def test_env_exists_matches_name(self): + mgr, console = self._mgr() + console.sh.return_value = ( + "# conda envs\nbase * /opt/conda\nmyenv /opt/conda/envs/myenv" + ) + assert mgr.env_exists("myenv") is True + assert mgr.env_exists("absent") is False + + def test_create_with_python_version(self): + mgr, console = self._mgr() + # env does not exist + console.sh.return_value = "# conda envs\nbase * /opt/conda" + name = mgr.create_or_update( + {"name": "m", "conda_env": "e", "python_version": "3.11"} + ) + assert name == "e" + # Find the create command among the calls. + create_calls = [ + c.args[0] for c in console.sh.call_args_list if "create" in c.args[0] + ] + assert any("python=3.11" in cmd and "-n e" in cmd for cmd in create_calls) + + def test_reuse_existing_env_skips_create(self): + mgr, console = self._mgr(bm_config={"reuse_env": True}) + console.sh.return_value = "myenv /opt/conda/envs/myenv" + mgr.create_or_update({"name": "m", "conda_env": "myenv"}) + create_calls = [ + c.args[0] + for c in console.sh.call_args_list + if "conda create" in c.args[0] or "env create" in c.args[0] + ] + assert create_calls == [] + + +# ---- Vendor-aware dependency file resolution ---- + + +class TestResolveEnvironmentFile: + def test_prefers_amd_variant(self, tmp_path): + base = tmp_path / "environment.yml" + base.write_text("name: base\n") + amd = tmp_path / "environment.amd.yml" + amd.write_text("name: amd\n") + assert resolve_environment_file(str(base), "AMD") == str(amd) + + def test_prefers_nvidia_variant(self, tmp_path): + base = tmp_path / "environment.yml" + base.write_text("name: base\n") + nvidia = tmp_path / "environment.nvidia.yml" + nvidia.write_text("name: nvidia\n") + assert resolve_environment_file(str(base), "nvidia") == str(nvidia) + + def test_falls_back_to_base_when_no_variant(self, tmp_path): + base = tmp_path / "environment.yml" + base.write_text("name: base\n") + assert resolve_environment_file(str(base), "AMD") == str(base) + + def test_unknown_vendor_returns_base(self, tmp_path): + base = tmp_path / "environment.yml" + base.write_text("name: base\n") + # Even if an amd variant exists, an unknown vendor keeps the base. + (tmp_path / "environment.amd.yml").write_text("name: amd\n") + assert resolve_environment_file(str(base), "") == str(base) + + def test_setup_script_suffix(self, tmp_path): + base = tmp_path / "setup_script.sh" + base.write_text("echo base\n") + amd = tmp_path / "setup_script.amd.sh" + amd.write_text("echo amd\n") + assert resolve_environment_file(str(base), "AMD") == str(amd) + + def test_create_or_update_uses_vendor_variant(self, tmp_path): + base = tmp_path / "environment.yml" + base.write_text("name: base\n") + amd = tmp_path / "environment.amd.yml" + amd.write_text("name: amd\n") + console = mock.MagicMock() + console.sh.return_value = "# conda envs\nbase * /opt/conda" + mgr = CondaEnvManager( + console=console, + bm_config={"conda_env": "e", "environment_file": str(base)}, + gpu_vendor="AMD", + ) + mgr._conda_bin = "/opt/conda/bin/conda" + mgr._dep_hash_path = lambda env_name: str(tmp_path / f"hash_{env_name}") + mgr.create_or_update({"name": "m"}) + env_calls = [ + c.args[0] for c in console.sh.call_args_list if "env create" in c.args[0] + ] + assert any(str(amd) in cmd for cmd in env_calls) + + +# ---- ROCm index URL resolution ---- + + +class TestResolveRocmIndexUrl: + def test_explicit_url_returned_unchanged(self): + assert ( + resolve_rocm_index_url("https://example.com/wheels/", "gfx942") + == "https://example.com/wheels/" + ) + + def test_auto_resolves_from_gfx_arch(self): + assert ( + resolve_rocm_index_url("auto", "gfx942") + == "https://rocm.nightlies.amd.com/v2/gfx942/" + ) + + def test_empty_resolves_from_gfx_arch(self): + assert ( + resolve_rocm_index_url("", "gfx90a") + == "https://rocm.nightlies.amd.com/v2/gfx90a/" + ) + + def test_auto_without_arch_raises(self): + with pytest.raises(RuntimeError): + resolve_rocm_index_url("auto", "") + + +# ---- ROCm wheel install ---- + + +class TestInstallRocmWheels: + def _mgr(self, bm_config=None, gpu_arch="gfx942"): + console = mock.MagicMock() + mgr = CondaEnvManager( + console=console, bm_config=bm_config or {}, gpu_arch=gpu_arch + ) + mgr._conda_bin = "/opt/conda/bin/conda" + return mgr, console + + def test_installs_default_packages_from_auto_index(self): + mgr, console = self._mgr() + mgr._install_rocm_wheels("e", {"enabled": True}, timeout=None) + cmds = [c.args[0] for c in console.sh.call_args_list] + assert any( + "pip install --index-url" in c + and "rocm[libraries,devel]" in c + and "gfx942" in c + for c in cmds + ) + # torch not requested -> no torch install + assert not any("torch torchvision" in c for c in cmds) + + def test_installs_torch_when_requested(self): + mgr, console = self._mgr() + mgr._install_rocm_wheels("e", {"enabled": True, "torch": True}, timeout=None) + cmds = [c.args[0] for c in console.sh.call_args_list] + assert any("torch torchvision" in c and "gfx942" in c for c in cmds) + + def test_custom_packages_and_explicit_index(self): + mgr, console = self._mgr(gpu_arch=None) + mgr._install_rocm_wheels( + "e", + {"enabled": True, "index_url": "https://x/whl/", "packages": ["rocm"]}, + timeout=None, + ) + cmds = [c.args[0] for c in console.sh.call_args_list] + assert any("https://x/whl/" in c and " rocm" in c for c in cmds) + + def test_create_or_update_installs_rocm_on_fresh_env(self): + mgr, console = self._mgr( + bm_config={"conda_env": "e", "rocm": {"enabled": True}} + ) + console.sh.return_value = "# conda envs\nbase * /opt/conda" + mgr.create_or_update({"name": "m"}) + cmds = [c.args[0] for c in console.sh.call_args_list] + assert any("pip install --index-url" in c and "gfx942" in c for c in cmds) + + def test_create_or_update_skips_rocm_on_reuse(self): + mgr, console = self._mgr( + bm_config={"conda_env": "e", "rocm": {"enabled": True}, "reuse_env": True} + ) + console.sh.return_value = "e /opt/conda/envs/e" + mgr.create_or_update({"name": "m"}) + cmds = [c.args[0] for c in console.sh.call_args_list] + assert not any("pip install --index-url" in c for c in cmds) + + +# ---- requirements_file install ---- + + +class TestRequirementsFile: + def _mgr(self, bm_config): + console = mock.MagicMock() + console.sh.return_value = "# conda envs\nbase * /opt/conda" + mgr = CondaEnvManager(console=console, bm_config=bm_config) + mgr._conda_bin = "/opt/conda/bin/conda" + return mgr, console + + def test_pip_install_requirements(self, tmp_path): + req = tmp_path / "requirements.txt" + req.write_text("numpy\n") + mgr, console = self._mgr({"conda_env": "e", "requirements_file": str(req)}) + with mock.patch.object(mgr, "_dep_hash_path", return_value=str(tmp_path / "h")): + mgr.create_or_update({"name": "m"}) + cmds = [c.args[0] for c in console.sh.call_args_list] + assert any(f"pip install -r {str(req)}" in c for c in cmds) + + def test_missing_requirements_file_raises(self): + mgr, _ = self._mgr( + {"conda_env": "e", "requirements_file": "/nonexistent/requirements.txt"} + ) + with pytest.raises(RuntimeError): + mgr.create_or_update({"name": "m"}) + + +# ---- Dependency-hash reuse invalidation ---- + + +class TestDepHashInvalidation: + def _mgr(self, bm_config, tmp_path): + console = mock.MagicMock() + mgr = CondaEnvManager(console=console, bm_config=bm_config) + mgr._conda_bin = "/opt/conda/bin/conda" + # Route the stamp file into tmp so tests don't touch the real cache. + mgr._dep_hash_path = lambda env_name: str(tmp_path / f"hash_{env_name}") + return mgr, console + + def test_env_update_on_changed_environment_file(self, tmp_path): + envf = tmp_path / "environment.yml" + envf.write_text("name: base\n") + mgr, console = self._mgr( + {"conda_env": "e", "environment_file": str(envf), "reuse_env": True}, + tmp_path, + ) + # First run: env does not exist -> creates and stamps hash. + console.sh.return_value = "# conda envs\nbase * /opt/conda" + mgr.create_or_update({"name": "m"}) + + # Second run: env now exists; unchanged file -> reuse (no env command). + console.sh.reset_mock() + console.sh.return_value = "e /opt/conda/envs/e" + mgr.create_or_update({"name": "m"}) + cmds = [c.args[0] for c in console.sh.call_args_list] + assert not any("env update" in c or "env create" in c for c in cmds) + + # Third run: change file content -> env update forced despite reuse. + envf.write_text("name: base\ndependencies: [numpy]\n") + console.sh.reset_mock() + console.sh.return_value = "e /opt/conda/envs/e" + mgr.create_or_update({"name": "m"}) + cmds = [c.args[0] for c in console.sh.call_args_list] + assert any("env update" in c for c in cmds) + + +# ---- micromamba bootstrap ---- + + +class TestBootstrapMicromamba: + def test_returns_cached_binary_without_download(self, tmp_path): + cache = tmp_path / ".cache" / "madengine" / "micromamba" + cache.mkdir(parents=True) + binary = cache / "micromamba" + binary.write_text("#!/bin/sh\n") + binary.chmod(0o755) + with mock.patch( + "madengine.execution.conda_env.os.path.expanduser", + return_value=str(tmp_path), + ), mock.patch("madengine.execution.conda_env.urllib.request.urlretrieve") as dl: + result = bootstrap_micromamba() + assert result == str(binary) + dl.assert_not_called() + + def test_detect_conda_bin_bootstraps_when_nothing_found(self, tmp_path): + console = mock.MagicMock() + mgr = CondaEnvManager(console=console, bm_config={}) + with mock.patch( + "madengine.execution.conda_env.shutil.which", return_value=None + ), mock.patch( + "madengine.execution.conda_env.bootstrap_micromamba", + return_value="/cache/micromamba", + ) as boot: + assert mgr.detect_conda_bin() == "/cache/micromamba" + boot.assert_called_once() + + +# ---- Opt-in env teardown ---- + + +class TestCleanupEnv: + def _runner(self, bm_config): + from madengine.execution.bare_metal_runner import BareMetalRunner + + runner = BareMetalRunner.__new__(BareMetalRunner) + runner.additional_context = {"bare_metal": bm_config} + runner.bm_config = bm_config + runner.perf_csv_path = "perf.csv" + runner.rich_console = mock.MagicMock() + runner.conda = mock.MagicMock() + runner.ensure_perf_csv_exists = mock.MagicMock() + runner._create_run_details = mock.MagicMock(return_value={}) + return runner + + def test_cleanup_env_removes_env(self): + runner = self._runner({"conda_env": "e", "cleanup_env": True}) + with mock.patch("madengine.execution.bare_metal_runner.write_perf_records"): + runner._record({"name": "m"}, {}, {"status": "SUCCESS"}, 1) + runner.conda.remove.assert_called_once_with("e") + + def test_no_cleanup_by_default(self): + runner = self._runner({"conda_env": "e"}) + with mock.patch("madengine.execution.bare_metal_runner.write_perf_records"): + runner._record({"name": "m"}, {}, {"status": "SUCCESS"}, 1) + runner.conda.remove.assert_not_called() + + +# ---- Preflight GPU validation in bare-metal build ---- + + +class TestBareMetalBuildPreflight: + def test_gpu_validation_failure_surfaces_as_build_error(self): + from madengine.core.errors import BuildError + from madengine.orchestration.build_orchestrator import BuildOrchestrator + from madengine.utils.gpu_validator import ( + GPUInstallationError, + GPUValidationResult, + GPUVendor, + ) + + orch = BuildOrchestrator.__new__(BuildOrchestrator) + orch.additional_context = {"bare_metal": {}} + orch.args = mock.MagicMock() + orch.console = mock.MagicMock() + orch.rich_console = mock.MagicMock() + orch.context = mock.MagicMock() + orch.context.init_gpu_context = mock.MagicMock() + orch.context.ctx = {"gpu_vendor": "AMD"} + orch.context._rocm_path = "/bogus/rocm" + + bad = GPUValidationResult(is_valid=False, vendor=GPUVendor.AMD) + bad.issues.append("ROCm not found") + with mock.patch( + "madengine.utils.gpu_validator.validate_gpu_installation", + side_effect=GPUInstallationError(bad), + ): + with pytest.raises(BuildError): + orch._execute_bare_metal_build() + + +# ---- Performance extraction ---- + + +class TestExtractPerformance: + def test_canonical_pattern(self, tmp_path): + log = tmp_path / "run.log" + log.write_text("some output\nperformance: 123.5 samples_per_second\n") + perf, metric = extract_performance_from_log(str(log)) + assert perf == "123.5" + assert metric == "samples_per_second" + + def test_missing_file(self): + perf, metric = extract_performance_from_log("/nonexistent/x.log") + assert perf is None and metric is None + + +class TestDetermineStatus: + def test_success_with_performance(self, tmp_path): + log = tmp_path / "run.log" + log.write_text("performance: 5 samples_per_second\n") + assert determine_status(str(log), "5", {}, {}) == "SUCCESS" + + def test_failure_no_performance(self, tmp_path): + log = tmp_path / "run.log" + log.write_text("nothing useful\n") + assert determine_status(str(log), None, {}, {}) == "FAILURE" + + def test_failure_on_error_pattern(self, tmp_path): + log = tmp_path / "run.log" + log.write_text("performance: 5 x\nTraceback (most recent call last)\n") + assert determine_status(str(log), "5", {}, {}) == "FAILURE"