diff --git a/Dockerfile.rocm b/Dockerfile.rocm index 339c4ecd919..b9c8a8c011a 100644 --- a/Dockerfile.rocm +++ b/Dockerfile.rocm @@ -79,17 +79,21 @@ COPY LICENSE CHANGELOG.md CODE_OF_CONDUCT.md CONTRIBUTING.md README.md versionee COPY tests ./tests COPY monai ./monai -# Use print_dependencies.py rather than -e .[all,testing] to filter CUDA-only packages: -# cucim-cu* pulls in cuda-toolkit (~1.2 GB); nvidia-ml-py fails at import on ROCm; nni depends on it. +# Set HIPCIM_INDEX_URL="" to build without WSI/cucim support. +ARG HIPCIM_INDEX_URL="https://pypi.amd.com/rocm-${ROCM_SERIES}.0/simple/" + +# print_dependencies.py resolves the ROCm variants of the GPU imaging stack (amd-hipcim, amd-cupy) +# and drops the NVIDIA-only packages, see monai/config/vendor_deps.py. # BUILD_MONAI=1 builds the C++/HIP extensions ahead of time; FORCE_CUDA=1 is required because the # build host has no GPU, so setup.py's `torch.cuda.is_available()` check would otherwise skip them. # Compilation itself needs only the toolkit and PYTORCH_ROCM_ARCH (set above), not a device. # pytest is not in the "testing" extra; it is added explicitly for running tests by hand. RUN python monai/config/print_dependencies.py build-system \ | xargs -d '\n' pip install --no-cache-dir --no-build-isolation \ - && python monai/config/print_dependencies.py all testing \ - | grep -vE '^cucim-cu|^nvidia-ml-py|^nni' > /tmp/rocm-requirements-$$.txt \ + && python monai/config/print_dependencies.py all testing > /tmp/rocm-requirements-$$.txt \ + && if [ -z "${HIPCIM_INDEX_URL}" ]; then sed -i '/^amd-hipcim/d' /tmp/rocm-requirements-$$.txt; fi \ && BUILD_MONAI=1 FORCE_CUDA=1 pip install --no-cache-dir --no-build-isolation \ + ${HIPCIM_INDEX_URL:+--extra-index-url "${HIPCIM_INDEX_URL}"} \ -r /tmp/rocm-requirements-$$.txt pytest -e . \ && rm -f /tmp/rocm-requirements-$$.txt @@ -97,15 +101,12 @@ RUN python monai/config/print_dependencies.py build-system \ # without it the extensions built above would be present but never used. ENV BUILD_MONAI=1 -# Set HIPCIM_INDEX_URL="" to build without WSI/cucim support. # CuImage is imported (not just cucim) because cucim uses lazy_loader -- a bare import # succeeds even when the native library is unresolvable. Failing here is deliberate: if # hipCIM was requested, an image where the cucim backends silently do not work is worse # than no image at all. -ARG HIPCIM_INDEX_URL="https://pypi.amd.com/rocm-${ROCM_SERIES}.0/simple/" RUN if [ -n "${HIPCIM_INDEX_URL}" ]; then \ - pip install --no-cache-dir --extra-index-url "${HIPCIM_INDEX_URL}" "amd-hipcim" \ - && python -c "from cucim import CuImage"; \ + python -c "from cucim import CuImage"; \ else \ echo "hipCIM not installed; whole-slide-image (cucim) backends are unavailable."; \ fi diff --git a/docs/source/installation.md b/docs/source/installation.md index 31c99a794c1..c83e3aef680 100644 --- a/docs/source/installation.md +++ b/docs/source/installation.md @@ -89,6 +89,62 @@ The `nvimgcodec` extra installs GPU-accelerated DICOM decoding for `NvImgCodecPy (`pip install 'monai[nvimgcodec]'`). It is Linux-only in the extra definition; CUDA 13 is the default. CUDA 12 users should install matching `cupy-cuda12x` and `nvidia-nvimgcodec-cu12` wheels. +On AMD ROCm the extras above keep their names. When MONAI is built or installed against a ROCm +build of PyTorch, the CUDA-only distributions they pull in are replaced by their AMD counterparts: +`cucim-cu12`/`cucim-cu13` become [`amd-hipcim`](https://rocm.docs.amd.com/projects/hipCIM/en/latest/) +(hipCIM), `cupy-cuda*` becomes `amd-cupy`, and `nvidia-ml-py`, `nni` and `nvidia-nvimgcodec-cu*` are +dropped. `amd-hipcim` ships the `cucim` Python namespace and `amd-cupy` ships `cupy`, so all existing +MONAI code importing from either (e.g. `WSIReader` with `backend="cucim"`, or `convert_to_cupy`) +works on ROCm without any code changes. The required dependencies additionally gain the AMD GPU +device extras (`torch[device-gfx942,device-gfx950]`) and a `rocm[libraries,devel,device-gfx*]` +requirement, which is what steers pip towards the ROCm build of PyTorch. `devel` is included because +`monai/_extensions` JIT-compiles its HIP sources on first use and the link step needs the ROCm +development tree. + +This is driven by a general accelerator-vendor mechanism: `monai/config/vendor_deps.py` holds the +generic rewriting and a registry of vendors, and each vendor contributes a `vendor_.py` plugin +listing its package substitutions. A plugin is imported only once its registry probe matches, so a +build for one vendor never runs another vendor's code, and a build with no vendor detected is left +exactly as declared in `pyproject.toml`. Set `MONAI_VENDOR` to a registered vendor name to force one, +or to `none` to disable rewriting. For ROCm, the GPU architectures come from `GPU_TARGETS` or +`AMDGPU_TARGETS` and the ROCm series from `MONAI_ROCM_SERIES`. + +Two AMD package indexes are involved, and neither is on PyPI. The ROCm SDK (`rocm`) and the ROCm +build of PyTorch come from the ROCm wheel index. `amd-hipcim` and `amd-cupy` come from the AMD +extensions index, which is published per ROCm series — at the time of writing, `rocm-10.0.0`. Those +wheels are supported on the series they were built for and the one after it, which is the range the +generated `rocm` requirement allows. + +Install the ROCm runtime and PyTorch first, selecting the architecture of your GPU (`gfx942` for +MI300X/MI325X, `gfx950` for MI350X/MI355X), then expand the development tree: + +```bash +pip install --index-url https://stable.repo.amd.com/rocm/whl-next/ \ + "rocm[libraries,devel,device-gfx942]" "torch[device-gfx942]" +rocm-sdk init +``` + +`rocm-sdk init` is required: the `rocm-sdk-devel` wheel ships its contents as an archive that this +command expands. Without it the ROCm development tree is absent, and `monai/_extensions` fails to +JIT-compile on first use because the link step cannot find `libamdhip64.so`. + +Then install MONAI with both indexes available, so the `rocm` requirement and the AMD imaging +packages can both resolve: + +```bash +pip install \ + --extra-index-url https://stable.repo.amd.com/rocm/whl-next/ \ + --extra-index-url https://pypi.amd.com/rocm-10.0.0/simple/ \ + 'monai[cucim]' +``` + +AMD currently publishes the imaging wheels for CPython 3.12 on x86_64 Linux only; on other +interpreters or architectures the install fails with `No matching distribution found for +amd-hipcim`. + +`Dockerfile.rocm` in the repo root is the authoritative reference for the complete environment setup +(library paths, compiler flags, `rocm-sdk init`, and the full container recipe). + The `hyena` extra pulls in [`nvsubquadratic`](https://github.com/NVIDIA-BioNeMo/nvSubquadratic), required by `HyenaNDUNETR` / `HyenaMixer` / `HyenaTransformerBlock` (subquadratic O(N log N) alternatives to windowed self-attention). Install with @@ -333,7 +389,7 @@ MONAI itself: ```bash git clone https://github.com/Project-MONAI/MONAI.git cd MONAI/ -python monai/config/print_dependencies.py \* > requirements.txt +python monai/config/print_dependencies.py all testing > requirements.txt pip install -r requirements.txt ``` diff --git a/monai/__init__.py b/monai/__init__.py index 2d6f5dfc37b..dfad1731a92 100644 --- a/monai/__init__.py +++ b/monai/__init__.py @@ -94,6 +94,9 @@ def filter(self, record): "(.*(__main__)$)", "(.*(video_dataset)$)", "(.*(nnunet).*$)", + # packaging-time only, and a vendor plugin must not be imported on another vendor's hardware + "(^(monai.config.vendor_))", + "(^(monai.config.print_dependencies)$)", ] ) diff --git a/monai/config/check_env.py b/monai/config/check_env.py index f00c5b8fcd3..372f0834720 100644 --- a/monai/config/check_env.py +++ b/monai/config/check_env.py @@ -123,7 +123,11 @@ def check_torch_cuda(): """ import torch - fprint("CUDA version:", torch.version.cuda) + # `torch.version.cuda` is None on a ROCm build, where the toolkit version is `torch.version.hip`. + if torch.version.hip: + fprint("HIP version:", torch.version.hip) + else: + fprint("CUDA version:", torch.version.cuda) try: dcount = torch.cuda.device_count() diff --git a/monai/config/deviceconfig.py b/monai/config/deviceconfig.py index 6a0d2378aa7..7ae1cbceeca 100644 --- a/monai/config/deviceconfig.py +++ b/monai/config/deviceconfig.py @@ -214,7 +214,12 @@ def get_gpu_info() -> OrderedDict: _dict_append(output, "Has CUDA", lambda: bool(torch.cuda.is_available())) if output["Has CUDA"]: - _dict_append(output, "CUDA version", lambda: torch.version.cuda) + # On a ROCm build `torch.version.cuda` is None and `torch.version.hip` carries the toolkit + # version. Reporting "CUDA version: None" is misleading, so the key name switches with the build. + if torch.version.hip: + _dict_append(output, "HIP version", lambda: torch.version.hip) + else: + _dict_append(output, "CUDA version", lambda: torch.version.cuda) cudnn_ver = torch.backends.cudnn.version() _dict_append(output, "cuDNN enabled", lambda: bool(cudnn_ver)) _dict_append(output, "NVIDIA_TF32_OVERRIDE", os.environ.get("NVIDIA_TF32_OVERRIDE")) diff --git a/monai/config/print_dependencies.py b/monai/config/print_dependencies.py index a099949ecac..e7e0bfeb8b9 100644 --- a/monai/config/print_dependencies.py +++ b/monai/config/print_dependencies.py @@ -14,6 +14,9 @@ be piped to a requirements file to work with pip. All required dependencies are always printed, those for builing are included in "build-system" is given as an argument, and all optional requirements are included if "*" is given. This assumes the pyproject.toml file is in the current working directory. + +On an accelerator-vendor PyTorch the NVIDIA-only distributions are swapped for that vendor's +equivalents, see ``monai.config.vendor_deps``. """ from __future__ import annotations @@ -21,6 +24,11 @@ import sys from collections.abc import Collection +try: + from .vendor_deps import active_vendor +except ImportError: # run as a script rather than imported from the package + from vendor_deps import active_vendor # type: ignore[no-redef] + BUILD_SYSTEM_KEY = "build-system" PROJ_KEY = "project" OPTS_KEY = "optional-dependencies" @@ -57,10 +65,16 @@ def parse_dependencies(filename: str | None = None, sections: Collection[str] | opts = proj[OPTS_KEY] dependencies = list(proj[DEP_KEY]) sections = set(sections or []) + vendor = active_vendor() + + if vendor is not None: + dependencies = vendor.apply_to_dependencies(dependencies) + opts = vendor.apply_to_optional_dependencies(opts) if BUILD_SYSTEM_KEY in sections: sections.remove(BUILD_SYSTEM_KEY) - dependencies += data[BUILD_SYSTEM_KEY][REQ_KEY] + build_requires = data[BUILD_SYSTEM_KEY][REQ_KEY] + dependencies += vendor.apply_to_dependencies(build_requires) if vendor else build_requires if "*" in sections: dependencies += sum(opts.values(), []) diff --git a/monai/config/vendor_deps.py b/monai/config/vendor_deps.py new file mode 100644 index 00000000000..d35137d1121 --- /dev/null +++ b/monai/config/vendor_deps.py @@ -0,0 +1,210 @@ +# Copyright (c) MONAI Consortium +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# http://www.apache.org/licenses/LICENSE-2.0 +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +""" +Accelerator-vendor rewriting of the dependency lists declared in ``pyproject.toml``. + +Those lists name the packages published for NVIDIA hardware (``cucim-cu*``, ``cupy-cuda*``, +``nvidia-*``). Other vendors publish equivalents under different distribution names, so a wheel +built against, say, a ROCm PyTorch needs different metadata from the same source tree. + +This module owns the generic machinery -- requirement parsing, substitution, de-duplication -- and a +registry of vendor plugins. It contains no vendor-specific knowledge beyond a one-line detection +probe per entry in :data:`REGISTRY`. + +A vendor plugin is a sibling module exposing: + +``NAME`` + Identifier used by the ``MONAI_VENDOR`` environment variable. +``SUBSTITUTIONS`` + Mapping of canonical distribution name to its replacement requirement, or to None to drop it. +``extras_for(name)`` + Extras to add to the requirement for the distribution ``name``, e.g. device selectors. +``extra_requirements()`` + Requirements to append to the required dependencies, e.g. a runtime version pin. + +Plugin modules are imported only once their registry probe matches, so a build for one vendor never +executes another vendor's code. + +``MONAI_VENDOR`` selects a vendor explicitly: a registered name, ``none`` to disable rewriting, or +``auto`` (the default) to detect. +""" + +from __future__ import annotations + +import importlib +import importlib.util +import os +import re +import sys +from collections.abc import Callable +from dataclasses import dataclass +from typing import Any + +__all__ = [ + "REGISTRY", + "ActiveVendor", + "VendorPluginError", + "active_vendor", + "apply_to_dependencies", + "apply_to_optional_dependencies", + "canonical_name", + "load_by_path", + "split_requirement", +] + + +class VendorPluginError(RuntimeError): + """A vendor was detected but its plugin could not be loaded. + + Callers must not treat this as "no vendor": the hardware is known, so falling back to the + unmodified lists would ship another vendor's packages in this vendor's wheel. + """ + + +def load_by_path(name: str, path: str) -> Any: + """Import the module at ``path`` under ``name``, without importing its package. + + It is registered in ``sys.modules`` before execution because ``dataclasses`` resolves a class's + module through it. + """ + spec = importlib.util.spec_from_file_location(name, path) + if spec is None or spec.loader is None: + raise ImportError(f"cannot load {name} from {path}") + module = importlib.util.module_from_spec(spec) + sys.modules[name] = module + spec.loader.exec_module(module) + return module + + +def _torch_version_attr(attr: str) -> str: + """``torch.version.`` as a string, or "" if torch is absent or has no such attribute.""" + try: + import torch + + return getattr(torch.version, attr, None) or "" + except Exception: # torch absent, or present but not importable in this environment + return "" + + +@dataclass(frozen=True) +class _Registration: + """A vendor plugin and the cheap probe deciding whether to load it.""" + + module: str + detect: Callable[[], bool] + + +# Probes must stay trivial and must not import the plugin module: they run on every build, including +# builds for other vendors. Add new vendors here and in the matching ``vendor_.py``. +REGISTRY: dict[str, _Registration] = {"rocm": _Registration("vendor_rocm", lambda: bool(_torch_version_attr("hip")))} + +# name, optional "[extras]", then everything else (version specifiers, markers, URLs). +_REQUIREMENT = re.compile(r"\s*(?P[A-Za-z0-9][A-Za-z0-9._-]*)\s*(?:\[(?P[^]]*)\])?(?P.*)", re.S) + + +def canonical_name(name: str) -> str: + """PEP 503 normalised distribution name.""" + return re.sub(r"[-_.]+", "-", name).lower() + + +def split_requirement(requirement: str) -> tuple[str, list[str], str]: + """Split ``requirement`` into its name, its extras and the remaining specifier/marker text.""" + match = _REQUIREMENT.match(requirement) + if match is None: + return "", [], requirement + extras = [e.strip() for e in (match["extras"] or "").split(",") if e.strip()] + return match["name"], extras, match["rest"] + + +def _import_plugin(module: str) -> Any: + """Import a sibling plugin module, whether this module was imported, run as a script or loaded + by path from ``setup.py``. + """ + if __package__: + return importlib.import_module(f".{module}", __package__) + path = os.path.join(os.path.dirname(os.path.abspath(__file__)), f"{module}.py") + return load_by_path(f"monai_{module}", path) + + +def _substitute(requirements: list[str], substitutions: dict[str, str | None]) -> list[str]: + """Apply ``substitutions``, dropping None entries and de-duplicating the result.""" + out: list[str] = [] + for requirement in requirements: + name, _, _ = split_requirement(requirement) + key = canonical_name(name) + if key in substitutions: + replacement = substitutions[key] + if replacement is None: + continue + requirement = replacement + if requirement not in out: + out.append(requirement) + return out + + +def apply_to_dependencies(dependencies: list[str], plugin: Any) -> list[str]: + """Rewrite required dependencies: substitute distributions, add the extras the plugin asks for + and append its extra requirements. + """ + out = [] + for requirement in _substitute(dependencies, plugin.SUBSTITUTIONS): + name, extras, rest = split_requirement(requirement) + added = plugin.extras_for(canonical_name(name)) + if added: + requirement = f"{name}[{','.join(sorted(set(extras) | set(added)))}]{rest}" + out.append(requirement) + for requirement in plugin.extra_requirements(): + if requirement not in out: + out.append(requirement) + return out + + +def apply_to_optional_dependencies(optional_dependencies: dict[str, list[str]], plugin: Any) -> dict[str, list[str]]: + """Rewrite each optional-dependency group, substituting distributions.""" + return {name: _substitute(list(group), plugin.SUBSTITUTIONS) for name, group in optional_dependencies.items()} + + +@dataclass(frozen=True) +class ActiveVendor: + """The vendor plugin selected for this build, with the generic rewrites bound to it.""" + + name: str + plugin: Any + + def apply_to_dependencies(self, dependencies: list[str]) -> list[str]: + return apply_to_dependencies(dependencies, self.plugin) + + def apply_to_optional_dependencies(self, optional_dependencies: dict[str, list[str]]) -> dict[str, list[str]]: + return apply_to_optional_dependencies(optional_dependencies, self.plugin) + + +def _load(name: str) -> ActiveVendor: + """Import the plugin for ``name``; a failure here is fatal, never a silent "no vendor".""" + try: + return ActiveVendor(name, _import_plugin(REGISTRY[name].module)) + except Exception as exc: + raise VendorPluginError(f"{name} was detected but its plugin could not be loaded: {exc}") from exc + + +def active_vendor() -> ActiveVendor | None: + """The vendor whose packages this build needs, or None to leave ``pyproject.toml`` alone.""" + requested = os.environ.get("MONAI_VENDOR", "").strip().lower() + if requested in ("none", "cpu"): + return None + if requested and requested != "auto": + if requested not in REGISTRY: + raise ValueError(f"unknown MONAI_VENDOR {requested!r}; expected one of {sorted(REGISTRY)} or 'none'") + return _load(requested) + for name, registration in REGISTRY.items(): + if registration.detect(): + return _load(name) + return None diff --git a/monai/config/vendor_rocm.py b/monai/config/vendor_rocm.py new file mode 100644 index 00000000000..c76cc00538a --- /dev/null +++ b/monai/config/vendor_rocm.py @@ -0,0 +1,181 @@ +# Copyright (c) MONAI Consortium +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# http://www.apache.org/licenses/LICENSE-2.0 +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +""" +AMD ROCm plugin for :mod:`monai.config.vendor_deps`. + +AMD publishes drop-in replacements for the NVIDIA imaging stack under different distribution names: +``amd-hipcim`` ships the ``cucim`` Python namespace and ``amd-cupy`` ships ``cupy``, so MONAI's +``optional_import("cucim"...)`` / ``optional_import("cupy")`` call sites work unchanged once they +are installed. + +Required dependencies additionally get the AMD GPU device extras (``torch[device-gfx942]``) and a +``rocm`` SDK requirement, which is what steers pip towards the ROCm build of PyTorch. + +Environment variables: + +* ``MONAI_ROCM_SERIES``: ROCm release ``MAJOR.MINOR`` to pin, bypassing autodetection. +* ``GPU_TARGETS`` / ``AMDGPU_TARGETS``: semicolon- or comma-separated AMD GPU architectures. + +This module is imported only when :data:`monai.config.vendor_deps.REGISTRY` detects ROCm, so it +never runs during a build for another vendor. +""" + +from __future__ import annotations + +import os +import re +import subprocess +from pathlib import Path + +NAME = "rocm" + +# ROCm release series (MAJOR.MINOR) to pin against when it cannot be detected. This tracks the ROCm +# *release* -- the version of the `rocm` pip package, e.g. 10.0.x -- NOT the HIP version, which is a +# separate internal number (ROCm 10.0 ships HIP 7.15, ROCm 10.1 ships HIP 7.16). +DEFAULT_ROCM_SERIES = "10.0" +DEFAULT_GPU_ARCHS = ("gfx942", "gfx950") + +# Classic non-pip ROCm install, consulted only after the SDK in the current environment. +SYSTEM_ROCM_ROOT = "/opt/rocm" + +# How many minor series beyond the detected one the pin accepts. Artifacts built against one ROCm +# minor have been verified to run on the next: amd-cupy/amd-hipcim/amd-monai built for 10.0 pass on +# a 10.1 runtime, including hiprtc JIT, hipFFT/hipBLAS/hipRAND and the cucim.skimage kernels. +ROCM_SERIES_SPAN = 2 + +# A value of None drops the requirement entirely: nvidia-ml-py raises at import time without an +# NVIDIA driver and nni depends on it; the nvidia-nvimgcodec wheels have no ROCm equivalent. +_HIPCIM = "amd-hipcim>=26.6.0; platform_system == 'Linux'" +_CUPY = "amd-cupy; platform_system == 'Linux'" + +SUBSTITUTIONS: dict[str, str | None] = { + "cucim-cu12": _HIPCIM, + "cucim-cu13": _HIPCIM, + "cupy-cuda12x": _CUPY, + "cupy-cuda13x": _CUPY, + "nvidia-nvimgcodec-cu12": None, + "nvidia-nvimgcodec-cu13": None, + "nvidia-ml-py": None, + "nni": None, +} + + +def _run(cmd: list[str]) -> str: + """stdout of ``cmd``, or "" if it is missing or fails.""" + try: + return subprocess.run(cmd, capture_output=True, text=True, check=False).stdout + except OSError: + return "" + + +def _installed_rocm_version() -> str: + """Version of the ``rocm`` package installed in this environment, or "" if absent.""" + try: + from importlib.metadata import version + + return version("rocm") + except Exception: + return "" + + +def _release_under(root: Path) -> str: + """ROCm release ``MAJOR.MINOR`` recorded under ``root``, or "" if not found there. + + Handles the pip ROCm SDK, whose release is the ``rocm`` package version in the ``*.dist-info`` + beside ``root``'s ``_rocm_sdk_*`` tree, and a classic install, whose release is in + ``.info/version``. + """ + for dist in sorted(root.parent.glob("rocm-*.dist-info")): + match = re.search(r"rocm-(\d+\.\d+)", dist.name) + if match: + return match.group(1) + try: + text = (root / ".info" / "version").read_text() + except OSError: + return "" + match = re.search(r"(\d+\.\d+)", text) + return match.group(1) if match else "" + + +def _rocm_release_from_disk() -> str: + """ROCm release read from the tree ``ROCM_PATH`` / ``ROCM_HOME`` point at, or "" if unset. + + Reads files only, so it still works under ``pip`` build isolation where the PATH-based + ``rocm-sdk`` console script may not be reachable. Mirrors hipCIM's detection. + """ + for root in (os.environ.get("ROCM_PATH"), os.environ.get("ROCM_HOME")): + if root: + release = _release_under(Path(root)) + if release: + return release + return "" + + +def detect_rocm_series() -> str: + """ + ROCm release ``MAJOR.MINOR`` to pin against. First match wins: ``MONAI_ROCM_SERIES``, the tree + ``ROCM_PATH``/``ROCM_HOME`` point at, the ``rocm`` package installed here, ``rocm-sdk version``, + a system install under :data:`SYSTEM_ROCM_ROOT`, else :data:`DEFAULT_ROCM_SERIES`. + + Two orderings matter. The HIP version is never consulted: ``torch.version.hip`` and + ``hipcc --version`` report HIP, a different number from the ROCm release the ``rocm`` package is + versioned by, so pinning to it yields a requirement no index can satisfy. And a system install + is consulted last, so an unrelated ``/opt/rocm`` cannot shadow the SDK being built against. + """ + for text in ( + os.environ.get("MONAI_ROCM_SERIES", ""), + _rocm_release_from_disk(), + _installed_rocm_version(), + _run(["rocm-sdk", "version"]), + _release_under(Path(SYSTEM_ROCM_ROOT)), + ): + match = re.search(r"(\d+\.\d+)", text) + if match: + return match.group(1) + return DEFAULT_ROCM_SERIES + + +def detect_gpu_archs() -> list[str]: + """ + AMD GPU architectures this build targets, used both for the ``torch[device-gfx...]`` extras and + for the ``rocm[device-gfx...]`` features. Read from ``GPU_TARGETS`` or ``AMDGPU_TARGETS``, + falling back to :data:`DEFAULT_GPU_ARCHS`. + """ + raw = os.environ.get("GPU_TARGETS") or os.environ.get("AMDGPU_TARGETS", "") + archs: list[str] = [] + for match in re.findall(r"gfx[0-9a-f]+", raw, re.IGNORECASE): + arch = match.lower() + if arch not in archs: + archs.append(arch) + return archs or list(DEFAULT_GPU_ARCHS) + + +def extras_for(name: str) -> list[str]: + """Device selector extras for ``name``; only ``torch`` carries them.""" + if name != "torch": + return [] + return [f"device-{arch}" for arch in detect_gpu_archs()] + + +def extra_requirements() -> list[str]: + """The ``rocm`` SDK requirement: runtime libraries, the development tree and a device feature per + targeted GPU arch, spanning :data:`ROCM_SERIES_SPAN` minor series from the detected one. + + ``devel`` is required after install, not just to build: ``monai/_extensions`` JIT-compiles its + HIP sources on first use through ``torch.utils.cpp_extension.load()``, which needs ``hipcc`` and + the ROCm headers. amd-cupy likewise JIT-compiles kernels through ``hipcc``. Whether the compiler + resolves without ``devel`` depends on where ``ROCM_PATH`` points -- it is in ``_rocm_sdk_core`` + on ROCm 10.x -- so do not drop it on the strength of a core-only layout happening to work. + """ + major, minor = (int(part) for part in detect_rocm_series().split(".")[:2]) + features = ["libraries", "devel"] + [f"device-{arch}" for arch in detect_gpu_archs()] + return [f"rocm[{','.join(features)}]>={major}.{minor}.0a0,<{major}.{minor + ROCM_SERIES_SPAN}"] diff --git a/pyproject.toml b/pyproject.toml index aa434b364fa..a0bc121d7e4 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -39,6 +39,8 @@ classifiers = [ "Topic :: Software Development :: Libraries", "Typing :: Typed", ] +# The NVIDIA distributions named below are swapped for the detected accelerator vendor's equivalents +# by monai/config/vendor_deps.py when MONAI is built or installed against that vendor's PyTorch. dependencies = [ "torch>=2.8.0", "numpy>=1.24,<3.0" diff --git a/setup.py b/setup.py index 73ff3460341..6a219df6d50 100644 --- a/setup.py +++ b/setup.py @@ -12,6 +12,7 @@ from __future__ import annotations import glob +import importlib.util import os import re import sys @@ -20,6 +21,7 @@ from packaging import version from setuptools import find_packages, setup +from setuptools.dist import Distribution import versioneer @@ -151,6 +153,46 @@ def get_cmds(): return cmds +def load_vendor_deps(): + """Import ``monai/config/vendor_deps.py`` by path, without importing the ``monai`` package.""" + path = os.path.join(os.path.dirname(os.path.abspath(__file__)), "monai", "config", "vendor_deps.py") + spec = importlib.util.spec_from_file_location("monai_vendor_deps", path) + module = importlib.util.module_from_spec(spec) + sys.modules[spec.name] = module # dataclasses resolves a class's module through sys.modules + spec.loader.exec_module(module) + return module + + +def vendor_distclass(): + """``Distribution`` subclass rewriting the dependency metadata for the detected accelerator + vendor, or None when there is none so the build keeps stock setuptools behaviour. + """ + try: + vendor_deps = load_vendor_deps() + vendor = vendor_deps.active_vendor() + except Exception as e: + # A vendor that was detected but could not be loaded is fatal: emitting another vendor's + # packages in this vendor's wheel is worse than failing the build. + if os.environ.get("MONAI_VENDOR") or type(e).__name__ == "VendorPluginError": + raise + warnings.warn(f"accelerator vendor detection skipped: {e}") + return None + if vendor is None: + return None + + class VendorDistribution(Distribution): + # pyproject.toml wins over anything passed to setup(), so the lists can only be adjusted once + # setuptools has applied them, and _finalize_requires() has to be re-run to propagate the + # result into the wheel metadata. + def parse_config_files(self, *args, **kwargs): + super().parse_config_files(*args, **kwargs) + self.install_requires = vendor.apply_to_dependencies(list(self.install_requires or [])) + self.extras_require = vendor.apply_to_optional_dependencies(dict(self.extras_require or {})) + self._finalize_requires() + + return VendorDistribution + + # Gathering source used for JIT extensions to include in package_data. jit_extension_source = [] @@ -160,9 +202,12 @@ def get_cmds(): jit_extension_source = [os.path.join("..", path) for path in jit_extension_source] +_distclass = vendor_distclass() + setup( version=versioneer.get_version(), cmdclass=get_cmds(), + **({"distclass": _distclass} if _distclass is not None else {}), packages=find_packages(exclude=("docs", "examples", "tests", "tests.*")), zip_safe=False, package_data=cast(Any, {"monai": ["py.typed", *jit_extension_source]}), diff --git a/tests/config/test_print_dependencies.py b/tests/config/test_print_dependencies.py index bbf8c4c7cda..f29299109a0 100644 --- a/tests/config/test_print_dependencies.py +++ b/tests/config/test_print_dependencies.py @@ -45,6 +45,10 @@ class TestPrintDependencies(unittest.TestCase): def setUp(self): + # pin the expectations to the NVIDIA/PyPI lists, whatever torch build the host has + vendor = patch.dict(os.environ, {"MONAI_VENDOR": "none"}) + vendor.start() + self.addCleanup(vendor.stop) self.toml = NamedTemporaryFile("w", delete=False) self.toml.write(TEST_TOML) self.toml.close() @@ -61,6 +65,21 @@ def test_missing_section(self): with self.assertRaises(KeyError): parse_dependencies(self.toml.name, ["nonexistent_section"]) + def test_vendor_dependencies(self): + env = {"MONAI_VENDOR": "rocm", "MONAI_ROCM_SERIES": "10.0", "GPU_TARGETS": "gfx942"} + with patch.dict(os.environ, env): + deps = parse_dependencies(self.toml.name, ["build-system"]) + self.assertEqual( + [ + "numpy", + "rocm[libraries,devel,device-gfx942]>=10.0.0a0,<10.2", + "setuptools", + "torch[device-gfx942]", + "wheel", + ], + deps, + ) + def test_print_dependencies(self): out = StringIO() with redirect_stdout(out), patch("monai.config.print_dependencies.TOML_FILE", self.toml.name): diff --git a/tests/config/test_print_info.py b/tests/config/test_print_info.py index da91c431498..5a67f15ddc4 100644 --- a/tests/config/test_print_info.py +++ b/tests/config/test_print_info.py @@ -11,10 +11,13 @@ from __future__ import annotations +import sys import unittest from io import StringIO +from unittest.mock import patch from monai.config import print_debug_info +from monai.config.deviceconfig import get_gpu_info class TestPrintInfo(unittest.TestCase): @@ -24,6 +27,38 @@ def test_print_info(self): print_debug_info(file=out) self.assertGreater(out.tell(), 0) + @patch("torch.backends.cudnn.version", lambda: 0) + @patch("torch.cuda.device_count", lambda: 0) + @patch("torch.cuda.is_available", lambda: True) + @patch("torch.version.hip", "7.0.0") + def test_hip_version(self): + """A ROCm build reports the HIP version, where `torch.version.cuda` would be None.""" + output = get_gpu_info() + self.assertEqual(output["HIP version"], "7.0.0") + self.assertNotIn("CUDA version", output) + + @patch("torch.backends.cudnn.version", lambda: 0) + @patch("torch.cuda.device_count", lambda: 0) + @patch("torch.cuda.is_available", lambda: True) + @patch("torch.version.hip", None) + @patch("torch.version.cuda", "12.4") + def test_cuda_version(self): + """A CUDA build is unaffected and still reports the CUDA version.""" + output = get_gpu_info() + self.assertEqual(output["CUDA version"], "12.4") + self.assertNotIn("HIP version", output) + + @patch("torch.version.hip", "7.0.0") + def test_check_env_hip_version(self): + """check_env.check_torch_cuda() prints 'HIP version:' on a ROCm build.""" + import monai.config.check_env as ce + + buf = StringIO() + with patch.object(sys, "stdout", buf): + ce.check_torch_cuda() + self.assertIn("HIP version:", buf.getvalue()) + self.assertNotIn("CUDA version:", buf.getvalue()) + if __name__ == "__main__": unittest.main() diff --git a/tests/config/test_vendor_deps.py b/tests/config/test_vendor_deps.py new file mode 100644 index 00000000000..41d0cc20938 --- /dev/null +++ b/tests/config/test_vendor_deps.py @@ -0,0 +1,141 @@ +# Copyright (c) MONAI Consortium +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# http://www.apache.org/licenses/LICENSE-2.0 +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from __future__ import annotations + +import types +import unittest +from unittest.mock import patch + +from monai.config import vendor_deps +from monai.config.vendor_deps import ( + active_vendor, + apply_to_dependencies, + apply_to_optional_dependencies, + canonical_name, + split_requirement, +) + + +def fake_plugin(substitutions=None, extras=None, extra_requirements=()): + """A minimal plugin satisfying the contract documented in vendor_deps.""" + return types.SimpleNamespace( + NAME="fake", + SUBSTITUTIONS=substitutions or {}, + extras_for=lambda name: (extras or {}).get(name, []), + extra_requirements=lambda: list(extra_requirements), + ) + + +class TestRequirementParsing(unittest.TestCase): + def test_canonical_name(self): + self.assertEqual("nvidia-ml-py", canonical_name("NVIDIA_ML.Py")) + + def test_split_plain(self): + self.assertEqual(("numpy", [], ">=1.24,<3.0"), split_requirement("numpy>=1.24,<3.0")) + + def test_split_extras_and_marker(self): + name, extras, rest = split_requirement("cupy-cuda13x[ctk]!=14.1.0; platform_system == 'Linux'") + self.assertEqual("cupy-cuda13x", name) + self.assertEqual(["ctk"], extras) + self.assertEqual("!=14.1.0; platform_system == 'Linux'", rest) + + def test_split_direct_url(self): + name, extras, _ = split_requirement("MetricsReloaded @ git+https://github.com/x/y@z") + self.assertEqual("MetricsReloaded", name) + self.assertEqual([], extras) + + +class TestGenericRewrites(unittest.TestCase): + def test_substitution_drop_and_dedupe(self): + plugin = fake_plugin({"a-cu12": "repl", "a-cu13": "repl", "gone": None}) + self.assertEqual( + {"g": ["repl", "kept"]}, apply_to_optional_dependencies({"g": ["a-cu12", "a-cu13", "gone", "kept"]}, plugin) + ) + + def test_group_can_become_empty(self): + self.assertEqual({"g": []}, apply_to_optional_dependencies({"g": ["gone"]}, fake_plugin({"gone": None}))) + + def test_extras_and_extra_requirements(self): + plugin = fake_plugin(extras={"torch": ["device-x"]}, extra_requirements=["runtime>=1"]) + self.assertEqual( + ["torch[device-x]>=2.8.0", "numpy>=1.24", "runtime>=1"], + apply_to_dependencies(["torch>=2.8.0", "numpy>=1.24"], plugin), + ) + + def test_existing_extras_preserved(self): + plugin = fake_plugin(extras={"torch": ["device-x"]}) + self.assertEqual( + ["torch[device-x,opt-einsum]>=2.8.0"], apply_to_dependencies(["torch[opt-einsum]>=2.8.0"], plugin) + ) + + def test_unrelated_requirements_untouched(self): + group = ["scikit-image>=0.19.0", "MetricsReloaded @ git+https://github.com/x/y@z"] + self.assertEqual({"all": group}, apply_to_optional_dependencies({"all": group}, fake_plugin())) + + +class TestVendorSelection(unittest.TestCase): + def test_none_disables_rewriting(self): + for value in ("none", "NONE", "cpu"): + with patch.dict("os.environ", {"MONAI_VENDOR": value}): + self.assertIsNone(active_vendor()) + + def test_unknown_vendor_is_rejected(self): + with patch.dict("os.environ", {"MONAI_VENDOR": "nosuchvendor"}), self.assertRaises(ValueError): + active_vendor() + + def test_explicit_vendor_bypasses_detection(self): + with patch.dict("os.environ", {"MONAI_VENDOR": "rocm"}): + vendor = active_vendor() + self.assertIsNotNone(vendor) + self.assertEqual("rocm", vendor.name) + + def test_no_vendor_when_no_probe_matches(self): + with patch.dict("os.environ", {}, clear=True), patch.dict(vendor_deps.REGISTRY, {}, clear=True): + self.assertIsNone(active_vendor()) + + def test_plugin_not_imported_when_probe_does_not_match(self): + probe = vendor_deps._Registration("vendor_rocm", lambda: False) + with ( + patch.dict("os.environ", {}, clear=True), + patch.dict(vendor_deps.REGISTRY, {"rocm": probe}, clear=True), + patch.object(vendor_deps, "_import_plugin", side_effect=AssertionError("plugin must not be imported")), + ): + self.assertIsNone(active_vendor()) + + def test_detection_selects_matching_plugin(self): + probe = vendor_deps._Registration("vendor_rocm", lambda: True) + with patch.dict("os.environ", {}, clear=True), patch.dict(vendor_deps.REGISTRY, {"rocm": probe}, clear=True): + vendor = active_vendor() + self.assertEqual("rocm", vendor.name) + + def test_detected_vendor_that_cannot_load_is_fatal(self): + # Falling back to "no vendor" here would ship another vendor's packages in this one's wheel. + probe = vendor_deps._Registration("vendor_broken", lambda: True) + with ( + patch.dict("os.environ", {}, clear=True), + patch.dict(vendor_deps.REGISTRY, {"rocm": probe}, clear=True), + patch.object(vendor_deps, "_import_plugin", side_effect=SyntaxError("invalid syntax")), + self.assertRaises(vendor_deps.VendorPluginError), + ): + active_vendor() + + def test_explicit_vendor_that_cannot_load_is_fatal(self): + with ( + patch.dict("os.environ", {"MONAI_VENDOR": "rocm"}), + patch.object(vendor_deps, "_import_plugin", side_effect=SyntaxError("invalid syntax")), + self.assertRaises(vendor_deps.VendorPluginError), + ): + active_vendor() + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/config/test_vendor_rocm.py b/tests/config/test_vendor_rocm.py new file mode 100644 index 00000000000..8bae2e9b2ae --- /dev/null +++ b/tests/config/test_vendor_rocm.py @@ -0,0 +1,197 @@ +# Copyright (c) MONAI Consortium +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# http://www.apache.org/licenses/LICENSE-2.0 +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from __future__ import annotations + +import tempfile +import unittest +from pathlib import Path +from unittest.mock import patch + +from monai.config import vendor_rocm +from monai.config.vendor_deps import apply_to_dependencies, apply_to_optional_dependencies +from monai.config.vendor_rocm import ( + DEFAULT_GPU_ARCHS, + DEFAULT_ROCM_SERIES, + detect_gpu_archs, + detect_rocm_series, + extra_requirements, + extras_for, +) + + +class TestReleaseFromDisk(unittest.TestCase): + """Exercises the real filesystem parsing rather than mocking it out.""" + + def setUp(self): + tmp = tempfile.TemporaryDirectory() + self.addCleanup(tmp.cleanup) + self.root = Path(tmp.name) + + def test_pip_sdk_layout(self): + # ROCM_PATH points into site-packages, where the `rocm` dist-info sits beside it. + site = self.root / "site-packages" + (site / "rocm-10.1.0.dist-info").mkdir(parents=True) + (site / "_rocm_sdk_core").mkdir() + with patch.dict("os.environ", {"ROCM_PATH": str(site / "_rocm_sdk_core")}, clear=True): + self.assertEqual("10.1", vendor_rocm._rocm_release_from_disk()) + + def test_pip_sdk_layout_ignores_sibling_rocm_packages(self): + site = self.root / "site-packages" + (site / "rocm_sdk_core-10.1.0.dist-info").mkdir(parents=True) + (site / "rocm_bootstrap-0.3.0.dist-info").mkdir() + (site / "rocm-10.0.0.dist-info").mkdir() + (site / "_rocm_sdk_core").mkdir() + with patch.dict("os.environ", {"ROCM_PATH": str(site / "_rocm_sdk_core")}, clear=True): + self.assertEqual("10.0", vendor_rocm._rocm_release_from_disk()) + + def test_classic_layout(self): + rocm = self.root / "opt" / "rocm" + (rocm / ".info").mkdir(parents=True) + (rocm / ".info" / "version").write_text("10.0.0-18\n") + with patch.dict("os.environ", {"ROCM_PATH": str(rocm)}, clear=True): + self.assertEqual("10.0", vendor_rocm._rocm_release_from_disk()) + + def test_rocm_home_is_consulted(self): + rocm = self.root / "rocm" + (rocm / ".info").mkdir(parents=True) + (rocm / ".info" / "version").write_text("7.14.1\n") + with patch.dict("os.environ", {"ROCM_HOME": str(rocm)}, clear=True): + self.assertEqual("7.14", vendor_rocm._rocm_release_from_disk()) + + def test_nothing_found(self): + with patch.dict("os.environ", {"ROCM_PATH": str(self.root / "absent")}, clear=True): + self.assertEqual("", vendor_rocm._rocm_release_from_disk()) + + def test_feeds_detect_rocm_series(self): + site = self.root / "site-packages" + (site / "rocm-10.1.0.dist-info").mkdir(parents=True) + (site / "_rocm_sdk_core").mkdir() + with patch.dict("os.environ", {"ROCM_PATH": str(site / "_rocm_sdk_core")}, clear=True): + self.assertEqual("10.1", detect_rocm_series()) + + def test_system_install_does_not_shadow_this_environment(self): + # A host /opt/rocm must not win over the SDK the build is actually using. + system = self.root / "opt" / "rocm" + (system / ".info").mkdir(parents=True) + (system / ".info" / "version").write_text("7.2.3\n") + with ( + patch.dict("os.environ", {}, clear=True), + patch.object(vendor_rocm, "SYSTEM_ROCM_ROOT", str(system)), + patch.object(vendor_rocm, "_installed_rocm_version", return_value="10.1.0"), + ): + self.assertEqual("10.1", detect_rocm_series()) + + def test_system_install_is_used_when_nothing_else_knows(self): + system = self.root / "opt" / "rocm" + (system / ".info").mkdir(parents=True) + (system / ".info" / "version").write_text("7.2.3\n") + with ( + patch.dict("os.environ", {}, clear=True), + patch.object(vendor_rocm, "SYSTEM_ROCM_ROOT", str(system)), + patch.object(vendor_rocm, "_installed_rocm_version", return_value=""), + patch.object(vendor_rocm, "_run", return_value=""), + ): + self.assertEqual("7.2", detect_rocm_series()) + + +class TestDetection(unittest.TestCase): + def test_series_from_env(self): + with patch.dict("os.environ", {"MONAI_ROCM_SERIES": "7.14.1"}): + self.assertEqual("7.14", detect_rocm_series()) + + def test_series_fallback(self): + with ( + patch.dict("os.environ", {}, clear=True), + patch.object(vendor_rocm, "_rocm_release_from_disk", return_value=""), + patch.object(vendor_rocm, "_installed_rocm_version", return_value=""), + patch.object(vendor_rocm, "SYSTEM_ROCM_ROOT", "/nonexistent-rocm"), + patch.object(vendor_rocm, "_run", return_value=""), + ): + self.assertEqual(DEFAULT_ROCM_SERIES, detect_rocm_series()) + + def test_series_from_disk(self): + with ( + patch.dict("os.environ", {}, clear=True), + patch.object(vendor_rocm, "_rocm_release_from_disk", return_value="10.0"), + ): + self.assertEqual("10.0", detect_rocm_series()) + + def test_series_never_uses_hip_version(self): + # ROCm releases and HIP carry different numbers; pinning `rocm` to HIP matches no index. + with ( + patch.dict("os.environ", {}, clear=True), + patch.object(vendor_rocm, "_rocm_release_from_disk", return_value=""), + patch.object(vendor_rocm, "_run", return_value="HIP version: 7.16.26385"), + ): + self.assertEqual("7.16", detect_rocm_series()) + self.assertFalse(hasattr(vendor_rocm, "_torch_hip_version")) + + def test_gpu_archs_from_env(self): + with patch.dict("os.environ", {"GPU_TARGETS": "gfx90a;GFX942;gfx942"}): + self.assertEqual(["gfx90a", "gfx942"], detect_gpu_archs()) + + def test_gpu_archs_fallback(self): + with patch.dict("os.environ", {}, clear=True): + self.assertEqual(list(DEFAULT_GPU_ARCHS), detect_gpu_archs()) + + +class TestPluginContract(unittest.TestCase): + def setUp(self): + env = patch.dict("os.environ", {"MONAI_ROCM_SERIES": "10.0", "GPU_TARGETS": "gfx942"}) + env.start() + self.addCleanup(env.stop) + + def test_extras_only_for_torch(self): + self.assertEqual(["device-gfx942"], extras_for("torch")) + self.assertEqual([], extras_for("numpy")) + + def test_rocm_requirement(self): + self.assertEqual(["rocm[libraries,devel,device-gfx942]>=10.0.0a0,<10.2"], extra_requirements()) + + def test_dependencies(self): + self.assertEqual( + ["torch[device-gfx942]>=2.8.0", "numpy>=1.24,<3.0", "rocm[libraries,devel,device-gfx942]>=10.0.0a0,<10.2"], + apply_to_dependencies(["torch>=2.8.0", "numpy>=1.24,<3.0"], vendor_rocm), + ) + + def test_cucim_variants_collapse(self): + extras = { + "cucim": [ + "cucim-cu12; platform_system == 'Linux' and python_version <= '3.10'", + "cucim-cu13; platform_system == 'Linux' and python_version >= '3.11'", + ] + } + self.assertEqual( + {"cucim": ["amd-hipcim>=26.6.0; platform_system == 'Linux'"]}, + apply_to_optional_dependencies(extras, vendor_rocm), + ) + + def test_cupy_and_dropped_packages(self): + extras = { + "cupy": ["cupy-cuda13x!=14.1.0"], + "nvimgcodec": ["pydicom", "cupy-cuda13x[ctk]!=14.1.0", "nvidia-nvimgcodec-cu13[all]>=0.8.0"], + "pynvml": ["nvidia-ml-py"], + "nni": ["nni; platform_system == 'Linux'", "filelock<3.12.0"], + } + self.assertEqual( + { + "cupy": ["amd-cupy; platform_system == 'Linux'"], + "nvimgcodec": ["pydicom", "amd-cupy; platform_system == 'Linux'"], + "pynvml": [], + "nni": ["filelock<3.12.0"], + }, + apply_to_optional_dependencies(extras, vendor_rocm), + ) + + +if __name__ == "__main__": + unittest.main()