Change CPU arch default value/validation (#4136)

* In the fleet spec, None/missing value now means "any arch", not
  "must be inferred and set"
* In the run spec, None/missing value now means "any arch supported by
  the image", not "must be inferred and set"
* The arch inferred from the image registry is set per-JobSpec (not
  RunSpec) via JobConfigurator. As a consequence, resources.cpu.arch
  must now be set in addition to commands and user to avoid contacting
  the image registry (e.g., when a local image not pushed to the
  registry is used)
* As the arch is no longer inferred from GPU models, the image/cpu.arch
  field combination is only validated when cpu.arch is explicitly set
  to arm. A configuration that requests an NVIDIA ARM chip via gpu but
  sets neither image nor cpu.arch is valid now
* Relaxing resources.cpu.arch from a specific value to null is treated
  as a compatible change, so re-applying an unchanged configuration
  after a server upgrade neither redeploys service replicas nor fails
  the in-place update
* image: "scratch" now has a special meaning -- it forces the server
  to use some dummy defaults instead of contacting the image registry
  to extract optional field values from the image config. Used by
  `dstack offer`

Fixes: #4056

Author: un-def

mkdocs/docs/concepts/dev-environments.md

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

mkdocs/docs/concepts/services.md

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

mkdocs/docs/concepts/tasks.md

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

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=[":"],

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,
@@ -524,7 +525,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:

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}")

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

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)
 

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