diff --git a/nf_core/modules/containers.py b/nf_core/modules/containers.py index 132e443499..a36570ce21 100644 --- a/nf_core/modules/containers.py +++ b/nf_core/modules/containers.py @@ -7,6 +7,7 @@ from collections.abc import Callable from concurrent.futures import ThreadPoolExecutor, as_completed from pathlib import Path +from typing import NamedTuple from urllib.parse import quote import requests @@ -28,7 +29,14 @@ scan_modules_dir, ) from nf_core.pipelines.lint_utils import run_prettier_on_file -from nf_core.utils import CONTAINER_PLATFORMS, CONTAINER_SYSTEMS, ContainerRegistryUrls +from nf_core.utils import ( + CONTAINER_PLATFORMS, + CONTAINER_SYSTEMS, + GPU_CONTAINER_PLATFORMS, + GPU_CONTAINER_SUFFIX, + GPU_CONTAINER_SYSTEMS, + ContainerRegistryUrls, +) log = logging.getLogger(__name__) @@ -39,6 +47,28 @@ # Wave container build `format` field, keyed by nf-core container system. WAVE_FORMAT = {"docker": "docker", "singularity": "sif"} +# Wave build templates. CPU environments build with pixi; GPU environments must use the +# micromamba v2 template, because the pixi build image has no `__cuda` virtual package and +# fails to solve the `__cuda`-gated CUDA packages (e.g. pytorch-gpu). +WAVE_BUILD_TEMPLATE = "conda/pixi:v1" +WAVE_GPU_BUILD_TEMPLATE = "conda/micromamba:v2" + +ENVIRONMENT_GPU_YML = "environment.gpu.yml" + +# All meta.yml container/conda fields (CPU + GPU), used when iterating to persist or list. +ALL_CONTAINER_SYSTEMS = CONTAINER_SYSTEMS + GPU_CONTAINER_SYSTEMS +ALL_CONDA_FIELDS = ["conda", f"conda{GPU_CONTAINER_SUFFIX}"] + + +class BuildTarget(NamedTuple): + """A single Wave build to run: which meta.yml field it fills and how to build it.""" + + field: str # meta.yml container field, e.g. "docker" or "docker_gpu" + system: str # base container system passed to Wave: "docker" or "singularity" + platform: str + env_file: Path + build_template: str # Wave build template (pixi for CPU, micromamba v2 for GPU) + def wave_send( method: str, url: str, json_body: dict | None = None, error_context: str = "Wave request" @@ -157,12 +187,17 @@ def __init__( self.module_directory: Path | None = self.nfcore_component.component_dir self.environment_yml: Path | None = self.nfcore_component.environment_yml self.meta_yml: Path | None = self.nfcore_component.meta_yml + # GPU modules (dual-container pattern) carry a second environment.gpu.yml + # alongside the CPU environment.yml. Not tracked by NFCoreComponent. + gpu_env = self.nfcore_component.component_dir / ENVIRONMENT_GPU_YML + self.environment_gpu_yml: Path | None = gpu_env if gpu_env.exists() else None else: # For all modules mode, these will be set per module during iteration self.nfcore_component = None self.module_directory = None self.environment_yml = None self.meta_yml = None + self.environment_gpu_yml = None self.containers: MetaYmlContainers | None = None self._module_lint: ModuleLint | None = None @@ -265,8 +300,12 @@ def cleanup_stale_conda_lock_files(self, new_lock_files: set[Path]) -> None: log.debug(f"Could not remove .conda-lock directory: {e}") def update_main_nf_container(self, force=False) -> None: - """Update the container name in main.nf using the docker amd64 image without registry. - Don't update if the container name is already correct. + """Update the container directive in main.nf using the amd64 images. + + For a normal module this writes the standard singularity/docker ternary. For a GPU + module (dual-container pattern) it writes the nested ternary that also switches on + ``task.accelerator`` between the CPU and GPU images. + Don't update if the directive is already correct. """ import re @@ -294,12 +333,29 @@ def update_main_nf_container(self, force=False) -> None: log.error(f"No singularity image found for {linux_amd64}") return + # GPU images (dual-container pattern), if this module has an environment.gpu.yml. + gpu_platform = GPU_CONTAINER_PLATFORMS[0] + _docker_gpu_entry = self.containers.docker_gpu.get(gpu_platform) + _singularity_gpu_entry = self.containers.singularity_gpu.get(gpu_platform) + docker_gpu_image = _docker_gpu_entry.name if _docker_gpu_entry else "" + singularity_gpu_image = ( + (_singularity_gpu_entry.https or _singularity_gpu_entry.name) if _singularity_gpu_entry else "" + ) + is_gpu = bool(docker_gpu_image and singularity_gpu_image) + if (docker_gpu_image or singularity_gpu_image) and not is_gpu: + log.warning( + f"Incomplete GPU containers for '{self.module}' - writing CPU-only container directive in main.nf" + ) + # Read main.nf main_nf_path = self.nfcore_component.main_nf content = main_nf_path.read_text() - # Check if the container directive is already correct - if docker_image in content and singularity_image in content and not force: + # Check if the container directive is already correct (all relevant images present) + required_images = [docker_image, singularity_image] + if is_gpu: + required_images += [docker_gpu_image, singularity_gpu_image] + if all(img in content for img in required_images) and not force: log.debug(f"Container directive in `{self.nfcore_component.component_name}/main.nf` is already correct") return @@ -309,6 +365,15 @@ def update_main_nf_container(self, force=False) -> None: def _replace(match: "re.Match[str]") -> str: indent = match.group(1) inner_indent = indent + " " + if is_gpu: + # Nested ternary: first pick engine (singularity/docker), then pick the + # CPU/GPU image within each engine based on task.accelerator. + return ( + f"{indent}container \"${{ workflow.containerEngine in ['singularity', 'apptainer'] " + "&& !task.ext.singularity_pull_docker_container ?\n" + f"{inner_indent}(task.accelerator ? '{singularity_gpu_image}' : '{singularity_image}') :\n" + f"{inner_indent}(task.accelerator ? '{docker_gpu_image}' : '{docker_image}') }}\"" + ) return ( f"{indent}container \"${{ workflow.containerEngine in ['singularity', 'apptainer'] " "&& !task.ext.singularity_pull_docker_container\n" @@ -321,10 +386,103 @@ def _replace(match: "re.Match[str]") -> str: ) main_nf_path.write_text(new_content) - log.debug( - f"Updated container in `{self.nfcore_component.component_name}/main.nf` " - f"(docker: `{docker_image}`, singularity: `{singularity_image}`)" - ) + if is_gpu: + log.debug( + f"Updated GPU container in `{self.nfcore_component.component_name}/main.nf` " + f"(docker: `{docker_image}` / `{docker_gpu_image}`, " + f"singularity: `{singularity_image}` / `{singularity_gpu_image}`)" + ) + else: + log.debug( + f"Updated container in `{self.nfcore_component.component_name}/main.nf` " + f"(docker: `{docker_image}`, singularity: `{singularity_image}`)" + ) + + def _build_targets(self) -> list[BuildTarget]: + """ + Enumerate the Wave builds to run for this module. + + Always includes the CPU targets (docker/singularity × all platforms) built from + ``environment.yml`` with the pixi template. GPU modules that also carry an + ``environment.gpu.yml`` add the GPU targets (docker/singularity × GPU platforms), + built with the micromamba v2 template and stored under the ``*_gpu`` fields. + """ + assert self.environment_yml is not None + targets = [ + BuildTarget(cs, cs, platform, self.environment_yml, WAVE_BUILD_TEMPLATE) + for cs in CONTAINER_SYSTEMS + for platform in CONTAINER_PLATFORMS + ] + if self.environment_gpu_yml is not None: + targets += [ + BuildTarget( + f"{cs}{GPU_CONTAINER_SUFFIX}", + cs, + platform, + self.environment_gpu_yml, + WAVE_GPU_BUILD_TEMPLATE, + ) + for cs in CONTAINER_SYSTEMS + for platform in GPU_CONTAINER_PLATFORMS + ] + return targets + + def _download_conda_locks( + self, + containers: MetaYmlContainers, + gpu: bool, + progress_bar: rich.progress.Progress | None = None, + task_id: rich.progress.TaskID | None = None, + ) -> set[Path]: + """ + Download the conda lock files for the CPU (``gpu=False``) or GPU (``gpu=True``) set, + writing them under ``.conda-lock/`` and registering the paths on ``containers``. + + Returns the set of lock file paths written (for stale-file cleanup). A missing lock + is not fatal: the other platforms' locks are kept. + """ + assert self.module_directory is not None + suffix = GPU_CONTAINER_SUFFIX if gpu else "" + docker_field = f"docker{suffix}" + conda_field = f"conda{suffix}" + platforms = GPU_CONTAINER_PLATFORMS if gpu else CONTAINER_PLATFORMS + label = "conda lock gpu" if gpu else "conda lock" + + written: set[Path] = set() + for platform in platforms: + # Get docker build ID for this platform (the conda lock is keyed off the docker build) + _docker_entry = getattr(containers, docker_field).get(platform) + build_id = _docker_entry.build_id if _docker_entry else "" + short_platform = platform.split("/")[-1] + if not build_id: + log.debug(f"Docker image for {docker_field} {platform} missing - Conda-lock skipped") + if progress_bar and task_id is not None: + progress_bar.update(task_id, status=f"{label} {short_platform} skipped") + continue + + conda_lock_path = self.module_directory / ".conda-lock" / f"{platform.replace('/', '_')}-{build_id}.txt" + conda_lock_path.parent.mkdir(parents=True, exist_ok=True) + + try: + log.debug(f"Downloading conda lock file for {docker_field} {platform} to {conda_lock_path}") + if progress_bar and task_id is not None: + progress_bar.update(task_id, status=f"{label} {short_platform}...") + conda_lock_path.write_text(self.get_conda_lock_file(platform, gpu=gpu)) + # Only register the entry once the lock has actually been written, so + # meta.yml never points at a missing lock file. + getattr(containers, conda_field)[platform] = CondaEntry(lock_file=str(conda_lock_path)) + written.add(conda_lock_path) + if progress_bar and task_id is not None: + progress_bar.update(task_id, status=f"{label} {short_platform} done") + + except (ValueError, RuntimeError, OSError, requests.RequestException) as e: + # Wave does not always expose a conda lock for every build (e.g. some + # freshly built arm64 images). A missing lock is not fatal: keep the + # other platforms' containers and locks instead of failing the module. + log.warning(f"No conda lock file available for {docker_field} {platform}: {e}") + if progress_bar and task_id is not None: + progress_bar.update(task_id, status=f"{label} {short_platform} unavailable") + return written def create( self, @@ -333,7 +491,12 @@ def create( force: bool = False, ) -> tuple[MetaYmlContainers, bool]: """ - Build docker and singularity containers for linux/amd64 and linux/arm64 using wave. + Build docker and singularity containers using wave. + + Builds the CPU containers (linux/amd64 and linux/arm64) from ``environment.yml``. + For GPU modules (dual-container pattern) that also carry an ``environment.gpu.yml``, + builds the parallel GPU containers (linux/amd64 only) with the micromamba v2 build + template and stores them under the ``*_gpu`` meta.yml fields. Args: progress_bar: Optional progress bar to use for tracking progress @@ -344,7 +507,6 @@ def create( """ containers = MetaYmlContainers() build_tasks = {} - threads = max(len(CONTAINER_SYSTEMS) * len(CONTAINER_PLATFORMS), 1) has_failures = False if not self.environment_yml: @@ -354,22 +516,26 @@ def create( raise RuntimeError("No environment.yml found.") assert self.module_directory is not None + targets = self._build_targets() + threads = max(len(targets), 1) + if self.environment_gpu_yml is not None: + log.debug(f"GPU environment found for '{self.module}' - building GPU containers as well") + # One spinner per build target — they run visually in parallel - build_task_ids: dict[tuple[str, str], rich.progress.TaskID] = {} + build_task_ids: dict[BuildTarget, rich.progress.TaskID] = {} if progress_bar: - for cs in CONTAINER_SYSTEMS: - for platform in CONTAINER_PLATFORMS: - short_platform = platform.split("/")[-1] - build_task_ids[(cs, platform)] = progress_bar.add_task( - f" {cs}/{short_platform}", - total=1, - completed=0, - status="submitting wave build request...", - ) - - def make_on_build_id(cs: str, platform: str) -> Callable[[str], None]: + for target in targets: + short_platform = target.platform.split("/")[-1] + build_task_ids[target] = progress_bar.add_task( + f" {target.field}/{short_platform}", + total=1, + completed=0, + status="submitting wave build request...", + ) + + def make_on_build_id(target: BuildTarget) -> Callable[[str], None]: def callback(build_id: str) -> None: - build_tid = build_task_ids.get((cs, platform)) + build_tid = build_task_ids.get(target) if progress_bar and build_tid is not None: url = f"{WAVE_URL}/view/builds/{build_id}" progress_bar.update(build_tid, status=f"building… {url}") @@ -382,38 +548,37 @@ def callback(build_id: str) -> None: # Submit all container build tasks with ThreadPoolExecutor(max_workers=threads) as pool: - for cs in CONTAINER_SYSTEMS: - for platform in CONTAINER_PLATFORMS: - fut = pool.submit( - self.request_container, - cs, - platform, - self.environment_yml, - self.verbose, - make_on_build_id(cs, platform) if progress_bar else None, - cancel_event, - ) - build_tasks[fut] = (cs, platform) + for target in targets: + fut = pool.submit( + self.request_container, + target.system, + target.platform, + target.env_file, + self.verbose, + make_on_build_id(target) if progress_bar else None, + cancel_event, + target.build_template, + ) + build_tasks[fut] = target # Process completed container builds try: for fut in as_completed(build_tasks): - cs, platform = build_tasks[fut] - short_platform = platform.split("/")[-1] - build_tid = build_task_ids.get((cs, platform)) + target = build_tasks[fut] + build_tid = build_task_ids.get(target) try: - getattr(containers, cs)[platform] = fut.result() + getattr(containers, target.field)[target.platform] = fut.result() if progress_bar and build_tid is not None: progress_bar.update(build_tid, completed=1, status="[green]done[/green]") except (ValueError, RuntimeError, OSError, AssertionError) as e: # make it a warning for arm (not required), but fail for other platforms - if platform == "linux/arm64": + if target.platform == "linux/arm64": log.warning( - f"Failed to build {cs} container for {platform}: {e}. This is only critical if the tool should support arm64." + f"Failed to build {target.field} container for {target.platform}: {e}. This is only critical if the tool should support arm64." ) else: - log.error(f"Failed to build {cs} container for {platform}: {e}") + log.error(f"Failed to build {target.field} container for {target.platform}: {e}") has_failures = True if progress_bar and build_tid is not None: progress_bar.update(build_tid, completed=1, status="[red]failed[/red]") @@ -435,49 +600,19 @@ def callback(build_id: str) -> None: # Set containers early so get_conda_lock_file can access it self.containers = containers - # Download conda lock files as separate tasks - new_lock_files = set() - for platform in CONTAINER_PLATFORMS: - # Get docker build ID for this platform - _docker_entry = containers.docker.get(platform) - build_id = _docker_entry.build_id if _docker_entry else "" - short_platform = platform.split("/")[-1] - if not build_id: - log.debug(f"Docker image for {platform} missing - Conda-lock skipped") - if progress_bar and task_id is not None: - progress_bar.update(task_id, status=f"conda lock {short_platform} skipped") - continue - - conda_lock_path = self.module_directory / ".conda-lock" / f"{platform.replace('/', '_')}-{build_id}.txt" - conda_lock_path.parent.mkdir(parents=True, exist_ok=True) - - try: - # Download conda lock file (it will look up build_id from docker container) - log.debug(f"Downloading conda lock file for {platform} to {conda_lock_path}") - if progress_bar and task_id is not None: - progress_bar.update(task_id, status=f"conda lock {short_platform}...") - conda_lock_path.write_text(self.get_conda_lock_file(platform)) - # Only register the entry once the lock has actually been written, so - # meta.yml never points at a missing lock file. - containers.conda[platform] = CondaEntry(lock_file=str(conda_lock_path)) - new_lock_files.add(conda_lock_path) - if progress_bar and task_id is not None: - progress_bar.update(task_id, status=f"conda lock {short_platform} done") - - except (ValueError, RuntimeError, OSError, requests.RequestException) as e: - # Wave does not always expose a conda lock for every build (e.g. some - # freshly built arm64 images). A missing lock is not fatal: keep the - # other platforms' containers and locks instead of failing the module. - log.warning(f"No conda lock file available for {platform}: {e}") - if progress_bar and task_id is not None: - progress_bar.update(task_id, status=f"conda lock {short_platform} unavailable") + # Download conda lock files (CPU, plus GPU for dual-container modules) + new_lock_files = self._download_conda_locks(containers, gpu=False, progress_bar=progress_bar, task_id=task_id) + if self.environment_gpu_yml is not None: + new_lock_files |= self._download_conda_locks( + containers, gpu=True, progress_bar=progress_bar, task_id=task_id + ) # Clean up stale conda-lock files self.cleanup_stale_conda_lock_files(new_lock_files) # Persist everything (docker, singularity and conda locks) to meta.yml in a # single write. Partial results are kept even if some builds failed. - if any(getattr(containers, cs) for cs in CONTAINER_SYSTEMS + ["conda"]): + if any(getattr(containers, cs) for cs in ALL_CONTAINER_SYSTEMS + ALL_CONDA_FIELDS): try: self.update_containers_in_meta() log.debug("Updated meta.yml with built containers") @@ -554,6 +689,7 @@ def request_container( verbose=False, on_build_id: Callable[[str], None] | None = None, cancel_event: threading.Event | None = None, + build_template: str = WAVE_BUILD_TEMPLATE, ) -> ContainerEntry: assert conda_file.exists() assert container_system in CONTAINER_SYSTEMS @@ -561,6 +697,8 @@ def request_container( # Submit the build via the Wave HTTP API (POST /v1alpha2/container). # `freeze` with no buildRepository pushes to the public community registry. + # GPU environments must pass WAVE_GPU_BUILD_TEMPLATE (micromamba v2) so Wave can + # resolve the `__cuda`-gated CUDA packages that the pixi template cannot solve. payload: dict = { "packages": { "type": "CONDA", @@ -568,7 +706,7 @@ def request_container( }, "containerPlatform": platform, "freeze": True, - "buildTemplate": "conda/pixi:v1", + "buildTemplate": build_template, "format": WAVE_FORMAT[container_system], "nameStrategy": "imageSuffix", } @@ -644,23 +782,27 @@ def get_conda_lock_url(build_id) -> str: url = f"{WAVE_API_ALPHA1}/builds/{build_id_safe}/condalock" return url - def get_conda_lock_file(self, platform: str) -> str: + def get_conda_lock_file(self, platform: str, gpu: bool = False) -> str: """ Get the conda lock file for an existing environment. Try (in that order): 1. reading from meta.yml 2. reading from cached containers 3. recreating with wave commands + + When ``gpu`` is set, the lock is resolved from the GPU docker build (``docker_gpu``) + rather than the CPU one. """ assert platform in CONTAINER_PLATFORMS containers = self.containers or self.get_containers_from_meta() or self.create()[0] - # Get build_id from docker container for this platform - _docker_entry = containers.docker.get(platform) if containers else None + # Get build_id from the docker (or docker_gpu) container for this platform + docker_field = f"docker{GPU_CONTAINER_SUFFIX}" if gpu else "docker" + _docker_entry = getattr(containers, docker_field).get(platform) if containers else None build_id = _docker_entry.build_id if _docker_entry else None if not build_id: - raise ValueError(f"No build_id found for docker container on platform {platform}") + raise ValueError(f"No build_id found for {docker_field} container on platform {platform}") # Generate the conda lock URL from the build_id conda_lock_url = self.get_conda_lock_url(build_id) @@ -680,17 +822,19 @@ def list_containers(self) -> list[tuple[str, str, str]]: if not containers_valid: return [] containers_flat = [] - for cs in CONTAINER_SYSTEMS + ["conda"]: - for p in CONTAINER_PLATFORMS: - if cs == "conda": - conda_entry = containers_valid.conda.get(p) + conda_fields = set(ALL_CONDA_FIELDS) + for cs in ALL_CONTAINER_SYSTEMS + ALL_CONDA_FIELDS: + platforms = GPU_CONTAINER_PLATFORMS if cs.endswith(GPU_CONTAINER_SUFFIX) else CONTAINER_PLATFORMS + for p in platforms: + if cs in conda_fields: + conda_entry = getattr(containers_valid, cs).get(p) if conda_entry: containers_flat.append((cs, p, conda_entry.lock_file)) else: entry = getattr(containers_valid, cs).get(p) if entry: containers_flat.append((cs, p, entry.name)) - if cs == "singularity" and entry.https: + if cs.startswith("singularity") and entry.https: containers_flat.append((cs, p, entry.https)) return containers_flat @@ -735,9 +879,9 @@ def update_containers_in_meta(self, module_lint: ModuleLint | None = None) -> No meta = read_meta_yml(self.meta_yml) meta_containers = meta.get("containers", {}) - # Remove stale entries for all known systems/platforms so old containers don't - # mix with new ones when a build partially fails. - for cs in CONTAINER_SYSTEMS + ["conda"]: + # Remove stale entries for all known systems/platforms (CPU and GPU) so old + # containers don't mix with new ones when a build partially fails. + for cs in ALL_CONTAINER_SYSTEMS + ALL_CONDA_FIELDS: for platform in CONTAINER_PLATFORMS: meta_containers.get(cs, {}).pop(platform, None) new_containers = self.containers.dump_for_meta_yml() diff --git a/nf_core/modules/modules_utils.py b/nf_core/modules/modules_utils.py index 6a3ed22604..f88bde768e 100644 --- a/nf_core/modules/modules_utils.py +++ b/nf_core/modules/modules_utils.py @@ -6,7 +6,11 @@ import requests from pydantic import BaseModel, ValidationInfo, field_validator, model_validator -from nf_core.utils import CONTAINER_PLATFORMS, CONTAINER_SYSTEMS, NFCORE_CACHE_DIR +from nf_core.utils import ( + CONTAINER_PLATFORMS, + CONTAINER_SYSTEMS, + NFCORE_CACHE_DIR, +) from ..components.nfcore_component import NFCoreComponent @@ -59,8 +63,14 @@ class MetaYmlContainers(BaseModel): docker: dict[str, ContainerEntry] = {} singularity: dict[str, ContainerEntry] = {} conda: dict[str, CondaEntry] = {} - - @field_validator("docker", "singularity", "conda", mode="after") + # GPU-capable modules (dual-container pattern) add a parallel set of containers built + # from environment.gpu.yml. These are only populated for GPU modules and stay empty + # (and are omitted from meta.yml) otherwise. + docker_gpu: dict[str, ContainerEntry] = {} + singularity_gpu: dict[str, ContainerEntry] = {} + conda_gpu: dict[str, CondaEntry] = {} + + @field_validator("docker", "singularity", "conda", "docker_gpu", "singularity_gpu", "conda_gpu", mode="after") @classmethod def check_keys_against_container_platforms(cls, value: dict) -> dict: # Check if all used keys are valid platforms @@ -75,7 +85,8 @@ def dump_for_meta_yml(self) -> dict: ``model_dump`` shaped for the meta.yml JSON schema: ``name``/``build_id`` are always emitted (the schema requires them, empty or not), while the optional fields (``https``, ``scan_id``) are omitted when empty — an empty ``https`` - would violate the schema's ``^https://`` pattern. + would violate the schema's ``^https://`` pattern. Empty container systems (e.g. the + GPU sections for non-GPU modules) are dropped entirely so they don't appear in meta.yml. """ dump = self.model_dump() for platforms in dump.values(): @@ -83,7 +94,7 @@ def dump_for_meta_yml(self) -> dict: for key in ("https", "scan_id"): if not entry.get(key, True): del entry[key] - return dump + return {system: platforms for system, platforms in dump.items() if platforms} @model_validator(mode="after") def check_container_systems_complete(self, info: ValidationInfo) -> "MetaYmlContainers": diff --git a/nf_core/utils.py b/nf_core/utils.py index 98db9aba31..e58364f579 100644 --- a/nf_core/utils.py +++ b/nf_core/utils.py @@ -110,6 +110,13 @@ def is_interactive() -> bool: CONTAINER_SYSTEMS = ["docker", "singularity"] CONTAINER_PLATFORMS = ["linux/amd64", "linux/arm64"] +# GPU-capable modules provide a second `environment.gpu.yml` and get a parallel set of +# containers stored under suffixed meta.yml keys (docker_gpu, singularity_gpu, conda_gpu). +# CUDA conda builds are effectively amd64-only, so the GPU path is not built for arm64. +GPU_CONTAINER_SUFFIX = "_gpu" +GPU_CONTAINER_SYSTEMS = [f"{cs}{GPU_CONTAINER_SUFFIX}" for cs in CONTAINER_SYSTEMS] +GPU_CONTAINER_PLATFORMS = ["linux/amd64"] + class ContainerRegistryUrls(Enum): SEQERA_DOCKER = "community.wave.seqera.io/library" diff --git a/tests/modules/test_containers.py b/tests/modules/test_containers.py index 5ba5c27d76..f3694b18d8 100644 --- a/tests/modules/test_containers.py +++ b/tests/modules/test_containers.py @@ -7,7 +7,7 @@ from nf_core.modules.containers import ModuleContainers from nf_core.modules.modules_utils import CondaEntry, ContainerEntry, MetaYmlContainers -from nf_core.utils import CONTAINER_PLATFORMS, CONTAINER_SYSTEMS +from nf_core.utils import CONTAINER_PLATFORMS, CONTAINER_SYSTEMS, GPU_CONTAINER_PLATFORMS from ..test_modules import TestModules @@ -308,6 +308,36 @@ def test_update_containers_in_meta_merges(self): meta = yaml.safe_load((self.bpipe_test_module_path / "meta.yml").read_text(encoding="utf-8")) assert meta["containers"] == containers.dump_for_meta_yml() + def test_build_targets_cpu_only(self): + """A module without environment.gpu.yml only builds CPU targets.""" + assert self.module_containers.environment_gpu_yml is None + targets = self.module_containers._build_targets() + assert len(targets) == len(CONTAINER_SYSTEMS) * len(CONTAINER_PLATFORMS) + assert all(not t.field.endswith("_gpu") for t in targets) + assert all(t.build_template == "conda/pixi:v1" for t in targets) + + def test_dump_for_meta_yml_includes_gpu_when_present(self): + """GPU sections are dumped when populated; empty ones are dropped.""" + amd64 = CONTAINER_PLATFORMS[0] + containers = MetaYmlContainers( + docker={p: ContainerEntry(name=f"d-{p}", build_id="b") for p in CONTAINER_PLATFORMS}, + docker_gpu={amd64: ContainerEntry(name="d-gpu", build_id="bg")}, + ) + dump = containers.dump_for_meta_yml() + assert dump["docker_gpu"][amd64]["name"] == "d-gpu" + # empty GPU sections must not leak into meta.yml + assert "singularity_gpu" not in dump + assert "conda_gpu" not in dump + + @mock.patch("nf_core.modules.containers.requests.post") + def test_request_container_passes_build_template(self, mock_post): + """request_container forwards the build template into the Wave payload.""" + mock_post.return_value = self._fake_response({"targetImage": "x:1", "buildId": "b", "cached": True}) + ModuleContainers.request_container( + "docker", CONTAINER_PLATFORMS[0], self.environment_yml, build_template="conda/micromamba:v2" + ) + assert mock_post.call_args[1]["json"]["buildTemplate"] == "conda/micromamba:v2" + class TestModuleContainersPipeline(TestModules): """Tests for ModuleContainers against a real pipeline repository""" @@ -381,3 +411,179 @@ def test_update_containers_in_meta_skips_dockerfile_module(self): original_meta = (self.module_dir / "meta.yml").read_text() self.module_containers.update_containers_in_meta() assert (self.module_dir / "meta.yml").read_text() == original_meta + + +class TestModuleContainersGPU(TestModules): + """Tests for GPU (dual-container) module handling in ModuleContainers""" + + def setUp(self): + super().setUp() + self.environment_yml = self.bpipe_test_module_path / "environment.yml" + # A GPU module carries a second environment.gpu.yml next to environment.yml. + # It must exist before ModuleContainers is constructed (detected in __init__). + (self.bpipe_test_module_path / "environment.gpu.yml").write_text( + "name: bpipe_test_gpu\nchannels:\n - conda-forge\ndependencies:\n - pytorch-gpu=2.1.0\n", + encoding="utf-8", + ) + self.module_containers = ModuleContainers("bpipe/test", directory=self.nfcore_modules) + + @staticmethod + def _fake_response(payload: dict | None = None, status_code: int = 200, text: str = "") -> mock.MagicMock: + resp = mock.MagicMock() + resp.status_code = status_code + resp.text = text + resp.json.return_value = payload or {} + return resp + + @staticmethod + def _fake_inspect(image: str) -> dict: + return { + "container": {"manifest": {"layers": [{"mediaType": "application/vnd.sif", "digest": "sha256:abcde12345"}]}} + } + + def test_detects_gpu_environment(self): + assert self.module_containers.environment_gpu_yml == self.bpipe_test_module_path / "environment.gpu.yml" + + def test_build_targets_includes_gpu(self): + targets = self.module_containers._build_targets() + cpu = [t for t in targets if not t.field.endswith("_gpu")] + gpu = [t for t in targets if t.field.endswith("_gpu")] + + # CPU: every system x every platform, built with pixi + assert len(cpu) == len(CONTAINER_SYSTEMS) * len(CONTAINER_PLATFORMS) + assert all(t.build_template == "conda/pixi:v1" for t in cpu) + + # GPU: every system x GPU platforms (amd64 only), built with micromamba v2 + assert {t.field for t in gpu} == {"docker_gpu", "singularity_gpu"} + assert {t.platform for t in gpu} == set(GPU_CONTAINER_PLATFORMS) + assert all(t.build_template == "conda/micromamba:v2" for t in gpu) + assert all(t.system in CONTAINER_SYSTEMS for t in gpu) + assert all(t.env_file == self.module_containers.environment_gpu_yml for t in gpu) + + @mock.patch("nf_core.modules.containers.requests.get") + @mock.patch.object(ModuleContainers, "request_image_inspect") + @mock.patch("nf_core.modules.containers.requests.post") + def test_create_builds_gpu_containers(self, mock_post, mock_inspect, mock_get): + submitted_templates = [] + + def fake_post(url, json=None, headers=None): + submitted_templates.append(json["buildTemplate"]) + system = "singularity" if json.get("format") == "sif" else "docker" + gpu = json["buildTemplate"] == "conda/micromamba:v2" + tag = "gpu" if gpu else "cpu" + image = f"community.wave.seqera.io/library/bpipe_test_{tag}:0.1.0--{tag}{system}" + meta = { + "buildId": f"bd-{tag}{system}", + "cached": True, + "containerImage": image, + "requestId": f"req-{tag}{system}", + "succeeded": True, + "targetImage": image, + } + if system == "docker": + meta["scanId"] = f"sc-{tag}{system}" + return self._fake_response(meta) + + mock_get.return_value = self._fake_response(text="# conda lock content") + mock_post.side_effect = fake_post + mock_inspect.side_effect = self._fake_inspect + + containers, success = self.module_containers.create() + assert success + + amd64 = GPU_CONTAINER_PLATFORMS[0] + # GPU containers built for amd64 only, stored under the *_gpu fields + assert set(containers.docker_gpu) == {amd64} + assert set(containers.singularity_gpu) == {amd64} + assert "gpu" in containers.docker_gpu[amd64].name + assert containers.docker_gpu[amd64].build_id == "bd-gpudocker" + assert containers.singularity_gpu[amd64].https # extracted from image inspect + # GPU conda lock recorded under conda_gpu (keyed off the docker_gpu build) + assert containers.conda_gpu[amd64].lock_file + + # CPU containers still built for all platforms + assert set(containers.docker) == set(CONTAINER_PLATFORMS) + + # Both build templates were used (pixi for CPU, micromamba v2 for GPU) + assert "conda/pixi:v1" in submitted_templates + assert "conda/micromamba:v2" in submitted_templates + + @mock.patch("nf_core.modules.containers.requests.get") + def test_get_conda_lock_file_gpu(self, mock_get): + """get_conda_lock_file(gpu=True) resolves the lock from the docker_gpu build.""" + mock_get.return_value = self._fake_response(text="# gpu conda lock") + amd64 = GPU_CONTAINER_PLATFORMS[0] + containers = MetaYmlContainers() + containers.docker_gpu[amd64] = ContainerEntry(name="x", build_id="gpu-build-1") + self.module_containers.containers = containers + + result = self.module_containers.get_conda_lock_file(amd64, gpu=True) + assert result == "# gpu conda lock" + assert mock_get.call_args[0][0] == "https://wave.seqera.io/v1alpha1/builds/gpu-build-1/condalock" + + def test_get_conda_lock_file_gpu_missing_build_id_raises(self): + """Without a docker_gpu build, the GPU conda lock cannot be resolved.""" + amd64 = GPU_CONTAINER_PLATFORMS[0] + self.module_containers.containers = MetaYmlContainers() + with pytest.raises(ValueError, match="No build_id found for docker_gpu container"): + self.module_containers.get_conda_lock_file(amd64, gpu=True) + + def test_update_main_nf_container_gpu_nested_ternary(self): + """A GPU module gets the nested task.accelerator ternary with CPU and GPU images.""" + gpu_amd64 = GPU_CONTAINER_PLATFORMS[0] + containers = MetaYmlContainers( + docker={p: ContainerEntry(name=f"docker-cpu-{p.split('/')[-1]}") for p in CONTAINER_PLATFORMS}, + singularity={ + p: ContainerEntry( + name=f"oras://sif-cpu-{p.split('/')[-1]}", https=f"https://sif-cpu-{p.split('/')[-1]}" + ) + for p in CONTAINER_PLATFORMS + }, + docker_gpu={gpu_amd64: ContainerEntry(name="docker-gpu-amd64")}, + singularity_gpu={gpu_amd64: ContainerEntry(name="oras://sif-gpu", https="https://sif-gpu-amd64")}, + ) + self.module_containers.containers = containers + self.module_containers.update_main_nf_container(force=True) + + content = self.module_containers.nfcore_component.main_nf.read_text() + # nested ternary: engine choice outer, task.accelerator inner + assert "workflow.containerEngine in ['singularity', 'apptainer']" in content + assert "(task.accelerator ? 'https://sif-gpu-amd64' : 'https://sif-cpu-amd64')" in content + assert "(task.accelerator ? 'docker-gpu-amd64' : 'docker-cpu-amd64')" in content + + def test_update_main_nf_container_cpu_only_when_no_gpu_containers(self): + """Without GPU containers the plain (non-accelerator) ternary is written.""" + containers = MetaYmlContainers( + docker={p: ContainerEntry(name=f"docker-cpu-{p.split('/')[-1]}") for p in CONTAINER_PLATFORMS}, + singularity={ + p: ContainerEntry( + name=f"oras://sif-cpu-{p.split('/')[-1]}", https=f"https://sif-cpu-{p.split('/')[-1]}" + ) + for p in CONTAINER_PLATFORMS + }, + ) + self.module_containers.containers = containers + self.module_containers.update_main_nf_container(force=True) + + content = self.module_containers.nfcore_component.main_nf.read_text() + assert "task.accelerator" not in content + assert "docker-cpu-amd64" in content + + def test_list_containers_includes_gpu(self): + amd64 = GPU_CONTAINER_PLATFORMS[0] + containers = MetaYmlContainers( + docker={p: ContainerEntry(name=f"d-{p}") for p in CONTAINER_PLATFORMS}, + singularity={p: ContainerEntry(name=f"s-{p}") for p in CONTAINER_PLATFORMS}, + docker_gpu={amd64: ContainerEntry(name="d-gpu")}, + singularity_gpu={amd64: ContainerEntry(name="s-gpu", https="https://s-gpu")}, + conda_gpu={amd64: CondaEntry(lock_file="/lock/gpu.txt")}, + ) + with mock.patch.object(self.module_containers, "get_containers_from_meta", return_value=containers): + listed = self.module_containers.list_containers() + + assert ("docker_gpu", amd64, "d-gpu") in listed + assert ("singularity_gpu", amd64, "s-gpu") in listed + assert ("singularity_gpu", amd64, "https://s-gpu") in listed + assert ("conda_gpu", amd64, "/lock/gpu.txt") in listed + # GPU sections are amd64-only: no arm64 GPU entries + assert not any(cs.endswith("_gpu") and p != amd64 for cs, p, _ in listed)