diff --git a/mkdocs/docs/concepts/dev-environments.md b/mkdocs/docs/concepts/dev-environments.md index e8ddb7045..0f99420b4 100644 --- a/mkdocs/docs/concepts/dev-environments.md +++ b/mkdocs/docs/concepts/dev-environments.md @@ -159,7 +159,8 @@ resources: The `cpu` property lets you set the architecture (`x86` or `arm`) and core count — e.g., `x86:16` (16 x86 cores), `arm:8..` (at least 8 ARM cores). -If not set, `dstack` infers it from the GPU or defaults to `x86`. +If the architecture is not set, `dstack` allows any architecture supported by the `image`, or `x86` if no `image` is set. +Since the default `dstack` image only supports `x86`, requesting `arm` requires setting `image` and is not compatible with `docker: true`. The `gpu` property lets you specify vendor, model, memory, and count — e.g., `nvidia` (one NVIDIA GPU), `A100` (one A100), `A10G,A100` (either), `A100:80GB` (one 80GB A100), `A100:2` (two A100), `24GB..40GB:2` (two GPUs with 24–40GB), `A100:40GB:2` (two 40GB A100s). diff --git a/mkdocs/docs/concepts/services.md b/mkdocs/docs/concepts/services.md index ea06d39be..7cd6034cb 100644 --- a/mkdocs/docs/concepts/services.md +++ b/mkdocs/docs/concepts/services.md @@ -879,7 +879,8 @@ resources: The `cpu` property lets you set the architecture (`x86` or `arm`) and core count — e.g., `x86:16` (16 x86 cores), `arm:8..` (at least 8 ARM cores). -If not set, `dstack` infers it from the GPU or defaults to `x86`. +If the architecture is not set, `dstack` allows any architecture supported by the `image`, or `x86` if no `image` is set. +Since the default `dstack` image only supports `x86`, requesting `arm` requires setting `image` and is not compatible with `docker: true`. The `gpu` property lets you specify vendor, model, memory, and count — e.g., `nvidia` (one NVIDIA GPU), `A100` (one A100), `A10G,A100` (either), `A100:80GB` (one 80GB A100), `A100:2` (two A100), `24GB..40GB:2` (two GPUs with 24–40GB), `A100:40GB:2` (two 40GB A100s). diff --git a/mkdocs/docs/concepts/tasks.md b/mkdocs/docs/concepts/tasks.md index 6592e431b..ea1715aa8 100644 --- a/mkdocs/docs/concepts/tasks.md +++ b/mkdocs/docs/concepts/tasks.md @@ -220,7 +220,8 @@ resources: The `cpu` property lets you set the architecture (`x86` or `arm`) and core count — e.g., `x86:16` (16 x86 cores), `arm:8..` (at least 8 ARM cores). -If not set, `dstack` infers it from the GPU or defaults to `x86`. +If the architecture is not set, `dstack` allows any architecture supported by the `image`, or `x86` if no `image` is set. +Since the default `dstack` image only supports `x86`, requesting `arm` requires setting `image` and is not compatible with `docker: true`. The `gpu` property lets you specify vendor, model, memory, and count — e.g., `nvidia` (one NVIDIA GPU), `A100` (one A100), `A10G,A100` (either), `A100:80GB` (one 80GB A100), `A100:2` (two A100), `24GB..40GB:2` (two GPUs with 24–40GB), `A100:40GB:2` (two 40GB A100s). diff --git a/src/dstack/_internal/cli/commands/offer.py b/src/dstack/_internal/cli/commands/offer.py index 6e52f4cfc..c884a178b 100644 --- a/src/dstack/_internal/cli/commands/offer.py +++ b/src/dstack/_internal/cli/commands/offer.py @@ -154,9 +154,13 @@ def _process_group_by_args(group_by_args: List[str]) -> List[str]: def _get_run_spec(args: argparse.Namespace) -> RunSpec: - # Set image and user so that the server (a) does not default gpu.vendor - # to nvidia — `dstack offer` should show all vendors, and (b) does not - # attempt to pull image config from the Docker registry. + # image="scratch" is a special value that forces the server to use some dummy default + # values for optional fields that otherwise would be extracted from the image config + # pulled from the image registry (commands/entrypoint, user, resources.cpu.arch). + # Additionally, it disables the server code path that sets gpu.vendor to nvidia when + # the image is not set. + # We still set `commands` and `user` for compatibility with older servers that don't treat + # "scratch" as a special "don't request the registry" value. conf = TaskConfiguration( resources=ResourcesSpec.unconstrained(), commands=[":"], diff --git a/src/dstack/_internal/server/background/pipeline_tasks/runs/active.py b/src/dstack/_internal/server/background/pipeline_tasks/runs/active.py index 1ad52e05f..ef47078bb 100644 --- a/src/dstack/_internal/server/background/pipeline_tasks/runs/active.py +++ b/src/dstack/_internal/server/background/pipeline_tasks/runs/active.py @@ -31,6 +31,7 @@ get_job_specs_from_run_spec, get_jobs_from_run_spec, group_jobs_by_replica_latest, + job_spec_updatable_in_place, ) from dstack._internal.server.services.runs import create_job_model_for_new_submission from dstack._internal.server.services.runs.replicas import ( @@ -515,7 +516,7 @@ async def _build_deployment_update_map( can_update_all_jobs = True for old_job_model, new_job_spec in zip(job_models, new_job_specs): old_job_spec = get_job_spec(old_job_model) - if new_job_spec != old_job_spec: + if not job_spec_updatable_in_place(old_job_spec, new_job_spec): can_update_all_jobs = False break if can_update_all_jobs: diff --git a/src/dstack/_internal/server/services/docker.py b/src/dstack/_internal/server/services/docker.py index 70c3c458a..8518f3001 100644 --- a/src/dstack/_internal/server/services/docker.py +++ b/src/dstack/_internal/server/services/docker.py @@ -1,7 +1,9 @@ +import contextlib import re from dataclasses import dataclass from typing import List, Optional +import gpuhunt import requests from dxf import DXF from dxf.exceptions import DXFError @@ -18,7 +20,6 @@ parse_image_name, ) -DEFAULT_PLATFORM = "linux/amd64" MAX_CONFIG_OBJECT_SIZE = 2**22 # 4 MiB REGISTRY_REQUEST_TIMEOUT = 20 @@ -50,6 +51,8 @@ def normalize_user(cls, v: Optional[str]) -> Optional[str]: class ImageConfigObject(CoreModel): + architecture: str + os: str config: ImageConfig = ImageConfig() @field_validator("config", mode="before") @@ -66,7 +69,9 @@ class ImageManifest(CoreModel): config: ImageManifestConfigField -def get_image_config(image_name: str, registry_auth: Optional[RegistryAuth]) -> ImageConfigObject: +def get_image_config_and_cpu_architectures( + image_name: str, registry_auth: Optional[RegistryAuth] +) -> tuple[ImageConfigObject, set[gpuhunt.CPUArchitecture]]: image = parse_image_name(image_name) registry = image.registry @@ -81,21 +86,54 @@ def get_image_config(image_name: str, registry_auth: Optional[RegistryAuth]) -> ) with registry_client: + cpu_architectures: Optional[set[gpuhunt.CPUArchitecture]] = None try: - manifest_resp = registry_client.get_manifest( - alias=image.digest or image.tag, platform=DEFAULT_PLATFORM - ) - assert isinstance(manifest_resp, str), ( - "get_manifest() returns the manifest JSON when `platform` is given" - ) - manifest = validate_json_extra_ignore(ImageManifest, manifest_resp) + # FIXME: get_manifest() makes N+1 requests when platform is not specified and alias + # points to an image index, where N is a number of images in the index, + # e.g., debian has 8 os/architecture[/variant] combinations + manifest_resp = registry_client.get_manifest(alias=image.digest or image.tag) + if isinstance(manifest_resp, dict): + # Image index (OCI) aka Manifest list (Docker) -- multi os/arch higher-level object + manifests: dict[gpuhunt.CPUArchitecture, ImageManifest] = {} + for platform, manifest_raw in manifest_resp.items(): + # os/architecture[/variant] + os_name, architecture, *_ = platform.split("/") + if not _os_supported(os_name): + continue + cpu_arch = _cpu_arch_from_string(architecture) + if cpu_arch is not None: + manifests[cpu_arch] = validate_json_extra_ignore( + ImageManifest, manifest_raw + ) + # ImageConfigs (User/Cmd/Entrypoint) may be different for different images + # within the same index; we assume that it's not the case but at least pick + # the manifest deterministically + for cpu_arch in [gpuhunt.CPUArchitecture.X86, gpuhunt.CPUArchitecture.ARM]: + with contextlib.suppress(KeyError): + manifest = manifests[cpu_arch] + break + else: + raise _no_supported_platforms_error(image_name) + cpu_architectures = set(manifests) + else: + # Image manifest -- one specific os/arch combination + manifest = validate_json_extra_ignore(ImageManifest, manifest_resp) + config_stream = registry_client.pull_blob(manifest.config.digest) config_resp = join_byte_stream_checked(config_stream, MAX_CONFIG_OBJECT_SIZE) # type: ignore[arg-type] if config_resp is None: raise DockerRegistryError( f"Image config object exceeds the size limit of {MAX_CONFIG_OBJECT_SIZE} bytes" ) - return validate_json_extra_ignore(ImageConfigObject, config_resp) + image_config = validate_json_extra_ignore(ImageConfigObject, config_resp) + + if cpu_architectures is None: + cpu_arch = _cpu_arch_from_string(image_config.architecture) + if not _os_supported(image_config.os) or cpu_arch is None: + raise _no_supported_platforms_error(image_name) + cpu_architectures = {cpu_arch} + + return image_config, cpu_architectures except (DXFError, requests.RequestException, ValidationError) as e: raise DockerRegistryError(e) @@ -130,3 +168,19 @@ def is_valid_docker_volume_target(path: str) -> bool: if path.endswith("/") and path != "/": return False return DOCKER_TARGET_PATH_PATTERN.match(path) is not None + + +def _cpu_arch_from_string(architecture: str) -> Optional[gpuhunt.CPUArchitecture]: + if architecture == "amd64": + return gpuhunt.CPUArchitecture.X86 + if architecture == "arm64": + return gpuhunt.CPUArchitecture.ARM + return None + + +def _os_supported(os_name: str) -> bool: + return os_name == "linux" + + +def _no_supported_platforms_error(image_name: str) -> DockerRegistryError: + return DockerRegistryError(f"No supported OS/architectures found: {image_name!r}") diff --git a/src/dstack/_internal/server/services/fleets.py b/src/dstack/_internal/server/services/fleets.py index ba0eae4b3..ac677b494 100644 --- a/src/dstack/_internal/server/services/fleets.py +++ b/src/dstack/_internal/server/services/fleets.py @@ -87,10 +87,7 @@ list_user_project_models, project_model_to_project, ) -from dstack._internal.server.services.resources import ( - set_default_cpu_spec_arch, - set_default_gpu_spec, -) +from dstack._internal.server.services.resources import set_default_gpu_spec from dstack._internal.utils import random_names from dstack._internal.utils import ssh as ssh_utils from dstack._internal.utils.common import ( @@ -1429,8 +1426,7 @@ def _validate_fleet_configuration_subtype_specific_fields(conf: FleetConfigurati def _set_fleet_spec_defaults(spec: FleetSpec): resources_spec = spec.configuration.resources if resources_spec is not None: - gpu_spec = set_default_gpu_spec(resources_spec) - set_default_cpu_spec_arch(resources_spec.cpu, gpu_spec) + set_default_gpu_spec(resources_spec) def _validate_all_ssh_params_specified(ssh_config: SSHParams): diff --git a/src/dstack/_internal/server/services/jobs/__init__.py b/src/dstack/_internal/server/services/jobs/__init__.py index 728b0df55..e9f8eeca2 100644 --- a/src/dstack/_internal/server/services/jobs/__init__.py +++ b/src/dstack/_internal/server/services/jobs/__init__.py @@ -298,6 +298,24 @@ def get_job_spec(job_model: JobModel) -> JobSpec: return validate_json_extra_ignore(JobSpec, job_model.job_spec_data) +def job_spec_updatable_in_place(old_job_spec: JobSpec, new_job_spec: JobSpec) -> bool: + """ + Check if a job running with `old_job_spec` already satisfies `new_job_spec`, that is, + the job can be marked as up-to-date without redeployment. + """ + if old_job_spec == new_job_spec: + return True + # Older servers always resolved `cpu.arch` to a specific value. Now an unset `arch` means + # "any architecture supported by the image", so a specific value -> None change only widens + # the requirements -- an already provisioned job still satisfies them. Without this check, + # re-applying an unchanged configuration after a server upgrade would trigger redeployment. + if new_job_spec.requirements.resources.cpu.arch is not None: + return False + new_job_spec = new_job_spec.model_copy(deep=True) + new_job_spec.requirements.resources.cpu.arch = old_job_spec.requirements.resources.cpu.arch + return old_job_spec == new_job_spec + + def delay_job_instance_termination(job_model: JobModel): job_model.remove_at = common.get_current_datetime() + timedelta(seconds=15) diff --git a/src/dstack/_internal/server/services/jobs/configurators/base.py b/src/dstack/_internal/server/services/jobs/configurators/base.py index 5ec38790d..7b4b711ef 100644 --- a/src/dstack/_internal/server/services/jobs/configurators/base.py +++ b/src/dstack/_internal/server/services/jobs/configurators/base.py @@ -6,6 +6,7 @@ from pathlib import PurePosixPath from typing import Dict, List, Optional +import gpuhunt from cachetools import TTLCache, cached from dstack._internal import settings @@ -55,7 +56,7 @@ from dstack._internal.server.services.docker import ( ImageConfig, apply_server_docker_defaults, - get_image_config, + get_image_config_and_cpu_architectures, ) from dstack._internal.utils import crypto from dstack._internal.utils.common import run_async @@ -69,6 +70,19 @@ DSTACK_DIR = "/dstack" DSTACK_PROFILE_PATH = f"{DSTACK_DIR}/profile" +# A non-existent image name used to signal that the image registry must never be requested +# and some dummy defaults should be used instead. +# As a job with such an image cannot be started, this special value only makes sense +# when used for offer collection (via `/runs/get_plan` with `for_offers_only`), not +# regular run planning/submission. +# Specifying a single "magic" value is still hacky but better than requiring clients to set +# an ever-growing list of optional configuration fields such as `commands`/`entrypoint`, +# `user`, `resources.cpu.arch`. +# In addition, it has a special effect on `resources.cpu.arch` -- unlike unset image, +# which defaults the arch to x86-only (as the default dstack image doesn't support ARM), +# this dummy image leaves the arch unset. +DUMMY_IMAGE_NAME = "scratch" + def get_default_python_verison() -> str: version_info = sys.version_info @@ -98,6 +112,7 @@ class JobConfigurator(ABC): TYPE: RunConfigurationType _image_config: Optional[ImageConfig] = None + _image_cpu_architectures: Optional[set[gpuhunt.CPUArchitecture]] = None # JobSSHKey should be shared for all jobs in a replica for inter-node communication. _job_ssh_key: Optional[JobSSHKey] = None @@ -139,8 +154,17 @@ def _ports(self) -> List[PortMapping]: pass async def _get_image_config(self) -> ImageConfig: + image_config, _ = await self._get_image_config_and_cpu_architectures() + return image_config + + async def _get_image_config_and_cpu_architectures( + self, + ) -> tuple[ImageConfig, set[gpuhunt.CPUArchitecture]]: if self._image_config is not None: - return self._image_config + assert self._image_cpu_architectures is not None + return self._image_config, self._image_cpu_architectures + image_name = self._image_name() + assert image_name != DUMMY_IMAGE_NAME interpolate = VariablesInterpolator({"secrets": self.secrets}).interpolate_or_error registry_auth = self.run_spec.configuration.registry_auth if registry_auth is not None: @@ -151,14 +175,15 @@ async def _get_image_config(self) -> ImageConfig: ) except InterpolatorError as e: raise ServerClientError(e.args[0]) - image_name, registry_auth = apply_server_docker_defaults(self._image_name(), registry_auth) - image_config = await run_async( - _get_image_config, + image_name, registry_auth = apply_server_docker_defaults(image_name, registry_auth) + image_config, cpu_architectures = await run_async( + _get_image_config_and_cpu_architectures, image_name, registry_auth, ) self._image_config = image_config - return image_config + self._image_cpu_architectures = cpu_architectures + return image_config, cpu_architectures async def _get_job_spec( self, @@ -184,7 +209,7 @@ async def _get_job_spec( stop_duration=self._stop_duration(), utilization_policy=self._utilization_policy(), registry_auth=self._registry_auth(), - requirements=self._requirements(jobs_per_replica), + requirements=await self._requirements(jobs_per_replica), retry=self._retry(), working_dir=self._working_dir(), volumes=self._volumes(job_num), @@ -219,6 +244,9 @@ async def _commands(self) -> List[str]: entrypoint = [self._shell(), "-i", "-c"] dstack_image_commands = self._dstack_image_commands() commands = [_join_shell_commands(dstack_image_commands + shell_commands)] + elif self._image_name() == DUMMY_IMAGE_NAME: + entrypoint = [] + commands = [":"] else: # custom docker image without commands image_config = await self._get_image_config() entrypoint = image_config.entrypoint or [] @@ -299,6 +327,8 @@ def _image_name(self) -> str: async def _user(self) -> Optional[UnixUser]: user = self.run_spec.configuration.user if user is None and self.run_spec.configuration.image is not None: + if self.run_spec.configuration.image == DUMMY_IMAGE_NAME: + return None image_config = await self._get_image_config() user = image_config.user if user is None: @@ -335,13 +365,29 @@ def _utilization_policy(self) -> Optional[UtilizationPolicy]: def _registry_auth(self) -> Optional[RegistryAuth]: return self.run_spec.configuration.registry_auth - def _requirements(self, jobs_per_replica: int) -> Requirements: + async def _requirements(self, jobs_per_replica: int) -> Requirements: resources = self.run_spec.configuration.resources + image = self.run_spec.configuration.image if self.run_spec.configuration.type == "service": for group in self.run_spec.configuration.replica_groups: if group.name == self.replica_group_name: resources = group.resources + if group.image is not None: + image = group.image break + resources = resources.model_copy(deep=True) + if resources.cpu.arch is None and image != DUMMY_IMAGE_NAME: + if image is None: + # dstackai/base or dstackai/dind image, both don't support ARM + resources.cpu.arch = gpuhunt.CPUArchitecture.X86 + else: + _, cpu_architectures = await self._get_image_config_and_cpu_architectures() + if len(cpu_architectures) == 1: + resources.cpu.arch = next(iter(cpu_architectures)) + # len(cpu_architectures) > 1 => multi-arch image, keep CPUSpec.arch unset. + # In the requirements, unset arch means "any architecture supported by the + # image", unlike the run configuration, where unset arch means "not specified, + # resolve it here" spot_policy = self._spot_policy() return Requirements( resources=resources, @@ -514,10 +560,15 @@ def _join_shell_commands(commands: List[str]) -> str: cache=TTLCache(maxsize=2048, ttl=80), lock=threading.Lock(), ) -def _get_image_config(image: str, registry_auth: Optional[RegistryAuth]) -> ImageConfig: +def _get_image_config_and_cpu_architectures( + image: str, registry_auth: Optional[RegistryAuth] +) -> tuple[ImageConfig, set[gpuhunt.CPUArchitecture]]: try: - return get_image_config(image, registry_auth).config + image_config, cpu_architectures = get_image_config_and_cpu_architectures( + image, registry_auth + ) except DockerRegistryError as e: raise ServerClientError( f"Error pulling configuration for image {image!r} from the docker registry: {e}" ) + return image_config.config, cpu_architectures diff --git a/src/dstack/_internal/server/services/jobs/configurators/service.py b/src/dstack/_internal/server/services/jobs/configurators/service.py index 45bc4c8f7..4150edbf4 100644 --- a/src/dstack/_internal/server/services/jobs/configurators/service.py +++ b/src/dstack/_internal/server/services/jobs/configurators/service.py @@ -9,6 +9,7 @@ from dstack._internal.core.models.profiles import SpotPolicy from dstack._internal.core.models.unix import UnixUser from dstack._internal.server.services.jobs.configurators.base import ( + DUMMY_IMAGE_NAME, JobConfigurator, get_default_image, ) @@ -94,6 +95,8 @@ async def _user(self) -> Optional[UnixUser]: if self.run_spec.configuration.user is None: group = self._current_replica_group() if group is not None and group.image is not None: + if group.image == DUMMY_IMAGE_NAME: + return None image_config = await self._get_image_config() if image_config.user is None: return None diff --git a/src/dstack/_internal/server/services/resources.py b/src/dstack/_internal/server/services/resources.py index 12a547e24..221f944bf 100644 --- a/src/dstack/_internal/server/services/resources.py +++ b/src/dstack/_internal/server/services/resources.py @@ -4,7 +4,6 @@ from dstack._internal.core.models.resources import ( DEFAULT_GPU_SPEC, - CPUSpec, GPUSpec, ResourcesSpec, ) @@ -17,18 +16,6 @@ def set_default_gpu_spec(resources_spec: ResourcesSpec) -> GPUSpec: return resources_spec.gpu -def set_default_cpu_spec_arch(cpu_spec: CPUSpec, gpu_spec: GPUSpec) -> None: - if cpu_spec.arch is None: - if ( - gpu_spec.vendor in [None, gpuhunt.AcceleratorVendor.NVIDIA] - and gpu_spec.name - and any(map(gpuhunt.is_nvidia_superchip, gpu_spec.name)) - ): - cpu_spec.arch = gpuhunt.CPUArchitecture.ARM - else: - cpu_spec.arch = gpuhunt.CPUArchitecture.X86 - - def set_default_gpu_spec_vendor( gpu_spec: GPUSpec, image: Optional[str], diff --git a/src/dstack/_internal/server/services/runs/spec.py b/src/dstack/_internal/server/services/runs/spec.py index e8eb53b45..08d9e00d5 100644 --- a/src/dstack/_internal/server/services/runs/spec.py +++ b/src/dstack/_internal/server/services/runs/spec.py @@ -21,7 +21,6 @@ from dstack._internal.server.models import UserModel from dstack._internal.server.services.docker import is_valid_docker_volume_target from dstack._internal.server.services.resources import ( - set_default_cpu_spec_arch, set_default_gpu_spec, set_default_gpu_spec_vendor, ) @@ -167,7 +166,6 @@ def _set_resources_defaults( resources_spec: ResourcesSpec, image: Optional[str], docker: Optional[bool] ) -> None: gpu_spec = set_default_gpu_spec(resources_spec) - set_default_cpu_spec_arch(cpu_spec=resources_spec.cpu, gpu_spec=gpu_spec) set_default_gpu_spec_vendor(gpu_spec=gpu_spec, image=image, docker=docker) @@ -447,6 +445,13 @@ def _check_can_update_configuration( ): # Allow switching between `https: ` and unset `https`. Has no effect. updatable_fields.append("https") + # Services allow updating any field in ResourcesSpec via rolling deployment. For other + # configuration types, only changes that an already provisioned job still satisfies + # are allowed. + if "resources" not in updatable_fields and _is_compatible_resources_update( + current.resources, new.resources + ): + updatable_fields.append("resources") diff = diff_models(current, new) changed_fields = list(diff.keys()) for key in changed_fields: @@ -455,3 +460,19 @@ def _check_can_update_configuration( f"Failed to update fields {changed_fields}. Can only update {updatable_fields}" ) return diff + + +def _is_compatible_resources_update(current: ResourcesSpec, new: ResourcesSpec) -> bool: + """Check if a job provisioned with the `current` resources still satisfies the `new` ones.""" + diff = diff_models(current, new) + if not diff: + return True + if set(diff) != {"cpu"}: + return False + if set(diff_models(current.cpu, new.cpu)) != {"arch"}: + return False + # Older servers always resolved `cpu.arch` to a specific value. Now an unset `arch` means + # "any architecture supported by the image", so a specific value -> None change only + # widens the requirements. Without this, re-applying an unchanged configuration after + # a server upgrade would be rejected as a non-updatable change. + return new.cpu.arch is None diff --git a/src/dstack/_internal/utils/common.py b/src/dstack/_internal/utils/common.py index a6c3828a6..33db805f0 100644 --- a/src/dstack/_internal/utils/common.py +++ b/src/dstack/_internal/utils/common.py @@ -161,8 +161,8 @@ def pretty_resources( cpu_arch_lower = str(cpu_arch.value).lower() elif isinstance(cpu_arch, str): cpu_arch_lower = cpu_arch.lower() - if cpu_arch_lower == "arm": - cpu_arch_prefix = "arm:" + if cpu_arch_lower is not None: + cpu_arch_prefix = f"{cpu_arch_lower}:" else: cpu_arch_prefix = "" parts.append(f"cpu={cpu_arch_prefix}{cpus}") diff --git a/src/tests/_internal/cli/utils/conftest.py b/src/tests/_internal/cli/utils/conftest.py index a374bf9d3..fb5f9dd14 100644 --- a/src/tests/_internal/cli/utils/conftest.py +++ b/src/tests/_internal/cli/utils/conftest.py @@ -1,5 +1,6 @@ from unittest.mock import Mock +import gpuhunt import pytest from dstack._internal.server.services.docker import ImageConfig, ImageConfigObject @@ -11,11 +12,17 @@ def image_config_mock(monkeypatch: pytest.MonkeyPatch) -> ImageConfig: {"User": None, "Entrypoint": None, "Cmd": ["/bin/bash"]} ) monkeypatch.setattr( - "dstack._internal.server.services.jobs.configurators.base._get_image_config", - Mock(return_value=image_config), + "dstack._internal.server.services.jobs.configurators.base" + "._get_image_config_and_cpu_architectures", + Mock(return_value=(image_config, {gpuhunt.CPUArchitecture.X86})), ) monkeypatch.setattr( - "dstack._internal.server.services.docker.get_image_config", - Mock(return_value=ImageConfigObject(config=image_config)), + "dstack._internal.server.services.docker.get_image_config_and_cpu_architectures", + Mock( + return_value=( + ImageConfigObject(architecture="amd64", os="linux", config=image_config), + {gpuhunt.CPUArchitecture.X86}, + ) + ), ) return image_config diff --git a/src/tests/_internal/server/background/pipeline_tasks/test_runs/test_active.py b/src/tests/_internal/server/background/pipeline_tasks/test_runs/test_active.py index ed085e09d..ba81a3975 100644 --- a/src/tests/_internal/server/background/pipeline_tasks/test_runs/test_active.py +++ b/src/tests/_internal/server/background/pipeline_tasks/test_runs/test_active.py @@ -3,6 +3,7 @@ from datetime import timedelta from unittest.mock import AsyncMock, patch +import gpuhunt import pytest from sqlalchemy import select from sqlalchemy.ext.asyncio import AsyncSession @@ -946,6 +947,71 @@ async def test_service_in_place_deployment_bump( await session.refresh(job) assert job.deployment_num == 1 + async def test_service_in_place_deployment_bump_on_cpu_arch_widening( + self, test_db, session: AsyncSession, worker: RunWorker, image_config_mock + ) -> None: + """ + A replica submitted when `cpu.arch` was always resolved to a specific value is not + redeployed after the arch becomes unset (`any arch supported by the image`). + """ + project = await create_project(session=session) + user = await create_user(session=session) + repo = await create_repo(session=session, project_id=project.id) + run_spec = get_run_spec( + repo_id=repo.name, + run_name="service-run", + configuration=ServiceConfiguration( + port=8080, + image="debian", + commands=["echo Hi!"], + ), + ) + run = await create_run( + session=session, + project=project, + repo=repo, + user=user, + run_name="service-run", + run_spec=run_spec, + status=RunStatus.RUNNING, + deployment_num=1, + ) + # `image_config_mock` reports a single-arch image, so the job spec gets `arch: x86` + job = await create_job( + session=session, + run=run, + status=JobStatus.RUNNING, + deployment_num=0, + registered=True, + ready=True, + ) + assert get_job_spec(job).requirements.resources.cpu.arch == gpuhunt.CPUArchitecture.X86 + lock_run(run) + await session.commit() + + # The same image now reports both architectures, so the new job spec leaves `arch` unset + with patch( + "dstack._internal.server.services.jobs.configurators.base" + "._get_image_config_and_cpu_architectures", + return_value=( + image_config_mock, + {gpuhunt.CPUArchitecture.X86, gpuhunt.CPUArchitecture.ARM}, + ), + ): + await worker.process(run_to_pipeline_item(run)) + + session.expire_all() + await session.refresh(run) + assert run.status == RunStatus.RUNNING + + res = await session.execute(select(JobModel).where(JobModel.run_id == run.id)) + jobs = list(res.scalars().all()) + # No surge replica created, the existing one is marked as up-to-date + assert len(jobs) == 1 + assert jobs[0].id == job.id + assert jobs[0].status == JobStatus.RUNNING + assert jobs[0].deployment_num == 1 + async def test_service_rolling_deployment_scale_up( self, test_db, session: AsyncSession, worker: RunWorker ) -> None: diff --git a/src/tests/_internal/server/background/pipeline_tasks/test_submitted_jobs.py b/src/tests/_internal/server/background/pipeline_tasks/test_submitted_jobs.py index 7358cb573..9b2bb2c39 100644 --- a/src/tests/_internal/server/background/pipeline_tasks/test_submitted_jobs.py +++ b/src/tests/_internal/server/background/pipeline_tasks/test_submitted_jobs.py @@ -4,6 +4,7 @@ from typing import cast from unittest.mock import AsyncMock, Mock, call, patch +import gpuhunt import pytest from sqlalchemy import select from sqlalchemy.ext.asyncio import AsyncSession @@ -2626,8 +2627,8 @@ async def test_interpolates_secrets_when_provisioning_new_capacity( fleet=fleet, run_spec=run_spec, ) - with patch.object(JobConfigurator, "_get_image_config") as m: - m.return_value = image_config_mock + with patch.object(JobConfigurator, "_get_image_config_and_cpu_architectures") as m: + m.return_value = (image_config_mock, {gpuhunt.CPUArchitecture.X86}) job = await create_job(session=session, run=run, instance_assigned=True) offer = get_instance_offer_with_availability(backend=BackendType.RUNPOD) @@ -2684,8 +2685,8 @@ async def test_terminates_job_when_secret_is_missing( fleet=fleet, run_spec=run_spec, ) - with patch.object(JobConfigurator, "_get_image_config") as m: - m.return_value = image_config_mock + with patch.object(JobConfigurator, "_get_image_config_and_cpu_architectures") as m: + m.return_value = (image_config_mock, {gpuhunt.CPUArchitecture.X86}) job = await create_job(session=session, run=run, instance_assigned=True) offer = get_instance_offer_with_availability(backend=BackendType.RUNPOD) diff --git a/src/tests/_internal/server/conftest.py b/src/tests/_internal/server/conftest.py index 5275b6be0..43ee83b73 100644 --- a/src/tests/_internal/server/conftest.py +++ b/src/tests/_internal/server/conftest.py @@ -2,6 +2,7 @@ from pathlib import Path from unittest.mock import AsyncMock, Mock, patch +import gpuhunt import httpx import pytest @@ -66,12 +67,18 @@ def image_config_mock(monkeypatch: pytest.MonkeyPatch) -> ImageConfig: {"User": None, "Entrypoint": None, "Cmd": ["/bin/bash"]} ) monkeypatch.setattr( - "dstack._internal.server.services.jobs.configurators.base._get_image_config", - Mock(return_value=image_config), + "dstack._internal.server.services.jobs.configurators.base" + "._get_image_config_and_cpu_architectures", + Mock(return_value=(image_config, {gpuhunt.CPUArchitecture.X86})), ) monkeypatch.setattr( - "dstack._internal.server.services.docker.get_image_config", - Mock(return_value=ImageConfigObject(config=image_config)), + "dstack._internal.server.services.docker.get_image_config_and_cpu_architectures", + Mock( + return_value=( + ImageConfigObject(architecture="amd64", os="linux", config=image_config), + {gpuhunt.CPUArchitecture.X86}, + ) + ), ) return image_config diff --git a/src/tests/_internal/server/services/jobs/configurators/test_service.py b/src/tests/_internal/server/services/jobs/configurators/test_service.py index 7f39facac..a4d24ddc5 100644 --- a/src/tests/_internal/server/services/jobs/configurators/test_service.py +++ b/src/tests/_internal/server/services/jobs/configurators/test_service.py @@ -1,5 +1,6 @@ from unittest.mock import Mock +import gpuhunt import pytest from dstack._internal import settings @@ -295,8 +296,9 @@ async def test_user_looks_up_group_image(self, monkeypatch: pytest.MonkeyPatch): """When a group sets its own `image`, _user() queries that image's config.""" image_config = ImageConfig.model_validate({"User": "nginx", "Entrypoint": None, "Cmd": []}) monkeypatch.setattr( - "dstack._internal.server.services.jobs.configurators.base._get_image_config", - Mock(return_value=image_config), + "dstack._internal.server.services.jobs.configurators.base" + "._get_image_config_and_cpu_architectures", + Mock(return_value=(image_config, {gpuhunt.CPUArchitecture.X86})), ) run_spec = _make_run_spec( replicas=[ @@ -316,7 +318,8 @@ async def test_user_does_not_lookup_for_group_docker(self, monkeypatch: pytest.M """`docker: true` should not trigger an image-config registry call.""" mock_get_image_config = Mock() monkeypatch.setattr( - "dstack._internal.server.services.jobs.configurators.base._get_image_config", + "dstack._internal.server.services.jobs.configurators.base" + "._get_image_config_and_cpu_architectures", mock_get_image_config, ) run_spec = _make_run_spec( diff --git a/src/tests/_internal/server/services/jobs/test_jobs.py b/src/tests/_internal/server/services/jobs/test_jobs.py index ce9b931db..089b5b54b 100644 --- a/src/tests/_internal/server/services/jobs/test_jobs.py +++ b/src/tests/_internal/server/services/jobs/test_jobs.py @@ -1,5 +1,6 @@ from unittest.mock import patch +import gpuhunt import pytest import dstack._internal.server.settings as server_settings @@ -7,9 +8,13 @@ from dstack._internal.core.models.configurations import TaskConfiguration from dstack._internal.core.models.profiles import Profile from dstack._internal.core.models.repos.local import LocalRunRepoData -from dstack._internal.core.models.runs import RunSpec +from dstack._internal.core.models.resources import ResourcesSpec +from dstack._internal.core.models.runs import JobSpec, RunSpec from dstack._internal.server.services.docker import ImageConfig -from dstack._internal.server.services.jobs import get_job_specs_from_run_spec +from dstack._internal.server.services.jobs import ( + get_job_specs_from_run_spec, + job_spec_updatable_in_place, +) @pytest.mark.parametrize( @@ -37,12 +42,24 @@ id="custom-image-with-user", ), pytest.param( - # Setting `commands` and `user` is a known hack that we advertised to some customers - # to avoid registry requests. + # `commands` and `user` cover the image config, but the registry is still requested + # to find out which CPU architectures the image supports. TaskConfiguration(image="ubuntu", commands=["sleep infinity"], user="root"), - 0, + 1, id="custom-image-with-commands-and-user", ), + pytest.param( + # Setting `commands`, `user`, and `resources.cpu.arch` is a known hack that we + # advertised to some customers to avoid registry requests. + TaskConfiguration( + image="ubuntu", + commands=["sleep infinity"], + user="root", + resources=ResourcesSpec.model_validate({"cpu": "x86:2"}), + ), + 0, + id="custom-image-with-commands-user-and-arch", + ), ], ) @pytest.mark.asyncio @@ -64,8 +81,9 @@ async def test_get_job_specs_from_run_spec_image_config_calls( ) fake_image_config = ImageConfig.model_validate({"Entrypoint": ["/bin/bash"]}) with patch( - "dstack._internal.server.services.jobs.configurators.base._get_image_config", - return_value=fake_image_config, + "dstack._internal.server.services.jobs.configurators.base" + "._get_image_config_and_cpu_architectures", + return_value=(fake_image_config, {gpuhunt.CPUArchitecture.X86}), ) as mock_get_image_config: await get_job_specs_from_run_spec(run_spec=run_spec, secrets={}, replica_num=0) assert mock_get_image_config.call_count == expected_calls @@ -85,8 +103,9 @@ async def test_get_image_config_uses_server_default_registry(monkeypatch) -> Non ) fake_image_config = ImageConfig.model_validate({"Entrypoint": ["/bin/bash"]}) with patch( - "dstack._internal.server.services.jobs.configurators.base._get_image_config", - return_value=fake_image_config, + "dstack._internal.server.services.jobs.configurators.base" + "._get_image_config_and_cpu_architectures", + return_value=(fake_image_config, {gpuhunt.CPUArchitecture.X86}), ) as mock_get_image_config: job_specs = await get_job_specs_from_run_spec(run_spec=run_spec, secrets={}, replica_num=0) mock_get_image_config.assert_called_once_with( @@ -99,3 +118,60 @@ async def test_get_image_config_uses_server_default_registry(monkeypatch) -> Non # especially the credentials, so as not to leak them in the API. assert job_specs[0].image_name == "ubuntu" assert job_specs[0].registry_auth is None + + +class TestJobSpecUpdatableInPlace: + def _job_spec(self, **requirements_overrides) -> JobSpec: + resources = {"cpu": {"count": 2}, **requirements_overrides} + return JobSpec( + job_num=0, + job_name="test-run-0-0", + commands=["sleep infinity"], + env={}, + image_name="ubuntu", + requirements={"resources": resources}, + ) + + def test_identical_specs(self): + assert job_spec_updatable_in_place(self._job_spec(), self._job_spec()) + + def test_unrelated_change(self): + old_job_spec = self._job_spec() + new_job_spec = self._job_spec() + new_job_spec.commands = ["sleep 10"] + + assert not job_spec_updatable_in_place(old_job_spec, new_job_spec) + + def test_arch_widened_to_any(self): + # A job submitted by an older server that always resolved `cpu.arch` + old_job_spec = self._job_spec(cpu={"arch": gpuhunt.CPUArchitecture.X86, "count": 2}) + new_job_spec = self._job_spec(cpu={"arch": None, "count": 2}) + + assert job_spec_updatable_in_place(old_job_spec, new_job_spec) + + def test_arch_widened_to_any_with_another_change(self): + old_job_spec = self._job_spec(cpu={"arch": gpuhunt.CPUArchitecture.X86, "count": 2}) + new_job_spec = self._job_spec(cpu={"arch": None, "count": 2}) + new_job_spec.commands = ["sleep 10"] + + assert not job_spec_updatable_in_place(old_job_spec, new_job_spec) + + def test_arch_narrowed_to_specific(self): + old_job_spec = self._job_spec(cpu={"arch": None, "count": 2}) + new_job_spec = self._job_spec(cpu={"arch": gpuhunt.CPUArchitecture.X86, "count": 2}) + + assert not job_spec_updatable_in_place(old_job_spec, new_job_spec) + + def test_arch_changed_to_another_specific(self): + old_job_spec = self._job_spec(cpu={"arch": gpuhunt.CPUArchitecture.X86, "count": 2}) + new_job_spec = self._job_spec(cpu={"arch": gpuhunt.CPUArchitecture.ARM, "count": 2}) + + assert not job_spec_updatable_in_place(old_job_spec, new_job_spec) + + def test_does_not_mutate_the_new_spec(self): + old_job_spec = self._job_spec(cpu={"arch": gpuhunt.CPUArchitecture.X86, "count": 2}) + new_job_spec = self._job_spec(cpu={"arch": None, "count": 2}) + + job_spec_updatable_in_place(old_job_spec, new_job_spec) + + assert new_job_spec.requirements.resources.cpu.arch is None diff --git a/src/tests/_internal/server/services/runs/test_spec.py b/src/tests/_internal/server/services/runs/test_spec.py index 1d23a469e..a73229676 100644 --- a/src/tests/_internal/server/services/runs/test_spec.py +++ b/src/tests/_internal/server/services/runs/test_spec.py @@ -279,6 +279,52 @@ def test_non_dynamo_image_change_passes_configuration_gate(self): _check_can_update_configuration(current, new, ignore_files=True) +class TestCheckCanUpdateRunSpecResources: + """Non-service configurations cannot be redeployed, so only compatible changes are allowed.""" + + def test_allows_relaxing_cpu_arch(self): + # Older servers always resolved `cpu.arch`, so re-applying an unchanged configuration + # after a server upgrade must not be rejected + current = _task_run_spec(resources={"cpu": "x86:2"}, image="ubuntu") + new = _task_run_spec(resources={"cpu": 2}, image="ubuntu") + + check_can_update_run_spec(current, new) + + def test_allows_unchanged_resources(self): + current = _task_run_spec(resources={"cpu": "x86:2"}, image="ubuntu") + new = _task_run_spec(resources={"cpu": "x86:2"}, image="ubuntu") + + check_can_update_run_spec(current, new) + + def test_rejects_setting_cpu_arch(self): + current = _task_run_spec(resources={"cpu": 2}, image="ubuntu") + new = _task_run_spec(resources={"cpu": "x86:2"}, image="ubuntu") + + with pytest.raises(ServerClientError, match="resources"): + check_can_update_run_spec(current, new) + + def test_rejects_changing_cpu_arch(self): + current = _task_run_spec(resources={"cpu": "x86:2"}, image="ubuntu") + new = _task_run_spec(resources={"cpu": "arm:2"}, image="ubuntu") + + with pytest.raises(ServerClientError, match="resources"): + check_can_update_run_spec(current, new) + + def test_rejects_changing_cpu_count(self): + current = _task_run_spec(resources={"cpu": "x86:2"}, image="ubuntu") + new = _task_run_spec(resources={"cpu": 4}, image="ubuntu") + + with pytest.raises(ServerClientError, match="resources"): + check_can_update_run_spec(current, new) + + def test_rejects_changing_other_resources(self): + current = _task_run_spec(resources={"cpu": "x86:2", "memory": "8GB"}, image="ubuntu") + new = _task_run_spec(resources={"cpu": 2, "memory": "16GB"}, image="ubuntu") + + with pytest.raises(ServerClientError, match="resources"): + check_can_update_run_spec(current, new) + + class TestSetRunSpecResourcesDefaultsGpuVendor: @pytest.mark.parametrize( ["gpu_spec", "expected_vendor"], @@ -367,34 +413,6 @@ def test_sets_default_gpu_spec_if_gpu_is_null(self): assert gpu_spec.vendor == gpuhunt.AcceleratorVendor.NVIDIA -class TestSetRunSpecResourcesDefaultsCpuArch: - @pytest.mark.parametrize( - ["gpu_spec", "expected_arch"], - [ - (None, gpuhunt.CPUArchitecture.X86), - ("H100", gpuhunt.CPUArchitecture.X86), - ("GH200", gpuhunt.CPUArchitecture.ARM), # an NVIDIA superchip - ("GB200:4", gpuhunt.CPUArchitecture.ARM), - ], - ) - def test_sets_arch_detected_by_gpu_names( - self, gpu_spec: Optional[str], expected_arch: gpuhunt.CPUArchitecture - ): - resources = {"gpu": gpu_spec} if gpu_spec is not None else None - run_spec = _task_run_spec(resources=resources, image="ubuntu") - - set_run_spec_resources_defaults(run_spec) - - assert run_spec.configuration.resources.cpu.arch == expected_arch - - def test_does_not_override_arch_set_by_the_user(self): - run_spec = _task_run_spec(resources={"cpu": "arm:2", "gpu": "H100"}, image="ubuntu") - - set_run_spec_resources_defaults(run_spec) - - assert run_spec.configuration.resources.cpu.arch == gpuhunt.CPUArchitecture.ARM - - class TestSetRunSpecResourcesDefaultsReplicaGroups: def test_sets_defaults_for_every_replica_group(self): run_spec = _service_run_spec( @@ -408,10 +426,10 @@ def test_sets_defaults_for_every_replica_group(self): set_run_spec_resources_defaults(run_spec) groups = run_spec.configuration.replicas - assert [(g.resources.gpu.vendor, g.resources.cpu.arch) for g in groups] == [ - (gpuhunt.AcceleratorVendor.AMD, gpuhunt.CPUArchitecture.X86), - (gpuhunt.AcceleratorVendor.NVIDIA, gpuhunt.CPUArchitecture.ARM), - (gpuhunt.AcceleratorVendor.NVIDIA, gpuhunt.CPUArchitecture.X86), + assert [g.resources.gpu.vendor for g in groups] == [ + gpuhunt.AcceleratorVendor.AMD, + gpuhunt.AcceleratorVendor.NVIDIA, + gpuhunt.AcceleratorVendor.NVIDIA, ] @pytest.mark.parametrize( @@ -449,7 +467,6 @@ def test_sets_defaults_for_top_level_resources(self): resources = run_spec.configuration.resources assert resources.gpu.vendor == gpuhunt.AcceleratorVendor.NVIDIA - assert resources.cpu.arch == gpuhunt.CPUArchitecture.X86 class TestValidateRunSpecGpuVendorAndImage: @@ -507,16 +524,17 @@ def test_allows_replica_group_with_its_own_image(self): class TestValidateRunSpecCpuArchAndImage: - @pytest.mark.parametrize( - "resources", - [ - {"cpu": "arm:2"}, # the arch is set by the user - {"gpu": "GH200"}, # the arch is inferred from the GPU name - ], - ) - def test_rejects_arm_without_image(self, resources: dict): + # NOTE: only an explicitly requested ARM arch is validated. The arch is not inferred from + # the GPU name, as the actual arch is only known once an offer is selected -- the same + # fleet may provide both ARM (e.g., GH200) and x86 (e.g., H200) instances. + def test_rejects_arm_without_image(self): with pytest.raises(ServerClientError, match="`image` must be set when ARM CPU requested"): - _validate(_task_run_spec(resources=resources)) + _validate(_task_run_spec(resources={"cpu": "arm:2"})) + + def test_allows_arm_gpu_without_image(self): + # The run gets `arch: x86` requirements (the default image is x86-only) and is expected + # to find no offers rather than to be rejected + _validate(_task_run_spec(resources={"gpu": "GH200"})) def test_allows_arm_with_image(self): _validate(_task_run_spec(resources={"cpu": "arm:2"}, image="ubuntu")) @@ -535,7 +553,7 @@ def test_reports_replica_groups_requiring_image(self): replicas=[ {"count": 1, "commands": ["echo"], "resources": {"cpu": "arm:2"}}, {"count": 1, "commands": ["echo"]}, - {"count": 1, "commands": ["echo"], "resources": {"gpu": "GH200"}}, + {"count": 1, "commands": ["echo"], "resources": {"cpu": "arm:4"}}, ], ) diff --git a/src/tests/_internal/server/services/test_docker.py b/src/tests/_internal/server/services/test_docker.py index 14720b08a..2f80a66fb 100644 --- a/src/tests/_internal/server/services/test_docker.py +++ b/src/tests/_internal/server/services/test_docker.py @@ -1,11 +1,19 @@ +import json +from typing import Any, Union +from unittest.mock import MagicMock, patch + +import gpuhunt import pytest import dstack._internal.server.settings as server_settings +from dstack._internal.core.errors import DockerRegistryError from dstack._internal.core.models.common import RegistryAuth, validate_extra_ignore +from dstack._internal.server.services import docker as docker_services from dstack._internal.server.services.docker import ( ImageConfigObject, ImageManifest, apply_server_docker_defaults, + get_image_config_and_cpu_architectures, is_valid_docker_volume_target, ) @@ -266,3 +274,125 @@ def test_invalid_paths(self, path): def test_trailing_slash(self): assert not is_valid_docker_volume_target("/invalid/path/") + + +def _image_manifest(digest: str) -> str: + return json.dumps({"config": {"digest": digest, "size": 7023}}) + + +class TestGetImageConfigAndCpuArchitectures: + """ + `get_manifest()` returns a dict of platform -> manifest JSON for an image index + (multi-platform image) and a manifest JSON string for a single-platform image. + """ + + def _get_image_config( + self, + manifest_resp: Union[str, dict[str, str]], + config_object: dict[str, Any], + ) -> tuple[ImageConfigObject, set[gpuhunt.CPUArchitecture], MagicMock]: + registry_client = MagicMock() + registry_client.__enter__.return_value = registry_client + registry_client.get_manifest.return_value = manifest_resp + registry_client.pull_blob.return_value = [json.dumps(config_object).encode()] + with patch.object(docker_services, "DXF", return_value=registry_client): + image_config, cpu_architectures = get_image_config_and_cpu_architectures( + "debian", None + ) + return image_config, cpu_architectures, registry_client + + def test_index_reports_all_supported_architectures(self, sample_image_config_object): + image_config, cpu_architectures, _ = self._get_image_config( + { + "linux/amd64": _image_manifest("sha256:amd64"), + "linux/arm64": _image_manifest("sha256:arm64"), + }, + sample_image_config_object, + ) + + assert cpu_architectures == {gpuhunt.CPUArchitecture.X86, gpuhunt.CPUArchitecture.ARM} + assert image_config.config.user == "alice" + + def test_index_picks_the_x86_manifest_regardless_of_the_response_order( + self, sample_image_config_object + ): + # The ImageConfigs are assumed to be the same for all images within the index, but the + # manifest must still be picked deterministically, not in the registry response order + _, _, registry_client = self._get_image_config( + { + "linux/arm64": _image_manifest("sha256:arm64"), + "linux/amd64": _image_manifest("sha256:amd64"), + }, + sample_image_config_object, + ) + + registry_client.pull_blob.assert_called_once_with("sha256:amd64") + + def test_index_falls_back_to_the_arm_manifest(self, sample_image_config_object): + _, cpu_architectures, registry_client = self._get_image_config( + {"linux/arm64": _image_manifest("sha256:arm64")}, + sample_image_config_object, + ) + + assert cpu_architectures == {gpuhunt.CPUArchitecture.ARM} + registry_client.pull_blob.assert_called_once_with("sha256:arm64") + + def test_index_ignores_unsupported_platforms(self, sample_image_config_object): + _, cpu_architectures, _ = self._get_image_config( + { + "linux/amd64": _image_manifest("sha256:amd64"), + "linux/386": _image_manifest("sha256:386"), + "linux/arm/v7": _image_manifest("sha256:armv7"), + "windows/amd64": _image_manifest("sha256:windows"), + }, + sample_image_config_object, + ) + + assert cpu_architectures == {gpuhunt.CPUArchitecture.X86} + + def test_rejects_index_without_supported_platforms(self, sample_image_config_object): + with pytest.raises(DockerRegistryError, match="No supported OS/architectures found"): + self._get_image_config( + { + "linux/386": _image_manifest("sha256:386"), + "windows/amd64": _image_manifest("sha256:windows"), + }, + sample_image_config_object, + ) + + @pytest.mark.parametrize( + ["architecture", "expected_arch"], + [ + ("amd64", gpuhunt.CPUArchitecture.X86), + ("arm64", gpuhunt.CPUArchitecture.ARM), + ], + ) + def test_single_platform_image_uses_the_config_object_platform( + self, + sample_image_config_object, + architecture: str, + expected_arch: gpuhunt.CPUArchitecture, + ): + sample_image_config_object["architecture"] = architecture + + _, cpu_architectures, _ = self._get_image_config( + _image_manifest("sha256:config"), sample_image_config_object + ) + + assert cpu_architectures == {expected_arch} + + @pytest.mark.parametrize( + ["architecture", "os_name"], + [ + ("386", "linux"), + ("amd64", "windows"), + ], + ) + def test_rejects_single_platform_image_with_unsupported_platform( + self, sample_image_config_object, architecture: str, os_name: str + ): + sample_image_config_object["architecture"] = architecture + sample_image_config_object["os"] = os_name + + with pytest.raises(DockerRegistryError, match="No supported OS/architectures found"): + self._get_image_config(_image_manifest("sha256:config"), sample_image_config_object)