diff --git a/src/dstack/_internal/cli/services/configurators/run.py b/src/dstack/_internal/cli/services/configurators/run.py index 718b92e29..0733e4f71 100644 --- a/src/dstack/_internal/cli/services/configurators/run.py +++ b/src/dstack/_internal/cli/services/configurators/run.py @@ -7,9 +7,7 @@ import sys import time from pathlib import Path -from typing import Dict, List, Optional, Set, TypeVar - -import gpuhunt +from typing import Dict, List, Optional, TypeVar from dstack._internal.cli.services.args import port_mapping from dstack._internal.cli.services.configurators.base import ( @@ -78,10 +76,6 @@ from dstack._internal.utils.path import is_absolute_posix_path from dstack.api._public.runs import Run -_KNOWN_AMD_GPUS = {gpu.name.lower() for gpu in gpuhunt.KNOWN_AMD_GPUS} -_KNOWN_NVIDIA_GPUS = {gpu.name.lower() for gpu in gpuhunt.KNOWN_NVIDIA_GPUS} -_KNOWN_TPU_VERSIONS = {gpu.name.lower() for gpu in gpuhunt.KNOWN_TPUS} -_KNOWN_TENSTORRENT_GPUS = {gpu.name.lower() for gpu in gpuhunt.KNOWN_TENSTORRENT_ACCELERATORS} _BIND_ADDRESS_ARG = "bind_address" logger = get_logger(__name__) @@ -123,8 +117,6 @@ def get_plan( raise CLIError("Either --repo or --no-repo can be specified") self.apply_args(conf, configurator_args) - self.validate_gpu_vendor_and_image(conf) - self.validate_cpu_arch_and_image(conf) if conf.working_dir is not None and not is_absolute_posix_path(conf.working_dir): raise ConfigurationError("working_dir must be absolute") @@ -453,102 +445,6 @@ def interpolate_env(self, conf: RunConfigurationT): except InterpolatorError as e: raise ConfigurationError(e.args[0]) - def validate_gpu_vendor_and_image(self, conf: RunConfigurationT) -> None: - """ - Infers GPU vendor if not set. Defaults to Nvidia when using the default - CUDA image. Requires explicit `image` if the vendor is AMD or Tenstorrent. - - When vendor is inferred from GPU name (e.g. A100 -> nvidia), it is written to - gpu_spec. When vendor is inferred from image context (no name, no vendor, default - CUDA image -> nvidia), it is NOT written to gpu_spec because 0.19.x servers - (gpuhunt <0.1.12) break on vendor=nvidia + min_gpu_count=0. The server applies - the same default in set_gpu_vendor_default(). - - TODO: This entire method should move to the server (set_resources_defaults) - so that defaults and validation are equal for CLI and API users. - """ - gpu_spec = conf.resources.gpu - if gpu_spec is None: - return - if gpu_spec.count.max == 0: - return - has_amd_gpu: bool - has_tt_gpu: bool - vendor = gpu_spec.vendor - if vendor is None: - names = gpu_spec.name - if names: - # None is a placeholder for an unknown vendor. - vendors: Set[Optional[gpuhunt.AcceleratorVendor]] = set() - for name in names: - name = name.lower() - if name in _KNOWN_NVIDIA_GPUS: - vendors.add(gpuhunt.AcceleratorVendor.NVIDIA) - elif name in _KNOWN_AMD_GPUS: - vendors.add(gpuhunt.AcceleratorVendor.AMD) - elif name in _KNOWN_TENSTORRENT_GPUS: - vendors.add(gpuhunt.AcceleratorVendor.TENSTORRENT) - else: - maybe_tpu_version, _, maybe_tpu_cores = name.partition("-") - if maybe_tpu_version in _KNOWN_TPU_VERSIONS and maybe_tpu_cores.isdigit(): - vendors.add(gpuhunt.AcceleratorVendor.GOOGLE) - else: - vendors.add(None) - if len(vendors) == 1: - # Only one vendor or all names are not known. - vendor = next(iter(vendors)) - else: - # More than one vendor or some names are not known; in either case, we - # cannot set the vendor to a specific value, will use only names for matching. - vendor = None - # If some names are unknown, let's assume they are _not_ AMD products, otherwise - # ConfigurationError message may be confusing. In worst-case scenario we'll try - # to execute a run on an instance with an AMD accelerator with a default - # CUDA image, not a big deal. - has_amd_gpu = gpuhunt.AcceleratorVendor.AMD in vendors - has_tt_gpu = gpuhunt.AcceleratorVendor.TENSTORRENT in vendors - # Set vendor inferred from name on the spec (server needs it for filtering). - gpu_spec.vendor = vendor - else: - # No vendor or name specified. Default to Nvidia if using the - # default CUDA image, since it's only compatible with Nvidia GPUs. - if conf.image is None and conf.docker is not True: - vendor = gpuhunt.AcceleratorVendor.NVIDIA - has_amd_gpu = False - has_tt_gpu = False - else: - has_amd_gpu = vendor == gpuhunt.AcceleratorVendor.AMD - has_tt_gpu = vendor == gpuhunt.AcceleratorVendor.TENSTORRENT - # When docker=True, the system uses Docker-in-Docker image, so no custom image is required - if has_amd_gpu and conf.image is None and conf.docker is not True: - raise ConfigurationError("`image` is required if `resources.gpu.vendor` is `amd`") - if has_tt_gpu and conf.image is None and conf.docker is not True: - raise ConfigurationError( - "`image` is required if `resources.gpu.vendor` is `tenstorrent`" - ) - - def validate_cpu_arch_and_image(self, conf: RunConfigurationT) -> None: - """ - Infers `resources.cpu.arch` if not set, requires `image` if the architecture is ARM. - """ - cpu_spec = conf.resources.cpu - arch = cpu_spec.arch - if arch is None: - gpu_spec = conf.resources.gpu - if ( - gpu_spec is not None - and gpu_spec.vendor in [None, gpuhunt.AcceleratorVendor.NVIDIA] - and gpu_spec.name - and any(map(gpuhunt.is_nvidia_superchip, gpu_spec.name)) - ): - arch = gpuhunt.CPUArchitecture.ARM - else: - arch = gpuhunt.CPUArchitecture.X86 - # NOTE: We don't set the inferred resources.cpu.arch for compatibility with older servers. - # Servers with ARM support set the arch using the same logic. - if arch == gpuhunt.CPUArchitecture.ARM and conf.image is None: - raise ConfigurationError("`image` is required if `resources.cpu.arch` is `arm`") - def get_repo( self, conf: RunConfigurationT, diff --git a/src/dstack/_internal/server/services/fleets.py b/src/dstack/_internal/server/services/fleets.py index 105e026ea..6c6d14f4d 100644 --- a/src/dstack/_internal/server/services/fleets.py +++ b/src/dstack/_internal/server/services/fleets.py @@ -87,7 +87,10 @@ list_user_project_models, project_model_to_project, ) -from dstack._internal.server.services.resources import set_resources_defaults +from dstack._internal.server.services.resources import ( + set_default_cpu_spec_arch, + 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 ( @@ -1423,8 +1426,10 @@ def _validate_fleet_configuration_subtype_specific_fields(conf: FleetConfigurati def _set_fleet_spec_defaults(spec: FleetSpec): - if spec.configuration.resources is not None: - set_resources_defaults(spec.configuration.resources) + 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) def _validate_all_ssh_params_specified(ssh_config: SSHParams): diff --git a/src/dstack/_internal/server/services/resources.py b/src/dstack/_internal/server/services/resources.py index f439713d4..12a547e24 100644 --- a/src/dstack/_internal/server/services/resources.py +++ b/src/dstack/_internal/server/services/resources.py @@ -2,48 +2,69 @@ import gpuhunt -from dstack._internal.core.models.resources import ResourcesSpec +from dstack._internal.core.models.resources import ( + DEFAULT_GPU_SPEC, + CPUSpec, + GPUSpec, + ResourcesSpec, +) +from dstack._internal.utils.gpu import detect_gpu_vendors_by_gpu_name -def set_resources_defaults(resources: ResourcesSpec) -> None: - cpu = resources.cpu - if cpu.arch is None: - gpu = resources.gpu +def set_default_gpu_spec(resources_spec: ResourcesSpec) -> GPUSpec: + if resources_spec.gpu is None: + resources_spec.gpu = DEFAULT_GPU_SPEC.model_copy(deep=True) + 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 is not None - and gpu.vendor in [None, gpuhunt.AcceleratorVendor.NVIDIA] - and gpu.name - and any(map(gpuhunt.is_nvidia_superchip, gpu.name)) + gpu_spec.vendor in [None, gpuhunt.AcceleratorVendor.NVIDIA] + and gpu_spec.name + and any(map(gpuhunt.is_nvidia_superchip, gpu_spec.name)) ): - cpu.arch = gpuhunt.CPUArchitecture.ARM + cpu_spec.arch = gpuhunt.CPUArchitecture.ARM else: - cpu.arch = gpuhunt.CPUArchitecture.X86 + cpu_spec.arch = gpuhunt.CPUArchitecture.X86 -def set_gpu_vendor_default( - resources: ResourcesSpec, +def set_default_gpu_spec_vendor( + gpu_spec: GPUSpec, image: Optional[str], docker: Optional[bool], ) -> None: - """Default GPU vendor to Nvidia when using the default CUDA image, - since it's only compatible with Nvidia GPUs. Only called for runs - (not fleets) since fleets don't have image context. - - The client infers the same default for display and validation - (see validate_gpu_vendor_and_image) but does not write it to the spec - for 0.19.x server compatibility. This server-side function is what - actually sets the vendor before offer matching. - - TODO: All resource defaults and validation (gpu vendor, cpu arch, memory, - disk, etc.) should be set here on the server, not split between client - and model-level defaults.""" - gpu = resources.gpu - if ( - gpu is not None - and gpu.vendor is None - and gpu.name is None - and gpu.count.max != 0 - and image is None - and docker is not True - ): - gpu.vendor = gpuhunt.AcceleratorVendor.NVIDIA + """ + Infers and sets the GPU vendor if possible. + + * If the vendor is already set, does nothing. + * If no GPU requested (max=0), does nothing. + * If no names are specified (e.g., `gpu: 4`), infers the vendor from the requested image: + * If the image is not specified and DinD is not requested, that is, the default dstack + image is used, defaults to NVIDIA, since the image is only compatible with NVIDIA GPUs. + * Otherwise (the image is set or DinD is requested), does nothing. + * If names are specified (e.g., `gpu: H100,A100:4`), detects GPU vendors by the names: + * If all GPU models are known and there is only one vendor, sets that vendor. + * Otherwise (e.g., `gpu: H100,MI300X` or `gpu: H100,UNKNOWN1000`), does nothing. + """ + if gpu_spec.vendor is not None: + return + if gpu_spec.count.max == 0: + return + if not gpu_spec.name: + if image is None and not docker: + gpu_spec.vendor = gpuhunt.AcceleratorVendor.NVIDIA + else: + # None is a placeholder for an unknown vendor. + vendors: set[Optional[gpuhunt.AcceleratorVendor]] = set() + for name in gpu_spec.name: + _vendors = detect_gpu_vendors_by_gpu_name(name) + if not _vendors: + vendors.add(None) + else: + vendors.update(_vendors) + # len(vendors) == 1: Only one vendor or all names are not known (a {None} set). + # len(vendors) > 1: More than one vendor or some names are not known; in either case, we + # cannot set the vendor to a specific value, will use only names for matching. + if len(vendors) == 1: + gpu_spec.vendor = next(iter(vendors)) diff --git a/src/dstack/_internal/server/services/runs/__init__.py b/src/dstack/_internal/server/services/runs/__init__.py index 5270de31b..d96eed563 100644 --- a/src/dstack/_internal/server/services/runs/__init__.py +++ b/src/dstack/_internal/server/services/runs/__init__.py @@ -68,10 +68,6 @@ from dstack._internal.server.services.pipelines import PipelineHinterProtocol from dstack._internal.server.services.plugins import apply_plugin_policies from dstack._internal.server.services.probes import is_probe_ready -from dstack._internal.server.services.resources import ( - set_gpu_vendor_default, - set_resources_defaults, -) from dstack._internal.server.services.runs.plan import get_job_plans from dstack._internal.server.services.runs.service_router_worker_sync import ( ensure_service_router_worker_sync_row, @@ -79,6 +75,7 @@ from dstack._internal.server.services.runs.spec import ( can_update_run_spec, check_can_update_run_spec, + set_run_spec_resources_defaults, validate_run_spec_and_set_defaults, ) from dstack._internal.server.services.secrets import get_project_secrets_mapping @@ -562,7 +559,7 @@ async def get_plan( if current_resource is not None: # For backward compatibility (current_resource may has been submitted before # some fields, e.g., CPUSpec.arch, gpu.vendor were added) - _set_run_resources_defaults(current_resource.run_spec) + set_run_spec_resources_defaults(current_resource.run_spec) if not current_resource.status.is_finished() and can_update_run_spec( current_resource.run_spec, effective_run_spec ): @@ -634,7 +631,7 @@ async def apply_plan( # For backward compatibility (current_resource may has been submitted before # some fields, e.g., CPUSpec.arch, gpu.vendor were added) - _set_run_resources_defaults(current_resource.run_spec) + set_run_spec_resources_defaults(current_resource.run_spec) try: spec_diff = check_can_update_run_spec(current_resource.run_spec, run_spec) except ServerClientError: @@ -644,7 +641,7 @@ async def apply_plan( raise if not force: if plan.current_resource is not None: - _set_run_resources_defaults(plan.current_resource.run_spec) + set_run_spec_resources_defaults(plan.current_resource.run_spec) if ( plan.current_resource is None or plan.current_resource.id != current_resource.id @@ -1017,16 +1014,6 @@ def run_model_to_run( return run -def _set_run_resources_defaults(run_spec: RunSpec) -> None: - """Apply resource defaults to a run spec, including GPU vendor inference.""" - set_resources_defaults(run_spec.configuration.resources) - set_gpu_vendor_default( - run_spec.configuration.resources, - image=run_spec.configuration.image, - docker=getattr(run_spec.configuration, "docker", None), - ) - - def _get_run_jobs_with_submissions( run_spec: RunSpec, job_models: List[JobModel], diff --git a/src/dstack/_internal/server/services/runs/spec.py b/src/dstack/_internal/server/services/runs/spec.py index 364f81769..e8eb53b45 100644 --- a/src/dstack/_internal/server/services/runs/spec.py +++ b/src/dstack/_internal/server/services/runs/spec.py @@ -1,11 +1,17 @@ +from typing import Optional + +import gpuhunt + from dstack._internal.core.errors import ServerClientError from dstack._internal.core.models.configurations import ( RUN_PRIORITY_DEFAULT, SERVICE_HTTPS_DEFAULT, + ReplicaGroup, ServiceConfiguration, ) from dstack._internal.core.models.profiles import ProfileRetry from dstack._internal.core.models.repos.virtual import DEFAULT_VIRTUAL_REPO_ID, VirtualRunRepoData +from dstack._internal.core.models.resources import GPUSpec, ResourcesSpec from dstack._internal.core.models.routers import RouterType from dstack._internal.core.models.runs import LEGACY_REPO_DIR, AnyRunConfiguration, RunSpec from dstack._internal.core.models.volumes import InstanceMountPoint @@ -15,9 +21,11 @@ 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_gpu_vendor_default, - set_resources_defaults, + set_default_cpu_spec_arch, + set_default_gpu_spec, + set_default_gpu_spec_vendor, ) +from dstack._internal.utils.gpu import detect_gpu_vendors_by_gpu_name from dstack._internal.utils.logging import get_logger logger = get_logger(__name__) @@ -125,12 +133,9 @@ def validate_run_spec_and_set_defaults( run_spec.configuration.priority = RUN_PRIORITY_DEFAULT # We do not reject top-level `resources` when `replicas` is a list. Adding strict checks # would be fragile because the spec may be changed later (for example by plugins). - set_resources_defaults(run_spec.configuration.resources) - set_gpu_vendor_default( - run_spec.configuration.resources, - image=run_spec.configuration.image, - docker=getattr(run_spec.configuration, "docker", None), - ) + set_run_spec_resources_defaults(run_spec) + _validate_gpu_vendor_and_image(run_spec) + _validate_cpu_arch_and_image(run_spec) if run_spec.ssh_key_pub is None: if user.ssh_public_key: run_spec.ssh_key_pub = user.ssh_public_key @@ -140,12 +145,137 @@ def validate_run_spec_and_set_defaults( run_spec.configuration.working_dir = LEGACY_REPO_DIR +def set_run_spec_resources_defaults(run_spec: RunSpec) -> None: + """Apply resource defaults to a run spec, including GPU vendor and CPU arch inference.""" + configuration = run_spec.configuration + _set_resources_defaults( + resources_spec=configuration.resources, + image=configuration.image, + docker=configuration.docker, + ) + if configuration.type == "service" and isinstance(configuration.replicas, list): + for replica_group in configuration.replicas: + image, docker = _get_replica_group_image_and_docker(replica_group, configuration) + _set_resources_defaults( + resources_spec=replica_group.resources, + image=image, + docker=docker, + ) + + +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) + + def _validate_retry_duration(run_spec: RunSpec) -> None: retry = run_spec.merged_profile.retry if isinstance(retry, ProfileRetry) and retry.duration is not None and retry.duration < 0: raise ServerClientError("retry.duration cannot be negative") +def _validate_gpu_vendor_and_image(run_spec: RunSpec) -> None: + configuration = run_spec.configuration + vendors: set[gpuhunt.AcceleratorVendor] = set() + invalid_replicas: list[int] = [] + if configuration.type == "service" and isinstance(configuration.replicas, list): + for idx, replica_group in enumerate(configuration.replicas): + image, docker = _get_replica_group_image_and_docker(replica_group, configuration) + _vendors = _detect_gpu_vendors_requiring_image( + gpu_spec=replica_group.resources.gpu, + image=image, + docker=docker, + ) + if _vendors: + vendors.update(_vendors) + invalid_replicas.append(idx) + else: + vendors = _detect_gpu_vendors_requiring_image( + gpu_spec=configuration.resources.gpu, + image=configuration.image, + docker=configuration.docker, + ) + if vendors: + sorted_vendors = sorted(v.value for v in vendors) + msg = ( + "`image` must be set when the requested accelerator is not supported by" + f" the default image: {sorted_vendors}" + ) + if invalid_replicas: + msg = f"replicas{invalid_replicas}: {msg}" + raise ServerClientError(msg) + + +def _detect_gpu_vendors_requiring_image( + gpu_spec: Optional[GPUSpec], image: Optional[str], docker: Optional[bool] +) -> set[gpuhunt.AcceleratorVendor]: + if image is not None or docker: + return set() + if gpu_spec is None or gpu_spec.count.max == 0: + return set() + vendors: set[gpuhunt.AcceleratorVendor] = set() + if gpu_spec.vendor is not None: + vendors.add(gpu_spec.vendor) + else: + # Unknown models are ignored (loose validation -- skips possible models that + # won't work with the default dstack image). The other option would be to treat them as + # non-NVIDIA, forcing the user to set `image`, even if they actually are NVIDIA (overly + # strict validation) + for gpu_name in gpu_spec.name or []: + vendors.update(detect_gpu_vendors_by_gpu_name(gpu_name)) + # * NVIDIA definitely works with our image -- it's built for NVIDIA + # * Google TPU should work with our image -- all dependencies may be installed from PyPI, there + # are no vendors dependencies that must be preinstalled/shipped with the image; basically, + # our image is just Ubuntu + pip (uv) for TPU workloads + # * AMD, Intel Gaudi, Tenstorrent rely on some pinned system packages and/or patched libraries + # and ship their own images -- we don't expect them to work on our generic + # Ubuntu + CUDA image + return vendors - {gpuhunt.AcceleratorVendor.NVIDIA, gpuhunt.AcceleratorVendor.GOOGLE} + + +def _validate_cpu_arch_and_image(run_spec: RunSpec) -> None: + image_msg = "`image` must be set when ARM CPU requested" + docker_msg = "`docker: true` is not supported on ARM CPU" + configuration = run_spec.configuration + if configuration.type == "service" and isinstance(configuration.replicas, list): + invalid_replicas_without_image: list[int] = [] + invalid_replicas_with_docker: list[int] = [] + for idx, replica_group in enumerate(configuration.replicas): + image, docker = _get_replica_group_image_and_docker(replica_group, configuration) + if replica_group.resources.cpu.arch == gpuhunt.CPUArchitecture.ARM: + if docker: + invalid_replicas_with_docker.append(idx) + elif image is None: + invalid_replicas_without_image.append(idx) + errors: list[str] = [] + if invalid_replicas_without_image: + errors.append(f"replicas{invalid_replicas_without_image}: {image_msg}") + if invalid_replicas_with_docker: + errors.append(f"replicas{invalid_replicas_with_docker}: {docker_msg}") + if errors: + raise ServerClientError("\n".join(errors)) + elif configuration.resources.cpu.arch == gpuhunt.CPUArchitecture.ARM: + if configuration.docker: + raise ServerClientError(docker_msg) + if configuration.image is None: + raise ServerClientError(image_msg) + + +def _get_replica_group_image_and_docker( + replica_group: ReplicaGroup, configuration: ServiceConfiguration +) -> tuple[Optional[str], Optional[bool]]: + image = replica_group.image + if image is None: + image = configuration.image + docker = replica_group.docker + if docker is None: + docker = configuration.docker + return image, docker + + def _check_dynamo_in_place_update_compatibility( current_run_spec: RunSpec, new_run_spec: RunSpec ) -> None: diff --git a/src/dstack/_internal/utils/gpu.py b/src/dstack/_internal/utils/gpu.py index d0fc94e50..2998a3aaa 100644 --- a/src/dstack/_internal/utils/gpu.py +++ b/src/dstack/_internal/utils/gpu.py @@ -1,5 +1,7 @@ import re +import gpuhunt + def convert_nvidia_gpu_name(name: str) -> str: """Convert gpu_name from nvidia-smi to short version""" @@ -58,3 +60,29 @@ def convert_intel_accelerator_name(name: str) -> str: "HL-325": "Gaudi3", # OAM "HL-338": "Gaudi3", # PCIe } + + +def detect_gpu_vendors_by_gpu_name(name: str) -> set[gpuhunt.AcceleratorVendor]: + vendors: set[gpuhunt.AcceleratorVendor] = set() + name = name.lower() + if name in _KNOWN_NVIDIA_GPUS: + vendors.add(gpuhunt.AcceleratorVendor.NVIDIA) + if name in _KNOWN_AMD_GPUS: + vendors.add(gpuhunt.AcceleratorVendor.AMD) + if name in _KNOWN_INTEL_ACCELERATORS: + vendors.add(gpuhunt.AcceleratorVendor.INTEL) + if name in _KNOWN_TENSTORRENT_ACCELERATORS: + vendors.add(gpuhunt.AcceleratorVendor.TENSTORRENT) + maybe_tpu_version, _, maybe_tpu_cores = name.partition("-") + if maybe_tpu_cores.isdigit() and maybe_tpu_version in _KNOWN_TPU_VERSIONS: + vendors.add(gpuhunt.AcceleratorVendor.GOOGLE) + return vendors + + +_KNOWN_NVIDIA_GPUS = {gpu.name.lower() for gpu in gpuhunt.KNOWN_NVIDIA_GPUS} +_KNOWN_AMD_GPUS = {gpu.name.lower() for gpu in gpuhunt.KNOWN_AMD_GPUS} +_KNOWN_INTEL_ACCELERATORS = {gpu.name.lower() for gpu in gpuhunt.KNOWN_INTEL_ACCELERATORS} +_KNOWN_TENSTORRENT_ACCELERATORS = { + gpu.name.lower() for gpu in gpuhunt.KNOWN_TENSTORRENT_ACCELERATORS +} +_KNOWN_TPU_VERSIONS = {gpu.name.lower() for gpu in gpuhunt.KNOWN_TPUS} diff --git a/src/tests/_internal/cli/services/configurators/test_run.py b/src/tests/_internal/cli/services/configurators/test_run.py index c0d468fea..1e86759a5 100644 --- a/src/tests/_internal/cli/services/configurators/test_run.py +++ b/src/tests/_internal/cli/services/configurators/test_run.py @@ -1,14 +1,13 @@ import argparse from textwrap import dedent -from typing import List, Optional, Tuple +from typing import List, Tuple from unittest.mock import Mock import pytest -from gpuhunt import KNOWN_TENSTORRENT_ACCELERATORS, AcceleratorVendor +from gpuhunt import KNOWN_TENSTORRENT_ACCELERATORS from dstack._internal.cli.services.configurators import get_run_configurator_class from dstack._internal.cli.services.configurators.run import ( - BaseRunConfigurator, ServiceConfigurator, render_run_spec_diff, ) @@ -110,229 +109,6 @@ def test_interpolates_env(self): ) -class TestValidateGPUVendorAndImage: - def prepare_conf( - self, - *, - image: Optional[str] = None, - gpu_spec: Optional[str] = None, - docker: Optional[bool] = None, - ) -> BaseRunConfiguration: - conf_dict = { - "type": "none", - } - if image is not None: - conf_dict["image"] = image - if gpu_spec is not None: - conf_dict["resources"] = { - "gpu": gpu_spec, - } - if docker is not None: - conf_dict["docker"] = docker - return BaseRunConfiguration.model_validate(conf_dict) - - def validate(self, conf: BaseRunConfiguration) -> None: - BaseRunConfigurator(api_client=Mock()).validate_gpu_vendor_and_image(conf) - - def test_no_gpu(self): - conf = self.prepare_conf() - self.validate(conf) - assert conf.resources.gpu is not None - # Vendor is not written to spec for compatibility with older servers. - # The server infers nvidia in set_resources_defaults(). - assert conf.resources.gpu.vendor is None - assert conf.resources.gpu.name is None - assert conf.resources.gpu.count.min == 0 - - def test_zero_gpu(self): - conf = self.prepare_conf(gpu_spec="0") - self.validate(conf) - assert conf.resources.gpu.vendor is None - - def test_gpu_no_vendor_no_image_defaults_to_nvidia(self): - """Vendor is inferred as nvidia for validation but NOT written to spec.""" - conf = self.prepare_conf(gpu_spec="1") - self.validate(conf) - assert conf.resources.gpu.vendor is None - - def test_gpu_no_vendor_with_image_no_default(self): - conf = self.prepare_conf(gpu_spec="1", image="my-custom-image") - self.validate(conf) - assert conf.resources.gpu.vendor is None - - def test_gpu_no_vendor_docker_true_no_default(self): - conf = self.prepare_conf(gpu_spec="1", docker=True) - self.validate(conf) - assert conf.resources.gpu.vendor is None - - @pytest.mark.parametrize( - ["gpu_spec", "expected_vendor"], - [ - ["nvidia", AcceleratorVendor.NVIDIA], - ["tpu", AcceleratorVendor.GOOGLE], - ["google", AcceleratorVendor.GOOGLE], - ], - ) - def test_non_amd_vendor_declared(self, gpu_spec, expected_vendor): - conf = self.prepare_conf(gpu_spec=gpu_spec) - self.validate(conf) - assert conf.resources.gpu.vendor == expected_vendor - - def test_amd_vendor_declared_with_image(self): - conf = self.prepare_conf(image="tgi:rocm", gpu_spec="AMD") - self.validate(conf) - assert conf.resources.gpu.vendor == AcceleratorVendor.AMD - - @pytest.mark.parametrize( - ["gpu_spec", "expected_vendor"], - [ - ["a40,l40", AcceleratorVendor.NVIDIA], # lowercase - ["V3-64", AcceleratorVendor.GOOGLE], # uppercase - ], - ) - def test_one_non_amd_vendor_inferred(self, gpu_spec, expected_vendor): - conf = self.prepare_conf(gpu_spec=gpu_spec) - self.validate(conf) - assert conf.resources.gpu.vendor == expected_vendor - - @pytest.mark.parametrize("gpu_spec", ["MI300X", "MI300x", "mi300x"]) - def test_amd_vendor_inferred_with_image(self, gpu_spec): - conf = self.prepare_conf(image="tgi:rocm", gpu_spec=gpu_spec) - self.validate(conf) - assert conf.resources.gpu.vendor == AcceleratorVendor.AMD - - @pytest.mark.parametrize("gpu_spec", ["foo", "foo,bar"]) - def test_one_unknown_vendor_inferred(self, gpu_spec): - conf = self.prepare_conf(gpu_spec=gpu_spec) - self.validate(conf) - assert conf.resources.gpu.vendor is None - - @pytest.mark.parametrize( - "gpu_spec", - [ - "A1000,v4", # Nvidia and Google - "v3-64,foo", # Google and unknown - ], - ) - def test_two_non_amd_vendors_inferred(self, gpu_spec): - conf = self.prepare_conf(gpu_spec=gpu_spec) - self.validate(conf) - assert conf.resources.gpu.vendor is None - - @pytest.mark.parametrize( - "gpu_spec", - [ - "A1000,mi300x", # Nvidia and AMD (lowercase) - "MI300x,v3-64", # AMD (mixedcase) and Google - "foo,MI300X", # unknown and AMD (uppercase) - ], - ) - def test_two_vendors_including_amd_inferred_with_image(self, gpu_spec): - conf = self.prepare_conf(image="tgi:rocm", gpu_spec=gpu_spec) - self.validate(conf) - assert conf.resources.gpu.vendor is None - - def test_amd_vendor_declared_no_image(self): - conf = self.prepare_conf(gpu_spec="AMD") - with pytest.raises( - ConfigurationError, match=r"`image` is required if `resources.gpu.vendor` is `amd`" - ): - self.validate(conf) - - @pytest.mark.parametrize("gpu_spec", ["AMD", "MI300X"]) - def test_amd_vendor_docker_true_no_image(self, gpu_spec): - conf = self.prepare_conf(gpu_spec=gpu_spec, docker=True) - self.validate(conf) - assert conf.resources.gpu.vendor == AcceleratorVendor.AMD - - @pytest.mark.parametrize("gpu_spec", ["MI300X", "MI300x", "mi300x"]) - def test_amd_vendor_inferred_no_image(self, gpu_spec): - conf = self.prepare_conf(gpu_spec=gpu_spec) - with pytest.raises( - ConfigurationError, match=r"`image` is required if `resources.gpu.vendor` is `amd`" - ): - self.validate(conf) - - @pytest.mark.parametrize( - "gpu_spec", - [ - "A1000,mi300x", # Nvidia and AMD (lowercase) - "MI300x,v3-64", # AMD (mixedcase) and Google - "foo,MI300X", # unknown and AMD (uppercase) - ], - ) - def test_two_vendors_including_amd_inferred_no_image(self, gpu_spec): - conf = self.prepare_conf(gpu_spec=gpu_spec) - with pytest.raises( - ConfigurationError, match=r"`image` is required if `resources.gpu.vendor` is `amd`" - ): - self.validate(conf) - - @pytest.mark.parametrize("gpu_spec", _TENSTORRENT_ACCELERATOR_NAMES) - def test_tenstorrent_docker_true_no_image(self, gpu_spec): - conf = self.prepare_conf(gpu_spec=gpu_spec, docker=True) - self.validate(conf) - assert conf.resources.gpu.vendor == AcceleratorVendor.TENSTORRENT - - @pytest.mark.parametrize("gpu_spec", _TENSTORRENT_ACCELERATOR_NAMES) - def test_tenstorrent_vendor_inferred_no_image(self, gpu_spec): - conf = self.prepare_conf(gpu_spec=gpu_spec) - with pytest.raises( - ConfigurationError, - match=r"`image` is required if `resources.gpu.vendor` is `tenstorrent`", - ): - self.validate(conf) - - -class TestValidateCPUArchAndImage: - def prepare_conf( - self, - *, - cpu_spec: str, - gpu_spec: Optional[str] = None, - image: Optional[str] = None, - ) -> BaseRunConfiguration: - conf_dict = { - "type": "none", - "resources": { - "cpu": cpu_spec, - }, - } - if image is not None: - conf_dict["image"] = image - if gpu_spec is not None: - conf_dict["resources"]["gpu"] = gpu_spec - return BaseRunConfiguration.model_validate(conf_dict) - - def validate(self, conf: BaseRunConfiguration) -> None: - # validate_gpu_vendor_and_image sets GPU vendor if not set - BaseRunConfigurator(api_client=Mock()).validate_gpu_vendor_and_image(conf) - BaseRunConfigurator(api_client=Mock()).validate_cpu_arch_and_image(conf) - - @pytest.mark.parametrize("gpu_spec", [None, "GH200", "H100"]) - def test_explicit_arm_with_image(self, gpu_spec: Optional[str]): - conf = self.prepare_conf(cpu_spec="arm:1..", gpu_spec=gpu_spec, image="ubuntu") - self.validate(conf) - - def test_inferred_arm_with_image(self): - conf = self.prepare_conf(cpu_spec="1..", gpu_spec="GH200", image="ubuntu") - self.validate(conf) - - @pytest.mark.parametrize("cpu_spec", ["1..", "arm:1.."]) - def test_arm_no_image(self, cpu_spec: str): - conf = self.prepare_conf(cpu_spec=cpu_spec, gpu_spec="GH200") - with pytest.raises( - ConfigurationError, match=r"`image` is required if `resources.cpu.arch` is `arm`" - ): - self.validate(conf) - - @pytest.mark.parametrize("cpu_spec", ["1..", "x86:1.."]) - @pytest.mark.parametrize("image", [None, "ubuntu"]) - def test_x86(self, cpu_spec: str, image: Optional[str]): - conf = self.prepare_conf(cpu_spec=cpu_spec, gpu_spec="H100", image=image) - self.validate(conf) - - class TestApplyConfiguration: def test_composes_get_plan_and_apply_plan(self, monkeypatch): run_plan, repo = Mock(), Mock() diff --git a/src/tests/_internal/server/routers/test_runs.py b/src/tests/_internal/server/routers/test_runs.py index 5a40f6159..08207c944 100644 --- a/src/tests/_internal/server/routers/test_runs.py +++ b/src/tests/_internal/server/routers/test_runs.py @@ -50,12 +50,11 @@ from dstack._internal.server.models import JobModel, RunModel from dstack._internal.server.schemas.runs import MAX_JOB_SUBMISSIONS_LIMIT, ApplyRunPlanRequest from dstack._internal.server.services.projects import add_project_member -from dstack._internal.server.services.resources import ( - set_gpu_vendor_default, - set_resources_defaults, -) from dstack._internal.server.services.runs import run_model_to_run -from dstack._internal.server.services.runs.spec import validate_run_spec_and_set_defaults +from dstack._internal.server.services.runs.spec import ( + set_run_spec_resources_defaults, + validate_run_spec_and_set_defaults, +) from dstack._internal.server.testing.common import ( create_backend, create_export, @@ -89,6 +88,32 @@ def disable_sshproxy(monkeypatch: pytest.MonkeyPatch): monkeypatch.setattr("dstack._internal.server.settings.SSHPROXY_ENABLED", False) +def get_default_gpu_dict(docker: Optional[bool] = None) -> Dict: + """ + The GPU spec the server sets when the client submits no GPU requirements. + The vendor is inferred from the image: the default dstack image is NVIDIA-only, + while the DinD image works with any vendor. + """ + return { + "vendor": None if docker else "nvidia", + "name": None, + "count": {"min": 0, "max": None}, + "memory": None, + "total_memory": None, + "compute_capability": None, + } + + +def get_submitted_run_spec_dict(run_spec: Dict) -> Dict: + """ + A copy of the run spec dict as submitted by the client, that is, without the resource + defaults that the server sets (see `set_run_spec_resources_defaults`). + """ + run_spec = copy.deepcopy(run_spec) + run_spec["configuration"]["resources"]["gpu"] = None + return run_spec + + def get_dev_env_run_plan_dict( project_name: str = "test_project", username: str = "test_user", @@ -185,7 +210,7 @@ def get_dev_env_run_plan_dict( "cpu": {"arch": "x86", "count": {"min": 2, "max": None}}, "memory": {"min": 8.0, "max": None}, "disk": None, - "gpu": None, + "gpu": get_default_gpu_dict(docker), "shm_size": None, }, "volumes": [json.loads(v.model_dump_json()) for v in volumes], @@ -267,7 +292,8 @@ def get_dev_env_run_plan_dict( return { "project_name": project_name, "user": username, - "run_spec": run_spec, + # `run_spec` is returned as submitted, `effective_run_spec` — with the server defaults + "run_spec": get_submitted_run_spec_dict(run_spec), "effective_run_spec": run_spec, "job_plans": [ { @@ -294,7 +320,7 @@ def get_dev_env_run_plan_dict( "cpu": {"arch": "x86", "count": {"min": 2, "max": None}}, "memory": {"min": 8.0, "max": None}, "disk": None, - "gpu": None, + "gpu": get_default_gpu_dict(docker), "shm_size": None, }, "max_price": None, @@ -437,7 +463,7 @@ def get_dev_env_run_dict( "cpu": {"arch": "x86", "count": {"min": 2, "max": None}}, "memory": {"min": 8.0, "max": None}, "disk": None, - "gpu": None, + "gpu": get_default_gpu_dict(docker), "shm_size": None, }, "volumes": [], @@ -541,7 +567,7 @@ def get_dev_env_run_dict( "cpu": {"arch": "x86", "count": {"min": 2, "max": None}}, "memory": {"min": 8.0, "max": None}, "disk": None, - "gpu": None, + "gpu": get_default_gpu_dict(docker), "shm_size": None, }, "max_price": None, @@ -3100,13 +3126,7 @@ async def test_returns_update_or_create_action_on_conf_change( run_spec=run_spec, ) run = run_model_to_run(run_model) - # Apply the same defaults the server applies to current_resource - set_resources_defaults(run.run_spec.configuration.resources) - set_gpu_vendor_default( - run.run_spec.configuration.resources, - image=run.run_spec.configuration.image, - docker=getattr(run.run_spec.configuration, "docker", None), - ) + set_run_spec_resources_defaults(run.run_spec) run_spec.configuration = new_conf response = await client.post( f"/api/project/{project.name}/runs/get_plan", @@ -3307,7 +3327,7 @@ async def test_submits_new_run_docker_true( headers=get_auth_headers(user.token), json={ "plan": { - "run_spec": run_dict["run_spec"], + "run_spec": get_submitted_run_spec_dict(run_dict["run_spec"]), "current_resource": None, }, "force": False, diff --git a/src/tests/_internal/server/services/runs/test_spec.py b/src/tests/_internal/server/services/runs/test_spec.py index 0c62ad721..1d23a469e 100644 --- a/src/tests/_internal/server/services/runs/test_spec.py +++ b/src/tests/_internal/server/services/runs/test_spec.py @@ -1,11 +1,13 @@ import re import uuid from types import SimpleNamespace +from typing import Any, Optional +import gpuhunt import pytest from dstack._internal.core.errors import ServerClientError -from dstack._internal.core.models.configurations import ServiceConfiguration +from dstack._internal.core.models.configurations import ServiceConfiguration, TaskConfiguration from dstack._internal.core.models.files import FileArchiveMapping from dstack._internal.core.models.profiles import Profile, ProfileRetry from dstack._internal.core.models.repos.local import LocalRunRepoData @@ -13,6 +15,7 @@ from dstack._internal.server.services.runs.spec import ( _check_can_update_configuration, check_can_update_run_spec, + set_run_spec_resources_defaults, validate_run_spec_and_set_defaults, ) from dstack._internal.server.testing.common import get_run_spec @@ -80,6 +83,51 @@ def _run_spec_with_overrides(configuration: ServiceConfiguration, **overrides) - return RunSpec.model_validate({**run_spec.model_dump(), **run_spec_overrides}) +def _task_run_spec( + *, + resources: Optional[dict] = None, + image: Optional[str] = None, + docker: Optional[bool] = None, +) -> RunSpec: + conf: dict[str, Any] = {"type": "task", "commands": ["echo hello"]} + if resources is not None: + conf["resources"] = resources + if image is not None: + conf["image"] = image + if docker is not None: + conf["docker"] = docker + return get_run_spec( + repo_id="test-repo", + run_name="test-run", + configuration=TaskConfiguration.model_validate(conf), + ) + + +def _service_run_spec( + *, + replicas: list[dict], + resources: Optional[dict] = None, + image: Optional[str] = None, + docker: Optional[bool] = None, +) -> RunSpec: + conf: dict[str, Any] = {"type": "service", "port": 8000, "replicas": replicas} + if resources is not None: + conf["resources"] = resources + if image is not None: + conf["image"] = image + if docker is not None: + conf["docker"] = docker + return get_run_spec( + repo_id="test-repo", + run_name="test-run", + configuration=ServiceConfiguration.model_validate(conf), + ) + + +def _validate(run_spec: RunSpec) -> None: + validate_run_spec_and_set_defaults(SimpleNamespace(ssh_public_key="ssh-rsa test"), run_spec) + + class TestValidateRunSpecRetryDuration: def test_model_accepts_negative_retry_duration_for_backward_compatibility(self): retry = ProfileRetry(duration=-1) @@ -229,3 +277,274 @@ def test_non_dynamo_image_change_passes_configuration_gate(self): current = _service_configuration(router_type="sglang", image="img:1") new = _service_configuration(router_type="sglang", image="img:2") _check_can_update_configuration(current, new, ignore_files=True) + + +class TestSetRunSpecResourcesDefaultsGpuVendor: + @pytest.mark.parametrize( + ["gpu_spec", "expected_vendor"], + [ + ("A100", gpuhunt.AcceleratorVendor.NVIDIA), + ("a40,l40", gpuhunt.AcceleratorVendor.NVIDIA), # different names, same vendor + ("Mi300X", gpuhunt.AcceleratorVendor.AMD), + ("Gaudi2", gpuhunt.AcceleratorVendor.INTEL), + ("n300", gpuhunt.AcceleratorVendor.TENSTORRENT), + ("v5litepod-8", gpuhunt.AcceleratorVendor.GOOGLE), + ], + ) + def test_sets_vendor_detected_by_gpu_names( + self, gpu_spec: str, expected_vendor: gpuhunt.AcceleratorVendor + ): + run_spec = _task_run_spec(resources={"gpu": gpu_spec}, image="ubuntu") + + set_run_spec_resources_defaults(run_spec) + + assert run_spec.configuration.resources.gpu.vendor == expected_vendor + + @pytest.mark.parametrize( + "gpu_spec", + [ + "UNKNOWN1000", # an unknown name + "A100,UNKNOWN1000", # known and unknown names + "A100,MI300X", # names of different vendors + ], + ) + def test_does_not_set_vendor_if_gpu_names_are_ambiguous(self, gpu_spec: str): + run_spec = _task_run_spec(resources={"gpu": gpu_spec}, image="ubuntu") + + set_run_spec_resources_defaults(run_spec) + + assert run_spec.configuration.resources.gpu.vendor is None + + def test_does_not_override_vendor_set_by_the_user(self): + run_spec = _task_run_spec( + resources={"gpu": {"vendor": "amd", "name": ["A100"]}}, image="ubuntu" + ) + + set_run_spec_resources_defaults(run_spec) + + assert run_spec.configuration.resources.gpu.vendor == gpuhunt.AcceleratorVendor.AMD + + def test_sets_nvidia_if_the_default_image_is_used(self): + run_spec = _task_run_spec(resources={"gpu": "1"}) + + set_run_spec_resources_defaults(run_spec) + + assert run_spec.configuration.resources.gpu.vendor == gpuhunt.AcceleratorVendor.NVIDIA + + @pytest.mark.parametrize( + ["image", "docker"], + [ + ("ubuntu", None), + (None, True), # the DinD image can run containers with any accelerator + ], + ) + def test_does_not_set_vendor_if_the_default_image_is_not_used( + self, image: Optional[str], docker: Optional[bool] + ): + run_spec = _task_run_spec(resources={"gpu": "1"}, image=image, docker=docker) + + set_run_spec_resources_defaults(run_spec) + + assert run_spec.configuration.resources.gpu.vendor is None + + def test_does_not_set_vendor_if_no_gpu_requested(self): + run_spec = _task_run_spec(resources={"gpu": "0"}) + + set_run_spec_resources_defaults(run_spec) + + assert run_spec.configuration.resources.gpu.vendor is None + + def test_sets_default_gpu_spec_if_gpu_is_null(self): + run_spec = _task_run_spec(resources={"gpu": None}) + + set_run_spec_resources_defaults(run_spec) + + gpu_spec = run_spec.configuration.resources.gpu + assert gpu_spec is not None + assert gpu_spec.name is None + assert gpu_spec.count.min == 0 + assert gpu_spec.count.max is None + 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( + replicas=[ + {"count": 1, "commands": ["echo"], "resources": {"gpu": "MI300X"}}, + {"count": 1, "commands": ["echo"], "resources": {"gpu": "GH200"}}, + {"count": 1, "commands": ["echo"]}, + ], + ) + + 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), + ] + + @pytest.mark.parametrize( + ["service_image", "group_image", "expected_vendor"], + [ + (None, None, gpuhunt.AcceleratorVendor.NVIDIA), + (None, "ubuntu", None), # the group image overrides the default image + ("ubuntu", None, None), # the group inherits the service-level image + ], + ) + def test_infers_vendor_from_the_image_used_by_the_group( + self, + service_image: Optional[str], + group_image: Optional[str], + expected_vendor: Optional[gpuhunt.AcceleratorVendor], + ): + group: dict = {"count": 1, "commands": ["echo"], "resources": {"gpu": "1"}} + if group_image is not None: + group["image"] = group_image + run_spec = _service_run_spec(replicas=[group], image=service_image) + + set_run_spec_resources_defaults(run_spec) + + assert run_spec.configuration.replicas[0].resources.gpu.vendor == expected_vendor + + def test_sets_defaults_for_top_level_resources(self): + # The top-level resources are ignored when replica groups are set, but they are still + # normalized so that resubmitting the same configuration produces no spec diff + run_spec = _service_run_spec( + replicas=[{"count": 1, "commands": ["echo"]}], + resources={"gpu": "H100"}, + ) + + set_run_spec_resources_defaults(run_spec) + + resources = run_spec.configuration.resources + assert resources.gpu.vendor == gpuhunt.AcceleratorVendor.NVIDIA + assert resources.cpu.arch == gpuhunt.CPUArchitecture.X86 + + +class TestValidateRunSpecGpuVendorAndImage: + UNSUPPORTED_GPU_SPECS = ["amd", "MI300X", "intel", "Gaudi2", "tenstorrent", "n300"] + + @pytest.mark.parametrize("gpu_spec", UNSUPPORTED_GPU_SPECS) + def test_rejects_gpu_not_supported_by_the_default_image(self, gpu_spec: str): + run_spec = _task_run_spec(resources={"gpu": gpu_spec}) + + with pytest.raises(ServerClientError, match="`image` must be set"): + _validate(run_spec) + + @pytest.mark.parametrize("gpu_spec", UNSUPPORTED_GPU_SPECS) + @pytest.mark.parametrize(["image", "docker"], [("rocm", None), (None, True)]) + def test_allows_any_gpu_if_the_default_image_is_not_used( + self, gpu_spec: str, image: Optional[str], docker: Optional[bool] + ): + _validate(_task_run_spec(resources={"gpu": gpu_spec}, image=image, docker=docker)) + + @pytest.mark.parametrize( + "gpu_spec", + [ + "nvidia", + "H100", + # TPU workloads install all dependencies from PyPI, so they work with the default image + "google", + "v5litepod-8", + "UNKNOWN1000", # unknown names are not validated + ], + ) + def test_allows_gpu_supported_by_the_default_image(self, gpu_spec: str): + _validate(_task_run_spec(resources={"gpu": gpu_spec})) + + def test_allows_any_vendor_if_no_gpu_requested(self): + _validate(_task_run_spec(resources={"gpu": {"vendor": "amd", "count": 0}})) + + def test_reports_replica_groups_requiring_image(self): + run_spec = _service_run_spec( + replicas=[ + {"count": 1, "commands": ["echo"], "resources": {"gpu": "MI300X"}}, + {"count": 1, "commands": ["echo"], "resources": {"gpu": "H100"}}, + {"count": 1, "commands": ["echo"], "resources": {"gpu": "n300"}}, + ], + ) + + with pytest.raises(ServerClientError, match=re.escape("replicas[0, 2]")): + _validate(run_spec) + + def test_allows_replica_group_with_its_own_image(self): + run_spec = _service_run_spec( + replicas=[{"count": 1, "image": "rocm", "resources": {"gpu": "MI300X"}}], + ) + + _validate(run_spec) + + +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): + with pytest.raises(ServerClientError, match="`image` must be set when ARM CPU requested"): + _validate(_task_run_spec(resources=resources)) + + def test_allows_arm_with_image(self): + _validate(_task_run_spec(resources={"cpu": "arm:2"}, image="ubuntu")) + + def test_rejects_arm_with_dind(self): + # `image` cannot be set with `docker: true`, and the DinD image is x86-only + with pytest.raises(ServerClientError, match="`docker: true` is not supported on ARM CPU"): + _validate(_task_run_spec(resources={"cpu": "arm:2"}, docker=True)) + + @pytest.mark.parametrize("resources", [None, {"cpu": "x86:2"}, {"gpu": "H100"}]) + def test_allows_x86_without_image(self, resources: Optional[dict]): + _validate(_task_run_spec(resources=resources)) + + def test_reports_replica_groups_requiring_image(self): + run_spec = _service_run_spec( + replicas=[ + {"count": 1, "commands": ["echo"], "resources": {"cpu": "arm:2"}}, + {"count": 1, "commands": ["echo"]}, + {"count": 1, "commands": ["echo"], "resources": {"gpu": "GH200"}}, + ], + ) + + with pytest.raises(ServerClientError, match=re.escape("replicas[0, 2]")): + _validate(run_spec) + + def test_allows_replica_group_with_its_own_image(self): + run_spec = _service_run_spec( + replicas=[{"count": 1, "image": "ubuntu", "resources": {"cpu": "arm:2"}}], + ) + + _validate(run_spec) diff --git a/src/tests/_internal/utils/test_gpu.py b/src/tests/_internal/utils/test_gpu.py index b649b8d61..1847ef207 100644 --- a/src/tests/_internal/utils/test_gpu.py +++ b/src/tests/_internal/utils/test_gpu.py @@ -1,9 +1,11 @@ +import gpuhunt import pytest from dstack._internal.utils.gpu import ( convert_amd_gpu_name, convert_intel_accelerator_name, convert_nvidia_gpu_name, + detect_gpu_vendors_by_gpu_name, ) @@ -59,3 +61,37 @@ def test_convert_amd_gpu_name(self, test_input, expected): ) def test_convert_intel_accelerator_name(self, test_input, expected): assert convert_intel_accelerator_name(test_input) == expected + + +class TestDetectGpuVendorsByGpuName: + @pytest.mark.parametrize( + ["name", "expected"], + [ + ("A100", gpuhunt.AcceleratorVendor.NVIDIA), + ("MI300X", gpuhunt.AcceleratorVendor.AMD), + ("Gaudi2", gpuhunt.AcceleratorVendor.INTEL), + ("n300", gpuhunt.AcceleratorVendor.TENSTORRENT), + ("v5litepod-8", gpuhunt.AcceleratorVendor.GOOGLE), + ], + ) + def test_detects_known_names(self, name: str, expected: gpuhunt.AcceleratorVendor): + assert detect_gpu_vendors_by_gpu_name(name) == {expected} + + @pytest.mark.parametrize("name", ["mi300x", "MI300X", "Mi300X"]) + def test_ignores_case(self, name: str): + assert detect_gpu_vendors_by_gpu_name(name) == {gpuhunt.AcceleratorVendor.AMD} + + @pytest.mark.parametrize("name", ["v2-8", "V3-64", "v5litepod-8", "v6e-8"]) + def test_detects_tpus_by_version_and_cores(self, name: str): + assert detect_gpu_vendors_by_gpu_name(name) == {gpuhunt.AcceleratorVendor.GOOGLE} + + @pytest.mark.parametrize( + "name", + [ + "UNKNOWN1000", + "v3", # a TPU version without the number of cores + "v3-x", # a TPU version with a non-numeric number of cores + ], + ) + def test_returns_empty_set_for_unknown_names(self, name: str): + assert detect_gpu_vendors_by_gpu_name(name) == set()