diff --git a/docs/configuration.md b/docs/configuration.md index 4831cc4f..dcdc4e3d 100644 --- a/docs/configuration.md +++ b/docs/configuration.md @@ -697,6 +697,42 @@ export MAD_DOCKERHUB_PASSWORD=mytoken export MAD_DOCKERHUB_REPO=myorg ``` +### Registry Authentication + +madengine reuses an existing `docker login` — including an organization access +token (OAT) — rather than requiring credentials to be duplicated into +`credential.json`. Ambient credentials are read from +`${DOCKER_CONFIG:-~/.docker}/config.json` exactly as the Docker CLI reads them, +covering `auths` entries, `credHelpers`, and `credsStore`. + +Blank values are treated as **not configured**, not as credentials. A placeholder +entry such as `{"username": "", "password": ""}` will never override or break a +working `docker login`. + +| Existing `docker login` | Credentials in `credential.json` / env | Behavior | +|---|---|---| +| no | yes (non-blank) | `docker login` with the configured credentials | +| yes | yes (non-blank) | `docker login` with the configured credentials (explicit wins) | +| yes | absent or blank | Reuse the existing login; no `docker login` is run | +| no | absent or blank | Push fails with an actionable error; pull warns and continues | + +Before `docker build`, madengine logs in to the base image's registry only when +that registry has no existing login and usable credentials are configured, so a +node authenticated with an OAT is never re-authenticated. + +Relevant environment variables: + +- `DOCKER_CONFIG` — directory holding `config.json` (default `~/.docker`) +- `MAD_SKIP_DOCKER_LOGIN=1` — never run `docker login`; always defer to the + credentials the machine already has + +If a base image pull is denied, madengine distinguishes the two causes: +credentials were rejected (authentication — supply credentials or run +`docker login`), versus credentials were accepted but the registry granted no +pull scope for that repository (`insufficient_scope` — an authorization problem, +where the access token needs to be scoped to the repository and re-running +`docker login` will not help). + ## Configuration Priority For Kubernetes/SLURM deployments: diff --git a/src/madengine/core/auth.py b/src/madengine/core/auth.py index 15f0a0a6..021aa6dc 100644 --- a/src/madengine/core/auth.py +++ b/src/madengine/core/auth.py @@ -11,7 +11,8 @@ import json import os import shlex -from typing import Dict, Optional +from pathlib import Path +from typing import Dict, Optional, Tuple from madengine.core.errors import ( ConfigurationError, @@ -19,6 +20,19 @@ handle_error, ) +# Keys under which the Docker CLI stores Docker Hub credentials in config.json. +# Docker has used several spellings over the years and any of them means +# "this machine is authenticated to Docker Hub". +_DOCKERHUB_CONFIG_KEYS: Tuple[str, ...] = ( + "https://index.docker.io/v1/", + "index.docker.io", + "registry-1.docker.io", + "docker.io", +) + +# Registry values that madengine treats as "Docker Hub" rather than a host. +_DOCKERHUB_ALIASES = ("docker.io", "dockerhub") + def load_credentials() -> Optional[Dict]: """Load credentials from credential.json and environment variables. @@ -81,6 +95,174 @@ def load_credentials() -> Optional[Dict]: return credentials +def _registry_config_keys(registry: Optional[str]) -> Tuple[str, ...]: + """Map a madengine registry value to the keys Docker uses in config.json. + + Args: + registry: Registry URL (e.g. ``"localhost:5000"``, ``"docker.io/rocm"``), + or ``None``/empty string for Docker Hub. + + Returns: + The config.json ``auths`` keys that would hold credentials for it. + """ + if not registry or registry.lower() in _DOCKERHUB_ALIASES: + return _DOCKERHUB_CONFIG_KEYS + # Downstream code derives the registry host the same way (docker login ). + host = registry.split("/")[0] + if host.lower() in _DOCKERHUB_ALIASES: + return _DOCKERHUB_CONFIG_KEYS + return (host,) + + +def has_ambient_docker_auth(registry: Optional[str]) -> bool: + """Report whether the local Docker CLI is already authenticated to ``registry``. + + Reads ``${DOCKER_CONFIG:-~/.docker}/config.json`` the same way the Docker CLI + does, so an existing ``docker login`` (e.g. an organisation access token) is + honoured instead of being overridden or reported as "no credentials". + + Args: + registry: Registry URL, or ``None``/empty string for Docker Hub. + + Returns: + ``True`` if a usable credential entry exists for the registry. Any read + or parse problem yields ``False``; this function never raises and never + logs credential material. + """ + config_dir = os.environ.get("DOCKER_CONFIG") or os.path.join( + os.path.expanduser("~"), ".docker" + ) + try: + config = json.loads(Path(config_dir, "config.json").read_text(encoding="utf-8")) + except (OSError, ValueError): + return False + if not isinstance(config, dict): + return False + + keys = _registry_config_keys(registry) + auths = config.get("auths") or {} + cred_helpers = config.get("credHelpers") or {} + creds_store = config.get("credsStore") + + if isinstance(cred_helpers, dict) and any(key in cred_helpers for key in keys): + return True + if not isinstance(auths, dict): + return False + for key in keys: + entry = auths.get(key) + if not isinstance(entry, dict): + continue + if entry.get("auth") or entry.get("identitytoken") or entry.get("username"): + return True + # Credential-store-managed entries are persisted as an empty object; + # the secret itself lives in the external store. + if creds_store: + return True + return False + + +def _usable_credentials(creds: object) -> bool: + """Report whether a credential entry carries a non-blank username and password. + + Placeholder entries such as ``{"username": "", "password": ""}`` are treated + as "not configured" rather than as credentials, matching + :func:`madengine.deployment.k8s_secrets.build_registry_secret_data`. + """ + if not isinstance(creds, dict): + return False + return bool(str(creds.get("username") or "").strip()) and bool( + str(creds.get("password") or "").strip() + ) + + +def _registry_from_image(image: Optional[str]) -> Optional[str]: + """Extract the registry host from an image reference. + + Args: + image: Image reference (e.g. ``"ghcr.io/org/app:tag"``), or ``None``. + + Returns: + The registry host, or ``None`` when the reference targets Docker Hub, + is malformed, or was not supplied. + """ + if not image or not image.strip(): + return None + ref = image.strip() + if "/" not in ref: + return None + # Docker treats the first component as a registry only when it looks like a + # host; otherwise it is a Docker Hub namespace (e.g. "rocm/private"). + host = ref.split("/")[0] + if not ("." in host or ":" in host or host == "localhost"): + return None + if host.lower() in _DOCKERHUB_ALIASES: + return None + return host + + +def explain_registry_denial( + log_text: str, image: Optional[str] = None +) -> Optional[str]: + """Turn a registry denial in Docker output into an actionable explanation. + + Distinguishes "authenticated but not authorized for this repository" from + "not authenticated at all", because the two need completely different fixes. + + Args: + log_text: Docker build/pull output to inspect. + image: Optional image reference the denial refers to, for the message. + + Returns: + A multi-line hint, or ``None`` if the output shows no registry denial. + """ + lowered = (log_text or "").lower() + subject = image.strip() if image and image.strip() else "the base image" + + if "insufficient_scope" in lowered or "authorization failed" in lowered: + return ( + f"Base image pull was denied: {subject}\n" + " The registry ACCEPTED the credentials but granted no pull scope " + "for this repository.\n" + " This is an authorization problem, not a login problem:\n" + " - the access token is not scoped to this repository, or\n" + " - the repository/tag does not exist under that namespace.\n" + " Re-running `docker login` will not fix it; widen the token's " + "repository scope instead." + ) + + denied = ( + "pull access denied" in lowered + or "requested access to the resource is denied" in lowered + or "authentication required" in lowered + or "unauthorized" in lowered + ) + if not denied: + return None + + registry = _registry_from_image(image) + if registry is None: + fixes = ( + " - `docker login` on this machine (madengine reuses an existing " + "login)\n" + ' - add {"dockerhub": {"username": "...", "password": "..."}} to ' + "credential.json\n" + " - export MAD_DOCKERHUB_USER and MAD_DOCKERHUB_PASSWORD" + ) + else: + fixes = ( + f" - `docker login {registry}` on this machine (madengine reuses " + "an existing login)\n" + f' - add {{"{registry}": {{"username": "...", "password": "..."}}}} ' + "to credential.json" + ) + + return ( + f"Base image pull was denied: {subject}\n" + " No usable credentials were presented to the registry.\n" + " Fix with any one of:\n" + fixes + ) + + def login_to_registry( registry: Optional[str], credentials: Optional[Dict], @@ -103,10 +285,17 @@ def login_to_registry( failure (missing key, invalid format, or docker login error). Set to ``False`` to log and return instead, allowing the caller to fall back to pulling public images. + + Precedence: explicit credentials (``credential.json`` / ``MAD_DOCKERHUB_*``) + win when they carry a non-blank username and password. Otherwise an existing + ``docker login`` on this machine is reused and no login is attempted, so a + placeholder credential entry never overrides or breaks working ambient auth. + Set ``MAD_SKIP_DOCKER_LOGIN=1`` to always defer to ambient credentials. """ - if not credentials: + if os.environ.get("MAD_SKIP_DOCKER_LOGIN") == "1": rich_console.print( - "[yellow]No credentials provided for registry login[/yellow]" + "[yellow]MAD_SKIP_DOCKER_LOGIN=1 - using existing docker login for " + f"{registry or 'DockerHub'}[/yellow]" ) return @@ -116,8 +305,35 @@ def login_to_registry( if registry and registry.lower() == "docker.io": registry_key = "dockerhub" - if registry_key not in credentials: - error_msg = f"No credentials found for registry: {registry_key}" + entry = (credentials or {}).get(registry_key) + creds: Dict = entry if isinstance(entry, dict) else {} + + if not _usable_credentials(creds): + # No explicit credentials configured for this registry. If the machine + # is already logged in (e.g. an organisation access token), reuse that + # instead of failing or clobbering it. + if has_ambient_docker_auth(registry): + rich_console.print( + f"[green]Using existing docker login for " + f"{registry or 'DockerHub'} (no explicit credentials " + f"configured)[/green]" + ) + return + + if not credentials: + rich_console.print( + "[yellow]No credentials provided for registry login[/yellow]" + ) + return + + if registry_key not in credentials: + error_msg = f"No credentials found for registry: {registry_key}" + else: + error_msg = ( + f"Invalid credentials format for registry: {registry_key}" + f"\nCredentials must contain non-empty 'username' and " + f"'password' fields" + ) if registry_key == "dockerhub": error_msg += ( f"\nPlease add dockerhub credentials to credential.json:\n" @@ -140,18 +356,7 @@ def login_to_registry( " }\n" "}" ) - rich_console.print(f"[red]{error_msg}[/red]") - if raise_on_failure: - raise RuntimeError(error_msg) - return - - creds = credentials[registry_key] - - if "username" not in creds or "password" not in creds: - error_msg = ( - f"Invalid credentials format for registry: {registry_key}" - f"\nCredentials must contain 'username' and 'password' fields" - ) + error_msg += "\nAlternatively, run `docker login` on this machine." rich_console.print(f"[red]{error_msg}[/red]") if raise_on_failure: raise RuntimeError(error_msg) diff --git a/src/madengine/core/console.py b/src/madengine/core/console.py index 71105c2f..36e1e161 100644 --- a/src/madengine/core/console.py +++ b/src/madengine/core/console.py @@ -229,6 +229,12 @@ def sh( # Check for failure success = proc.returncode == 0 + # When output is captured rather than streamed it is discarded on + # failure, and the RuntimeError below carries only the command and the + # exit code. Echo it so the log records why the command actually failed. + if not success and not canFail and not secret and not self.live_output and outs: + print(redact_secrets(outs), flush=True) + # Show docker operation completion status if not secret: self._show_docker_completion(command, success) diff --git a/src/madengine/execution/docker_builder.py b/src/madengine/execution/docker_builder.py index bc44aa44..5ee83c47 100644 --- a/src/madengine/execution/docker_builder.py +++ b/src/madengine/execution/docker_builder.py @@ -9,6 +9,7 @@ import os import shlex +import sys from pathlib import Path import time import json @@ -16,7 +17,11 @@ import typing from contextlib import redirect_stdout, redirect_stderr from rich.console import Console as RichConsole -from madengine.core.auth import login_to_registry +from madengine.core.auth import ( + explain_registry_denial, + has_ambient_docker_auth, + login_to_registry, +) from madengine.core.console import Console from madengine.core.context import Context from madengine.utils.ops import PythonicTee @@ -98,6 +103,72 @@ def get_build_arg(self, run_build_arg: typing.Optional[typing.Dict] = None) -> s return build_args + def _resolve_base_docker(self, dockerfile: str) -> str: + """Resolve the base image the Dockerfile builds ``FROM``. + + Prefers a ``BASE_DOCKER`` override from ``docker_build_arg`` context, + otherwise reads ``ARG BASE_DOCKER=`` from the Dockerfile. + + Args: + dockerfile: Path to the Dockerfile. + + Returns: + str: The base image reference, or ``""`` if it cannot be determined. + """ + if ( + "docker_build_arg" in self.context.ctx + and "BASE_DOCKER" in self.context.ctx["docker_build_arg"] + ): + return str(self.context.ctx["docker_build_arg"]["BASE_DOCKER"]) + try: + return str( + self.console.sh( + f"grep '^ARG BASE_DOCKER=' {shlex.quote(dockerfile)} | sed -E 's/ARG BASE_DOCKER=//g'" + ) + ) + except Exception: + return "" + + @staticmethod + def _registry_of(image: str) -> str: + """Return the registry an image reference points at. + + Args: + image: An image reference such as ``rocm/pytorch:latest`` or + ``myhost:5000/team/img:tag``. + + Returns: + str: The registry host, or ``"docker.io"`` for Docker Hub references. + """ + first_segment = (image or "").strip().split("/")[0] + if "." in first_segment or ":" in first_segment or first_segment == "localhost": + return first_segment + return "docker.io" + + def _report_registry_denial(self, log_file_path: str, base_docker: str) -> None: + """Print an actionable hint if a build log shows a registry denial. + + Called from inside the build's stdout redirection, so the hint is written + to the real stdout as well to make sure it reaches the terminal and not + only the build log. + + Args: + log_file_path: Path to the build log written by this build. + base_docker: The base image the build was pulling. + """ + try: + with open(log_file_path, encoding="utf-8", errors="replace") as log: + log_text = log.read() + except OSError: + return + hint = explain_registry_denial(log_text, base_docker) + if not hint: + return + message = f"[bold red]❌ {hint}[/bold red]" + self.rich_console.print(f"\n{message}") + if not self.live_output and sys.__stdout__ is not None: + RichConsole(file=sys.__stdout__).print(f"\n{message}") + def build_image( self, model_info: typing.Dict, @@ -195,8 +266,28 @@ def build_image( with redirect_stdout( PythonicTee(outlog, self.live_output) ), redirect_stderr(PythonicTee(outlog, self.live_output)): + # `docker build --pull` resolves the base image itself, so it only + # works if this machine can authenticate to the base image's + # registry. Log in when — and only when — there is no existing + # login to reuse, so an ambient `docker login` (e.g. an + # organisation access token) is left alone. + base_docker = self._resolve_base_docker(dockerfile) + base_registry = self._registry_of(base_docker) + if credentials and not has_ambient_docker_auth(base_registry): + login_to_registry( + base_registry, + credentials, + console=self.console, + rich_console=self.rich_console, + raise_on_failure=False, + ) + print(f"🔨 Executing build command...") - self.console.sh(build_command, timeout=None) + try: + self.console.sh(build_command, timeout=None) + except Exception: + self._report_registry_denial(log_file_path, base_docker) + raise build_duration = time.time() - build_start_time @@ -206,17 +297,6 @@ def build_image( self.rich_console.print(f"[dim]{'='*80}[/dim]") # Get base docker info - base_docker = "" - if ( - "docker_build_arg" in self.context.ctx - and "BASE_DOCKER" in self.context.ctx["docker_build_arg"] - ): - base_docker = self.context.ctx["docker_build_arg"]["BASE_DOCKER"] - else: - base_docker = self.console.sh( - f"grep '^ARG BASE_DOCKER=' {shlex.quote(dockerfile)} | sed -E 's/ARG BASE_DOCKER=//g'" - ) - print(f"BASE DOCKER is {base_docker}") # Get docker SHA diff --git a/src/madengine/orchestration/build_orchestrator.py b/src/madengine/orchestration/build_orchestrator.py index 246701cf..17e836e8 100644 --- a/src/madengine/orchestration/build_orchestrator.py +++ b/src/madengine/orchestration/build_orchestrator.py @@ -20,7 +20,7 @@ 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.auth import has_ambient_docker_auth, load_credentials from madengine.core.errors import ( BuildError, ConfigurationError, @@ -975,16 +975,33 @@ def _execute_build_on_compute( ) if any(pub_reg in registry_lower for pub_reg in public_registries): if not dockerhub_user or not dockerhub_password: - raise ConfigurationError( - f"Registry credentials required for pushing to {registry}", - context=create_error_context( - operation="build_on_compute", - component="BuildOrchestrator", - additional_info={"registry": registry}, - ), - suggestions=_matched_hints, - ) - self.rich_console.print(f" Auth: Will login to registry before push") + # An existing `docker login` (e.g. an organisation access token) + # is a valid substitute for explicit credentials, so do not hard + # fail on it. The generated sbatch script already falls back to + # "assume pre-authenticated" when no credentials are supplied. + if has_ambient_docker_auth(registry.split("/")[0]): + self.rich_console.print( + " [yellow]Auth: No explicit credentials; relying on the " + "existing docker login[/yellow]" + ) + self.rich_console.print( + " [dim]Note: this login was detected on the submit node. " + "The compute node only shares it if $HOME (or $DOCKER_CONFIG) " + "is on shared storage.[/dim]" + ) + else: + raise ConfigurationError( + f"Registry credentials required for pushing to {registry}", + context=create_error_context( + operation="build_on_compute", + component="BuildOrchestrator", + additional_info={"registry": registry}, + ), + suggestions=_matched_hints + + ["Or run `docker login` on the build node"], + ) + else: + 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)") diff --git a/src/madengine/scripts/common/post_scripts/trace.sh b/src/madengine/scripts/common/post_scripts/trace.sh index 1e489861..1dbdaaf5 100644 --- a/src/madengine/scripts/common/post_scripts/trace.sh +++ b/src/madengine/scripts/common/post_scripts/trace.sh @@ -33,8 +33,12 @@ rpd) if [ -f "./rocmProfileData/tools/rpd2tracing.py" ]; then echo "RPD post-script: rpd2tracing.py found" if [ -f "trace.rpd" ] && [ -s "trace.rpd" ]; then - python3 ./rocmProfileData/tools/rpd2tracing.py trace.rpd trace.json - mv trace.rpd trace.json "$OUTPUT" + if python3 ./rocmProfileData/tools/rpd2tracing.py trace.rpd trace.json; then + mv trace.rpd trace.json "$OUTPUT" + else + echo "RPD post-script: rpd2tracing.py failed (likely no captured trace data); saving raw trace.rpd only" + mv trace.rpd "$OUTPUT" + fi else echo "RPD post-script: Skipping rpd2tracing.py because trace.rpd is missing or empty" # Create empty files so the directory structure exists diff --git a/src/madengine/scripts/common/pre_scripts/gpu_info_pre.sh b/src/madengine/scripts/common/pre_scripts/gpu_info_pre.sh index 60bd60a0..7903e378 100644 --- a/src/madengine/scripts/common/pre_scripts/gpu_info_pre.sh +++ b/src/madengine/scripts/common/pre_scripts/gpu_info_pre.sh @@ -5,15 +5,19 @@ # gpu_vendor="" -if [ -f "/usr/bin/nvidia-smi" ]; then +if command -v nvidia-smi >/dev/null 2>&1; then echo "NVIDIA GPU detected." gpu_vendor="NVIDIA" gpu_architecture=$(nvidia-smi --query-gpu=name --format=csv,noheader | grep -m 1 -E -o ".{0,1}100"| xargs ) python3 -m pip install nvidia-ml-py -elif [ -f "/opt/rocm/bin/rocm-smi" ]; then +elif command -v rocm-smi >/dev/null 2>&1 || command -v amd-smi >/dev/null 2>&1; then echo "AMD GPU detected." gpu_vendor="AMD" - gpu_architecture=$(rocminfo | grep -o -m 1 'gfx.*' | xargs ) + if command -v rocminfo >/dev/null 2>&1; then + gpu_architecture=$(rocminfo | grep -o -m 1 'gfx.*' | xargs ) + else + echo "rocminfo not found; skipping AMD GPU architecture detection." + fi MI200="gfx90a" MI100="gfx908" MI50="gfx906" diff --git a/src/madengine/scripts/common/tools.json b/src/madengine/scripts/common/tools.json index 82869087..b237c63c 100644 --- a/src/madengine/scripts/common/tools.json +++ b/src/madengine/scripts/common/tools.json @@ -7,10 +7,8 @@ "args": "rpd" } ], - "cmd": "./rocmProfileData/rpd_tracer/runTracer.sh", - "env_vars": { - "LD_LIBRARY_PATH": "./rocmProfileData/rpd_tracer:/opt/rocm/lib" - }, + "cmd": "LD_LIBRARY_PATH=\"./rocmProfileData/rpd_tracer:${ROCM_PATH:-/opt/rocm}/lib:${LD_LIBRARY_PATH}\" ./rocmProfileData/rpd_tracer/runTracer.sh", + "env_vars": {}, "post_scripts": [ { "path": "scripts/common/post_scripts/trace.sh", diff --git a/src/madengine/scripts/common/tools/amd_smi_utils.py b/src/madengine/scripts/common/tools/amd_smi_utils.py index e0e48096..ab245c8e 100644 --- a/src/madengine/scripts/common/tools/amd_smi_utils.py +++ b/src/madengine/scripts/common/tools/amd_smi_utils.py @@ -6,16 +6,18 @@ Copyright (c) Advanced Micro Devices, Inc. All rights reserved. """ +import os import sys import logging from typing import List, Optional, Dict, Any -sys.path.append("/opt/rocm/libexec/amdsmi_cli/") +_ROCM_PATH = os.environ.get("ROCM_PATH", "/opt/rocm") +sys.path.append(f"{_ROCM_PATH}/libexec/amdsmi_cli/") try: from amdsmi_init import amdsmi_interface from amdsmi_init import amdsmi_cli_init, amdsmi_cli_shutdown except ImportError: - raise ImportError("Could not import /opt/rocm/libexec/amdsmi_cli/amdsmi_init.py") + raise ImportError(f"Could not import {_ROCM_PATH}/libexec/amdsmi_cli/amdsmi_init.py") class ProfUtils: diff --git a/src/madengine/scripts/common/tools/gpu_info_profiler.py b/src/madengine/scripts/common/tools/gpu_info_profiler.py index 111f655d..14949954 100644 --- a/src/madengine/scripts/common/tools/gpu_info_profiler.py +++ b/src/madengine/scripts/common/tools/gpu_info_profiler.py @@ -14,6 +14,7 @@ import sys import csv import os +import shutil import logging import typing import signal @@ -27,10 +28,11 @@ def check_amd_smi_available() -> bool: bool: True if amd-smi is available, False otherwise. """ # First check for Python bindings (more reliable for programmatic access) + rocm_path = os.environ.get("ROCM_PATH", "/opt/rocm") try: - sys.path.append("/opt/rocm/libexec/amdsmi_cli/") + sys.path.append(f"{rocm_path}/libexec/amdsmi_cli/") from amdsmi_init import amdsmi_interface - logging.debug("amd-smi Python bindings found at /opt/rocm/libexec/amdsmi_cli/") + logging.debug(f"amd-smi Python bindings found at {rocm_path}/libexec/amdsmi_cli/") return True except ImportError: logging.debug("amd-smi Python bindings not found") @@ -74,9 +76,11 @@ def get_rocm_version() -> Optional[float]: logging.debug(f"hipconfig check failed: {e}") try: - # Fallback to /opt/rocm/.info/version - if os.path.exists("/opt/rocm/.info/version"): - result = subprocess.run(['cat', '/opt/rocm/.info/version'], + # Fallback to $ROCM_PATH/.info/version + rocm_path = os.environ.get("ROCM_PATH", "/opt/rocm") + version_file = f"{rocm_path}/.info/version" + if os.path.exists(version_file): + result = subprocess.run(['cat', version_file], capture_output=True, text=True, timeout=10) if result.returncode == 0: version_str = result.stdout.strip().split('-')[0] # Remove build suffix @@ -96,15 +100,21 @@ def detect_gpu_vendor() -> tuple[bool, bool]: Raises: ValueError: If no GPU management tools are found. """ - if os.path.exists("/usr/bin/nvidia-smi"): + rocm_path = os.environ.get("ROCM_PATH", "/opt/rocm") + if shutil.which("nvidia-smi"): return True, False - elif os.path.exists("/opt/rocm/bin/rocm-smi") or check_amd_smi_available(): + elif ( + shutil.which("rocm-smi") + or os.path.exists(f"{rocm_path}/bin/rocm-smi") + or check_amd_smi_available() + ): return False, True else: error_msg = ( "Unable to detect GPU vendor. No GPU management tools found.\n" - "For NVIDIA: /usr/bin/nvidia-smi not found\n" - "For AMD: /opt/rocm/bin/rocm-smi and amd-smi not found\n\n" + "For NVIDIA: nvidia-smi not found on PATH\n" + f"For AMD: rocm-smi not found on PATH or in {rocm_path}/bin, " + "and amd-smi not available\n\n" "Please ensure:\n" " 1. GPU drivers are installed\n" " 2. For AMD GPUs: ROCm is properly installed (https://rocm.docs.amd.com)\n" diff --git a/src/madengine/scripts/common/tools/rocm_smi_utils.py b/src/madengine/scripts/common/tools/rocm_smi_utils.py index dd73219b..0ef5772f 100644 --- a/src/madengine/scripts/common/tools/rocm_smi_utils.py +++ b/src/madengine/scripts/common/tools/rocm_smi_utils.py @@ -6,16 +6,18 @@ Copyright (c) Advanced Micro Devices, Inc. All rights reserved. """ +import os import sys import logging from typing import List -sys.path.append("/opt/rocm/libexec/rocm_smi/") +_ROCM_PATH = os.environ.get("ROCM_PATH", "/opt/rocm") +sys.path.append(f"{_ROCM_PATH}/libexec/rocm_smi/") try: import rocm_smi from rsmiBindings import * except ImportError: - raise ImportError("Could not import /opt/rocm/libexec/rocm_smi/rocm_smi.py") + raise ImportError(f"Could not import {_ROCM_PATH}/libexec/rocm_smi/rocm_smi.py") class ProfUtils: diff --git a/tests/unit/test_auth.py b/tests/unit/test_auth.py index 47d4a6c7..d2bf2fd8 100644 --- a/tests/unit/test_auth.py +++ b/tests/unit/test_auth.py @@ -1,10 +1,16 @@ """Unit tests for madengine.core.auth module.""" +import json import os import pytest from unittest.mock import MagicMock, mock_open, patch -from madengine.core.auth import load_credentials, login_to_registry +from madengine.core.auth import ( + explain_registry_denial, + has_ambient_docker_auth, + load_credentials, + login_to_registry, +) class TestLoadCredentials: @@ -115,21 +121,23 @@ def test_load_credentials_non_dockerhub_registry(self, mock_file, mock_exists): assert result["custom_registry"]["token"] == "abc123" +@patch.dict(os.environ, {"MAD_SKIP_DOCKER_LOGIN": ""}, clear=False) +@patch("madengine.core.auth.has_ambient_docker_auth", return_value=False) class TestLoginToRegistry: - """Tests for login_to_registry().""" + """Tests for login_to_registry() when the machine has no existing docker login.""" def _mocks(self): console = MagicMock() rich_console = MagicMock() return console, rich_console - def test_no_credentials_returns_early(self): + def test_no_credentials_returns_early(self, mock_ambient): """Passing None credentials logs a warning and returns without error.""" console, rich_console = self._mocks() login_to_registry("docker.io", None, console, rich_console) console.sh.assert_not_called() - def test_missing_registry_key_raises_when_raise_on_failure(self): + def test_missing_registry_key_raises_when_raise_on_failure(self, mock_ambient): """RuntimeError raised when registry key absent and raise_on_failure=True.""" console, rich_console = self._mocks() credentials = {"other_registry": {"username": "u", "password": "p"}} @@ -137,14 +145,14 @@ def test_missing_registry_key_raises_when_raise_on_failure(self): login_to_registry("myregistry.io", credentials, console, rich_console, raise_on_failure=True) console.sh.assert_not_called() - def test_missing_registry_key_returns_when_not_raise_on_failure(self): + def test_missing_registry_key_returns_when_not_raise_on_failure(self, mock_ambient): """Returns silently when registry key absent and raise_on_failure=False.""" console, rich_console = self._mocks() credentials = {"other_registry": {"username": "u", "password": "p"}} login_to_registry("myregistry.io", credentials, console, rich_console, raise_on_failure=False) console.sh.assert_not_called() - def test_invalid_credentials_format_raises(self): + def test_invalid_credentials_format_raises(self, mock_ambient): """RuntimeError raised when username/password fields missing.""" console, rich_console = self._mocks() credentials = {"dockerhub": {"token": "abc"}} @@ -152,14 +160,22 @@ def test_invalid_credentials_format_raises(self): login_to_registry("docker.io", credentials, console, rich_console, raise_on_failure=True) console.sh.assert_not_called() - def test_invalid_credentials_format_returns_when_not_raise_on_failure(self): + def test_invalid_credentials_format_returns_when_not_raise_on_failure(self, mock_ambient): """Returns silently when credentials format invalid and raise_on_failure=False.""" console, rich_console = self._mocks() credentials = {"dockerhub": {"token": "abc"}} login_to_registry("docker.io", credentials, console, rich_console, raise_on_failure=False) console.sh.assert_not_called() - def test_docker_io_normalised_to_dockerhub(self): + def test_blank_credentials_raise_without_ambient_auth(self, mock_ambient): + """Placeholder credentials are treated as absent, not as credentials.""" + console, rich_console = self._mocks() + credentials = {"dockerhub": {"repository": "r", "username": "", "password": ""}} + with pytest.raises(RuntimeError, match="username|password"): + login_to_registry("docker.io", credentials, console, rich_console, raise_on_failure=True) + console.sh.assert_not_called() + + def test_docker_io_normalised_to_dockerhub(self, mock_ambient): """docker.io registry is looked up under the 'dockerhub' key.""" console, rich_console = self._mocks() credentials = {"dockerhub": {"username": "user", "password": "pass"}} @@ -169,7 +185,7 @@ def test_docker_io_normalised_to_dockerhub(self): # docker.io should not appear in the login command (uses default DockerHub endpoint) assert "docker.io" not in cmd - def test_custom_registry_included_in_command(self): + def test_custom_registry_included_in_command(self, mock_ambient): """Non-DockerHub registry URL is included in the login command.""" console, rich_console = self._mocks() credentials = {"myregistry.io": {"username": "user", "password": "pass"}} @@ -178,7 +194,7 @@ def test_custom_registry_included_in_command(self): cmd = console.sh.call_args[0][0] assert "myregistry.io" in cmd - def test_login_failure_raises_when_raise_on_failure(self): + def test_login_failure_raises_when_raise_on_failure(self, mock_ambient): """docker login error is re-raised when raise_on_failure=True.""" console, rich_console = self._mocks() console.sh.side_effect = RuntimeError("auth failed") @@ -186,10 +202,182 @@ def test_login_failure_raises_when_raise_on_failure(self): with pytest.raises(RuntimeError, match="auth failed"): login_to_registry(None, credentials, console, rich_console, raise_on_failure=True) - def test_login_failure_suppressed_when_not_raise_on_failure(self): + def test_login_failure_suppressed_when_not_raise_on_failure(self, mock_ambient): """docker login error is suppressed when raise_on_failure=False.""" console, rich_console = self._mocks() console.sh.side_effect = RuntimeError("auth failed") credentials = {"dockerhub": {"username": "user", "password": "pass"}} login_to_registry(None, credentials, console, rich_console, raise_on_failure=False) # Should not propagate the exception + + +@patch.dict(os.environ, {"MAD_SKIP_DOCKER_LOGIN": ""}, clear=False) +class TestLoginToRegistryWithAmbientAuth: + """Tests for login_to_registry() when the machine already has a docker login.""" + + def _mocks(self): + return MagicMock(), MagicMock() + + @patch("madengine.core.auth.has_ambient_docker_auth", return_value=True) + def test_blank_credentials_defer_to_ambient_auth(self, mock_ambient): + """Blank credentials never override or break an existing docker login.""" + console, rich_console = self._mocks() + credentials = {"dockerhub": {"repository": "r", "username": "", "password": ""}} + # No raise even with raise_on_failure=True: the machine is authenticated. + login_to_registry("docker.io", credentials, console, rich_console, raise_on_failure=True) + console.sh.assert_not_called() + + @patch("madengine.core.auth.has_ambient_docker_auth", return_value=True) + def test_missing_registry_key_defers_to_ambient_auth(self, mock_ambient): + """A registry with no credential.json entry falls back to the existing login.""" + console, rich_console = self._mocks() + credentials = {"other_registry": {"username": "u", "password": "p"}} + login_to_registry("myregistry.io", credentials, console, rich_console, raise_on_failure=True) + console.sh.assert_not_called() + + @patch("madengine.core.auth.has_ambient_docker_auth", return_value=True) + def test_explicit_credentials_win_over_ambient_auth(self, mock_ambient): + """Usable explicit credentials still trigger a login (explicit wins).""" + console, rich_console = self._mocks() + credentials = {"dockerhub": {"username": "user", "password": "pass"}} + login_to_registry("docker.io", credentials, console, rich_console) + console.sh.assert_called_once() + assert "--username user" in console.sh.call_args[0][0] + + @patch("madengine.core.auth.has_ambient_docker_auth", return_value=False) + def test_whitespace_only_credentials_are_not_credentials(self, mock_ambient): + """Whitespace-only values are treated as blank.""" + console, rich_console = self._mocks() + credentials = {"dockerhub": {"username": " ", "password": "\t"}} + with pytest.raises(RuntimeError, match="username|password"): + login_to_registry("docker.io", credentials, console, rich_console, raise_on_failure=True) + console.sh.assert_not_called() + + +class TestSkipDockerLogin: + """Tests for the MAD_SKIP_DOCKER_LOGIN escape hatch.""" + + @patch.dict(os.environ, {"MAD_SKIP_DOCKER_LOGIN": "1"}, clear=False) + def test_skip_env_var_bypasses_login(self): + """MAD_SKIP_DOCKER_LOGIN=1 defers to ambient credentials unconditionally.""" + console, rich_console = MagicMock(), MagicMock() + credentials = {"dockerhub": {"username": "user", "password": "pass"}} + login_to_registry("docker.io", credentials, console, rich_console, raise_on_failure=True) + console.sh.assert_not_called() + + +class TestHasAmbientDockerAuth: + """Tests for has_ambient_docker_auth().""" + + def _write_config(self, tmp_path, config): + (tmp_path / "config.json").write_text(json.dumps(config), encoding="utf-8") + return {"DOCKER_CONFIG": str(tmp_path)} + + def test_dockerhub_auth_entry_detected(self, tmp_path): + """A Docker Hub entry with an auth blob counts as authenticated.""" + env = self._write_config( + tmp_path, {"auths": {"https://index.docker.io/v1/": {"auth": "abc123"}}} + ) + with patch.dict(os.environ, env, clear=False): + assert has_ambient_docker_auth(None) is True + assert has_ambient_docker_auth("docker.io") is True + assert has_ambient_docker_auth("docker.io/rocm/mad-private") is True + assert has_ambient_docker_auth("myregistry.io") is False + + def test_identity_token_entry_detected(self, tmp_path): + """An identitytoken-only entry counts as authenticated.""" + env = self._write_config( + tmp_path, {"auths": {"index.docker.io": {"identitytoken": "tok"}}} + ) + with patch.dict(os.environ, env, clear=False): + assert has_ambient_docker_auth("docker.io") is True + + def test_cred_helper_detected(self, tmp_path): + """A credHelpers entry counts as authenticated.""" + env = self._write_config(tmp_path, {"credHelpers": {"myregistry.io": "ecr-login"}}) + with patch.dict(os.environ, env, clear=False): + assert has_ambient_docker_auth("myregistry.io/team/img") is True + assert has_ambient_docker_auth("docker.io") is False + + def test_creds_store_with_empty_auth_entry(self, tmp_path): + """credsStore-managed entries are stored empty but are still credentials.""" + env = self._write_config( + tmp_path, + {"credsStore": "desktop", "auths": {"https://index.docker.io/v1/": {}}}, + ) + with patch.dict(os.environ, env, clear=False): + assert has_ambient_docker_auth("docker.io") is True + assert has_ambient_docker_auth("other.io") is False + + def test_empty_auth_entry_without_creds_store(self, tmp_path): + """An empty entry with no credential store is not usable.""" + env = self._write_config(tmp_path, {"auths": {"https://index.docker.io/v1/": {}}}) + with patch.dict(os.environ, env, clear=False): + assert has_ambient_docker_auth("docker.io") is False + + def test_registry_with_port_matches_host(self, tmp_path): + """host:port registries are matched on the host segment.""" + env = self._write_config(tmp_path, {"auths": {"localhost:5000": {"auth": "x"}}}) + with patch.dict(os.environ, env, clear=False): + assert has_ambient_docker_auth("localhost:5000/team/img") is True + + def test_missing_config_returns_false(self, tmp_path): + """A missing config.json yields False rather than raising.""" + with patch.dict(os.environ, {"DOCKER_CONFIG": str(tmp_path)}, clear=False): + assert has_ambient_docker_auth("docker.io") is False + + def test_corrupt_config_returns_false(self, tmp_path): + """A corrupt config.json yields False rather than raising.""" + (tmp_path / "config.json").write_text("not json{{{", encoding="utf-8") + with patch.dict(os.environ, {"DOCKER_CONFIG": str(tmp_path)}, clear=False): + assert has_ambient_docker_auth("docker.io") is False + + +class TestExplainRegistryDenial: + """Tests for explain_registry_denial().""" + + def test_insufficient_scope_is_an_authorization_problem(self): + """insufficient_scope is reported as authorization, not authentication.""" + log = ( + "#6 ERROR: pull access denied, repository does not exist or may require " + "authorization: server message: insufficient_scope: authorization failed" + ) + hint = explain_registry_denial(log, "rocm/triton-inference-server-dev:x") + assert hint is not None + assert "authorization problem" in hint + assert "rocm/triton-inference-server-dev:x" in hint + assert "will not fix it" in hint + + def test_plain_denial_points_at_login(self): + """A denial without insufficient_scope points at supplying credentials.""" + log = "Error response from daemon: pull access denied for rocm/private" + hint = explain_registry_denial(log, "rocm/private:latest") + assert hint is not None + assert "docker login" in hint + assert "authorization problem" not in hint + + def test_non_dockerhub_denial_names_that_registry(self): + """A denial for another registry does not suggest Docker Hub credentials.""" + log = "Error response from daemon: pull access denied for ghcr.io/org/app" + hint = explain_registry_denial(log, "ghcr.io/org/app:latest") + assert hint is not None + assert "docker login ghcr.io" in hint + assert '"ghcr.io": {"username"' in hint + assert "dockerhub" not in hint + assert "MAD_DOCKERHUB" not in hint + + def test_dockerhub_namespace_still_suggests_dockerhub(self): + """A bare namespace/repo reference is Docker Hub, not a registry host.""" + log = "Error response from daemon: pull access denied for rocm/private" + hint = explain_registry_denial(log, "rocm/private:latest") + assert hint is not None + assert '"dockerhub"' in hint + assert "MAD_DOCKERHUB_USER" in hint + + def test_unrelated_failure_returns_none(self): + """Non-registry build failures produce no hint.""" + assert explain_registry_denial("RUN apt-get install failed: exit code 100") is None + + def test_empty_log_returns_none(self): + """Empty output produces no hint.""" + assert explain_registry_denial("") is None