Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
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
3 changes: 2 additions & 1 deletion mkdocs/docs/concepts/dev-environments.md
Original file line number Diff line number Diff line change
Expand Up @@ -159,7 +159,8 @@ resources:
</div>

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).

Expand Down
3 changes: 2 additions & 1 deletion mkdocs/docs/concepts/services.md
Original file line number Diff line number Diff line change
Expand Up @@ -879,7 +879,8 @@ resources:
</div>

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).

Expand Down
3 changes: 2 additions & 1 deletion mkdocs/docs/concepts/tasks.md
Original file line number Diff line number Diff line change
Expand Up @@ -220,7 +220,8 @@ resources:
</div>

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).

Expand Down
10 changes: 7 additions & 3 deletions src/dstack/_internal/cli/commands/offer.py
Original file line number Diff line number Diff line change
Expand Up @@ -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=[":"],
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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 (
Expand Down Expand Up @@ -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:
Expand Down
74 changes: 64 additions & 10 deletions src/dstack/_internal/server/services/docker.py
Original file line number Diff line number Diff line change
@@ -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
Expand All @@ -18,7 +20,6 @@
parse_image_name,
)

DEFAULT_PLATFORM = "linux/amd64"
MAX_CONFIG_OBJECT_SIZE = 2**22 # 4 MiB
REGISTRY_REQUEST_TIMEOUT = 20

Expand Down Expand Up @@ -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")
Expand All @@ -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
Expand All @@ -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)
Expand Down Expand Up @@ -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}")
8 changes: 2 additions & 6 deletions src/dstack/_internal/server/services/fleets.py
Original file line number Diff line number Diff line change
Expand Up @@ -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 (
Expand Down Expand Up @@ -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):
Expand Down
18 changes: 18 additions & 0 deletions src/dstack/_internal/server/services/jobs/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)

Expand Down
71 changes: 61 additions & 10 deletions src/dstack/_internal/server/services/jobs/configurators/base.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand All @@ -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
Expand Down Expand Up @@ -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

Expand Down Expand Up @@ -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:
Expand All @@ -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,
Expand All @@ -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),
Expand Down Expand Up @@ -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 []
Expand Down Expand Up @@ -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:
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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
Loading
Loading