Skip to content

Rework GPU vendor/CPU arch and image validation #4114

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
un-def wants to merge 1 commit into
base: master
Choose a base branch
from
Open

Rework GPU vendor/CPU arch and image validation #4114

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
106 changes: 1 addition & 105 deletions src/dstack/_internal/cli/services/configurators/run.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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 (
Expand Down Expand Up @@ -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__)
Expand Down Expand Up @@ -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")
Expand Down Expand Up @@ -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,
Expand Down
11 changes: 8 additions & 3 deletions src/dstack/_internal/server/services/fleets.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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 (
Expand Down Expand Up @@ -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):
Expand Down
91 changes: 56 additions & 35 deletions src/dstack/_internal/server/services/resources.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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))
21 changes: 4 additions & 17 deletions src/dstack/_internal/server/services/runs/__init__.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -68,17 +68,14 @@
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,
)
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
Expand Down Expand Up @@ -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
):
Expand Down Expand Up @@ -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:
Expand All @@ -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
Expand Down Expand Up @@ -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],
Expand Down
Loading
Loading