dstack 0.19.17__py3-none-any.whl → 0.19.19__py3-none-any.whl

dstack/_internal/core/backends/gcp/compute.py

@@ -8,6 +8,7 @@ import google.api_core.exceptions
 import google.cloud.compute_v1 as compute_v1
 from cachetools import TTLCache, cachedmethod
 from google.cloud import tpu_v2
+from google.cloud.compute_v1.types.compute import Instance
 from gpuhunt import KNOWN_TPUS
 
 import dstack._internal.core.backends.gcp.auth as auth
@@ -19,6 +20,7 @@ from dstack._internal.core.backends.base.compute import (
     ComputeWithGatewaySupport,
     ComputeWithMultinodeSupport,
     ComputeWithPlacementGroupSupport,
+    ComputeWithPrivateGatewaySupport,
     ComputeWithVolumeSupport,
     generate_unique_gateway_instance_name,
     generate_unique_instance_name,
@@ -83,6 +85,7 @@ class GCPCompute(
     ComputeWithMultinodeSupport,
     ComputeWithPlacementGroupSupport,
     ComputeWithGatewaySupport,
+    ComputeWithPrivateGatewaySupport,
     ComputeWithVolumeSupport,
     Compute,
 ):
@@ -395,11 +398,7 @@ class GCPCompute(
         if instance.status in ["PROVISIONING", "STAGING"]:
             return
         if instance.status == "RUNNING":
-            if allocate_public_ip:
-                hostname = instance.network_interfaces[0].access_configs[0].nat_i_p
-            else:
-                hostname = instance.network_interfaces[0].network_i_p
-            provisioning_data.hostname = hostname
+            provisioning_data.hostname = _get_instance_ip(instance, allocate_public_ip)
             provisioning_data.internal_ip = instance.network_interfaces[0].network_i_p
             return
         raise ProvisioningError(
@@ -500,7 +499,7 @@ class GCPCompute(
         request.instance_resource = gcp_resources.create_instance_struct(
             disk_size=10,
             image_id=_get_gateway_image_id(),
-            machine_type="e2-small",
+            machine_type="e2-medium",
             accelerators=[],
             spot=False,
             user_data=get_gateway_user_data(configuration.ssh_key_pub),
@@ -512,6 +511,7 @@ class GCPCompute(
             service_account=self.config.vm_service_account,
             network=self.config.vpc_resource_name,
             subnetwork=subnetwork,
+            allocate_public_ip=configuration.public_ip,
         )
         operation = self.instances_client.insert(request=request)
         gcp_resources.wait_for_extended_operation(operation, "instance creation")
@@ -522,7 +522,7 @@ class GCPCompute(
             instance_id=instance_name,
             region=configuration.region,  # used for instance termination
             availability_zone=zone,
-            ip_address=instance.network_interfaces[0].access_configs[0].nat_i_p,
+            ip_address=_get_instance_ip(instance, configuration.public_ip),
             backend_data=json.dumps({"zone": zone}),
         )
 
@@ -1024,3 +1024,9 @@ def _is_tpu_provisioning_data(provisioning_data: JobProvisioningData) -> bool:
         backend_data_dict = json.loads(provisioning_data.backend_data)
         is_tpu = backend_data_dict.get("is_tpu", False)
     return is_tpu
+
+
+def _get_instance_ip(instance: Instance, public_ip: bool) -> str:
+    if public_ip:
+        return instance.network_interfaces[0].access_configs[0].nat_i_p
+    return instance.network_interfaces[0].network_i_p

dstack/_internal/core/backends/lambdalabs/compute.py

@@ -1,4 +1,5 @@
 import hashlib
+import shlex
 import subprocess
 import tempfile
 from threading import Thread
@@ -98,7 +99,7 @@ class LambdaCompute(
                 arch=provisioning_data.instance_type.resources.cpu_arch,
             )
             # shim is assumed to be run under root
-            launch_command = "sudo sh -c '" + "&& ".join(commands) + "'"
+            launch_command = "sudo sh -c " + shlex.quote(" && ".join(commands))
             thread = Thread(
                 target=_start_runner,
                 kwargs={

dstack/_internal/core/compatibility/fleets.py

@@ -1,19 +1,20 @@
-from typing import Any, Dict, Optional
+from typing import Optional
 
+from dstack._internal.core.models.common import IncludeExcludeDictType, IncludeExcludeSetType
 from dstack._internal.core.models.fleets import ApplyFleetPlanInput, FleetSpec
 from dstack._internal.core.models.instances import Instance
 
 
-def get_get_plan_excludes(fleet_spec: FleetSpec) -> Dict:
-    get_plan_excludes = {}
+def get_get_plan_excludes(fleet_spec: FleetSpec) -> IncludeExcludeDictType:
+    get_plan_excludes: IncludeExcludeDictType = {}
     spec_excludes = get_fleet_spec_excludes(fleet_spec)
     if spec_excludes:
         get_plan_excludes["spec"] = spec_excludes
     return get_plan_excludes
 
 
-def get_apply_plan_excludes(plan_input: ApplyFleetPlanInput) -> Dict:
-    apply_plan_excludes = {}
+def get_apply_plan_excludes(plan_input: ApplyFleetPlanInput) -> IncludeExcludeDictType:
+    apply_plan_excludes: IncludeExcludeDictType = {}
     spec_excludes = get_fleet_spec_excludes(plan_input.spec)
     if spec_excludes:
         apply_plan_excludes["spec"] = spec_excludes
@@ -28,23 +29,23 @@ def get_apply_plan_excludes(plan_input: ApplyFleetPlanInput) -> Dict:
     return {"plan": apply_plan_excludes}
 
 
-def get_create_fleet_excludes(fleet_spec: FleetSpec) -> Dict:
-    create_fleet_excludes = {}
+def get_create_fleet_excludes(fleet_spec: FleetSpec) -> IncludeExcludeDictType:
+    create_fleet_excludes: IncludeExcludeDictType = {}
     spec_excludes = get_fleet_spec_excludes(fleet_spec)
     if spec_excludes:
         create_fleet_excludes["spec"] = spec_excludes
     return create_fleet_excludes
 
 
-def get_fleet_spec_excludes(fleet_spec: FleetSpec) -> Optional[Dict]:
+def get_fleet_spec_excludes(fleet_spec: FleetSpec) -> Optional[IncludeExcludeDictType]:
     """
     Returns `fleet_spec` exclude mapping to exclude certain fields from the request.
     Use this method to exclude new fields when they are not set to keep
     clients backward-compatibility with older servers.
     """
-    spec_excludes: Dict[str, Any] = {}
-    configuration_excludes: Dict[str, Any] = {}
-    profile_excludes: set[str] = set()
+    spec_excludes: IncludeExcludeDictType = {}
+    configuration_excludes: IncludeExcludeDictType = {}
+    profile_excludes: IncludeExcludeSetType = set()
     profile = fleet_spec.profile
     if profile.fleets is None:
         profile_excludes.add("fleets")

dstack/_internal/core/compatibility/gateways.py

@@ -1,34 +1,35 @@
-from typing import Dict
-
+from dstack._internal.core.models.common import IncludeExcludeDictType
 from dstack._internal.core.models.gateways import GatewayConfiguration, GatewaySpec
 
 
-def get_gateway_spec_excludes(gateway_spec: GatewaySpec) -> Dict:
+def get_gateway_spec_excludes(gateway_spec: GatewaySpec) -> IncludeExcludeDictType:
     """
     Returns `gateway_spec` exclude mapping to exclude certain fields from the request.
     Use this method to exclude new fields when they are not set to keep
     clients backward-compatibility with older servers.
     """
-    spec_excludes = {}
+    spec_excludes: IncludeExcludeDictType = {}
     spec_excludes["configuration"] = _get_gateway_configuration_excludes(
         gateway_spec.configuration
     )
     return spec_excludes
 
 
-def get_create_gateway_excludes(configuration: GatewayConfiguration) -> Dict:
+def get_create_gateway_excludes(configuration: GatewayConfiguration) -> IncludeExcludeDictType:
     """
     Returns an exclude mapping to exclude certain fields from the create gateway request.
     Use this method to exclude new fields when they are not set to keep
     clients backward-compatibility with older servers.
     """
-    create_gateway_excludes = {}
+    create_gateway_excludes: IncludeExcludeDictType = {}
     create_gateway_excludes["configuration"] = _get_gateway_configuration_excludes(configuration)
     return create_gateway_excludes
 
 
-def _get_gateway_configuration_excludes(configuration: GatewayConfiguration) -> Dict:
-    configuration_excludes = {}
+def _get_gateway_configuration_excludes(
+    configuration: GatewayConfiguration,
+) -> IncludeExcludeDictType:
+    configuration_excludes: IncludeExcludeDictType = {}
     if configuration.tags is None:
         configuration_excludes["tags"] = True
     return configuration_excludes

dstack/_internal/core/compatibility/logs.py

@@ -1,15 +1,16 @@
-from typing import Dict, Optional
+from typing import Optional
 
+from dstack._internal.core.models.common import IncludeExcludeDictType
 from dstack._internal.server.schemas.logs import PollLogsRequest
 
 
-def get_poll_logs_excludes(request: PollLogsRequest) -> Optional[Dict]:
+def get_poll_logs_excludes(request: PollLogsRequest) -> Optional[IncludeExcludeDictType]:
     """
     Returns exclude mapping to exclude certain fields from the request.
     Use this method to exclude new fields when they are not set to keep
     clients backward-compatibility with older servers.
     """
-    excludes = {}
+    excludes: IncludeExcludeDictType = {}
     if request.next_token is None:
         excludes["next_token"] = True
     return excludes if excludes else None

dstack/_internal/core/compatibility/runs.py

@@ -1,29 +1,39 @@
-from typing import Any, Dict, Optional
+from typing import Optional
 
+from dstack._internal.core.models.common import IncludeExcludeDictType, IncludeExcludeSetType
 from dstack._internal.core.models.configurations import ServiceConfiguration
 from dstack._internal.core.models.runs import ApplyRunPlanInput, JobSpec, JobSubmission, RunSpec
-from dstack._internal.server.schemas.runs import GetRunPlanRequest
+from dstack._internal.server.schemas.runs import GetRunPlanRequest, ListRunsRequest
 
 
-def get_apply_plan_excludes(plan: ApplyRunPlanInput) -> Optional[Dict]:
+def get_list_runs_excludes(list_runs_request: ListRunsRequest) -> IncludeExcludeSetType:
+    excludes = set()
+    if list_runs_request.include_jobs:
+        excludes.add("include_jobs")
+    if list_runs_request.job_submissions_limit is None:
+        excludes.add("job_submissions_limit")
+    return excludes
+
+
+def get_apply_plan_excludes(plan: ApplyRunPlanInput) -> Optional[IncludeExcludeDictType]:
     """
     Returns `plan` exclude mapping to exclude certain fields from the request.
     Use this method to exclude new fields when they are not set to keep
     clients backward-compatibility with older servers.
     """
-    apply_plan_excludes = {}
+    apply_plan_excludes: IncludeExcludeDictType = {}
     run_spec_excludes = get_run_spec_excludes(plan.run_spec)
     if run_spec_excludes is not None:
         apply_plan_excludes["run_spec"] = run_spec_excludes
     current_resource = plan.current_resource
     if current_resource is not None:
-        current_resource_excludes = {}
+        current_resource_excludes: IncludeExcludeDictType = {}
         current_resource_excludes["status_message"] = True
         if current_resource.deployment_num == 0:
             current_resource_excludes["deployment_num"] = True
         apply_plan_excludes["current_resource"] = current_resource_excludes
         current_resource_excludes["run_spec"] = get_run_spec_excludes(current_resource.run_spec)
-        job_submissions_excludes = {}
+        job_submissions_excludes: IncludeExcludeDictType = {}
         current_resource_excludes["jobs"] = {
             "__all__": {
                 "job_spec": get_job_spec_excludes([job.job_spec for job in current_resource.jobs]),
@@ -45,7 +55,7 @@ def get_apply_plan_excludes(plan: ApplyRunPlanInput) -> Optional[Dict]:
             job_submissions_excludes["deployment_num"] = True
         latest_job_submission = current_resource.latest_job_submission
         if latest_job_submission is not None:
-            latest_job_submission_excludes = {}
+            latest_job_submission_excludes: IncludeExcludeDictType = {}
             current_resource_excludes["latest_job_submission"] = latest_job_submission_excludes
             if _should_exclude_job_submission_jpd_cpu_arch(latest_job_submission):
                 latest_job_submission_excludes["job_provisioning_data"] = {
@@ -62,12 +72,12 @@ def get_apply_plan_excludes(plan: ApplyRunPlanInput) -> Optional[Dict]:
     return {"plan": apply_plan_excludes}
 
 
-def get_get_plan_excludes(request: GetRunPlanRequest) -> Optional[Dict]:
+def get_get_plan_excludes(request: GetRunPlanRequest) -> Optional[IncludeExcludeDictType]:
     """
     Excludes new fields when they are not set to keep
     clients backward-compatibility with older servers.
     """
-    get_plan_excludes = {}
+    get_plan_excludes: IncludeExcludeDictType = {}
     run_spec_excludes = get_run_spec_excludes(request.run_spec)
     if run_spec_excludes is not None:
         get_plan_excludes["run_spec"] = run_spec_excludes
@@ -76,15 +86,15 @@ def get_get_plan_excludes(request: GetRunPlanRequest) -> Optional[Dict]:
     return get_plan_excludes
 
 
-def get_run_spec_excludes(run_spec: RunSpec) -> Optional[Dict]:
+def get_run_spec_excludes(run_spec: RunSpec) -> IncludeExcludeDictType:
     """
     Returns `run_spec` exclude mapping to exclude certain fields from the request.
     Use this method to exclude new fields when they are not set to keep
     clients backward-compatibility with older servers.
     """
-    spec_excludes: dict[str, Any] = {}
-    configuration_excludes: dict[str, Any] = {}
-    profile_excludes: set[str] = set()
+    spec_excludes: IncludeExcludeDictType = {}
+    configuration_excludes: IncludeExcludeDictType = {}
+    profile_excludes: IncludeExcludeSetType = set()
     configuration = run_spec.configuration
     profile = run_spec.profile
 
@@ -121,18 +131,16 @@ def get_run_spec_excludes(run_spec: RunSpec) -> Optional[Dict]:
         spec_excludes["configuration"] = configuration_excludes
     if profile_excludes:
         spec_excludes["profile"] = profile_excludes
-    if spec_excludes:
-        return spec_excludes
-    return None
+    return spec_excludes
 
 
-def get_job_spec_excludes(job_specs: list[JobSpec]) -> Optional[dict]:
+def get_job_spec_excludes(job_specs: list[JobSpec]) -> IncludeExcludeDictType:
     """
     Returns `job_spec` exclude mapping to exclude certain fields from the request.
     Use this method to exclude new fields when they are not set to keep
     clients backward-compatibility with older servers.
     """
-    spec_excludes: dict[str, Any] = {}
+    spec_excludes: IncludeExcludeDictType = {}
 
     if all(s.repo_code_hash is None for s in job_specs):
         spec_excludes["repo_code_hash"] = True
@@ -140,10 +148,10 @@ def get_job_spec_excludes(job_specs: list[JobSpec]) -> Optional[dict]:
         spec_excludes["repo_data"] = True
     if all(not s.file_archives for s in job_specs):
         spec_excludes["file_archives"] = True
+    if all(s.service_port is None for s in job_specs):
+        spec_excludes["service_port"] = True
 
-    if spec_excludes:
-        return spec_excludes
-    return None
+    return spec_excludes
 
 
 def _should_exclude_job_submission_jpd_cpu_arch(job_submission: JobSubmission) -> bool:

dstack/_internal/core/compatibility/volumes.py

@@ -1,32 +1,35 @@
-from typing import Dict
-
+from dstack._internal.core.models.common import IncludeExcludeDictType
 from dstack._internal.core.models.volumes import VolumeConfiguration, VolumeSpec
 
 
-def get_volume_spec_excludes(volume_spec: VolumeSpec) -> Dict:
+def get_volume_spec_excludes(volume_spec: VolumeSpec) -> IncludeExcludeDictType:
     """
     Returns `volume_spec` exclude mapping to exclude certain fields from the request.
     Use this method to exclude new fields when they are not set to keep
     clients backward-compatibility with older servers.
     """
-    spec_excludes = {}
+    spec_excludes: IncludeExcludeDictType = {}
     spec_excludes["configuration"] = _get_volume_configuration_excludes(volume_spec.configuration)
     return spec_excludes
 
 
-def get_create_volume_excludes(configuration: VolumeConfiguration) -> Dict:
+def get_create_volume_excludes(configuration: VolumeConfiguration) -> IncludeExcludeDictType:
     """
     Returns an exclude mapping to exclude certain fields from the create volume request.
     Use this method to exclude new fields when they are not set to keep
     clients backward-compatibility with older servers.
     """
-    create_volume_excludes = {}
+    create_volume_excludes: IncludeExcludeDictType = {}
     create_volume_excludes["configuration"] = _get_volume_configuration_excludes(configuration)
     return create_volume_excludes
 
 
-def _get_volume_configuration_excludes(configuration: VolumeConfiguration) -> Dict:
-    configuration_excludes = {}
+def _get_volume_configuration_excludes(
+    configuration: VolumeConfiguration,
+) -> IncludeExcludeDictType:
+    configuration_excludes: IncludeExcludeDictType = {}
     if configuration.tags is None:
         configuration_excludes["tags"] = True
+    if configuration.auto_cleanup_duration is None:
+        configuration_excludes["auto_cleanup_duration"] = True
     return configuration_excludes

dstack/_internal/core/errors.py

@@ -110,6 +110,10 @@ class PlacementGroupInUseError(ComputeError):
     pass
 
 
+class PlacementGroupNotSupportedError(ComputeError):
+    pass
+
+
 class NotYetTerminated(ComputeError):
     """
     Used by Compute.terminate_instance to signal that instance termination is not complete

dstack/_internal/core/models/common.py

@@ -1,11 +1,21 @@
 import re
 from enum import Enum
-from typing import Union
+from typing import Any, Callable, Optional, Union
 
+import orjson
 from pydantic import Field
 from pydantic_duality import DualBaseModel
 from typing_extensions import Annotated
 
+from dstack._internal.utils.json_utils import pydantic_orjson_dumps
+
+IncludeExcludeFieldType = Union[int, str]
+IncludeExcludeSetType = set[IncludeExcludeFieldType]
+IncludeExcludeDictType = dict[
+    IncludeExcludeFieldType, Union[bool, IncludeExcludeSetType, "IncludeExcludeDictType"]
+]
+IncludeExcludeType = Union[IncludeExcludeSetType, IncludeExcludeDictType]
+
 
 # DualBaseModel creates two classes for the model:
 # one with extra = "forbid" (CoreModel/CoreModel.__request__),
@@ -13,7 +23,40 @@ from typing_extensions import Annotated
 # This allows to use the same model both for a strict parsing of the user input and
 # for a permissive parsing of the server responses.
 class CoreModel(DualBaseModel):
-    pass
+    class Config:
+        json_loads = orjson.loads
+        json_dumps = pydantic_orjson_dumps
+
+    def json(
+        self,
+        *,
+        include: Optional[IncludeExcludeType] = None,
+        exclude: Optional[IncludeExcludeType] = None,
+        by_alias: bool = False,
+        skip_defaults: Optional[bool] = None,  # ignore as it's deprecated
+        exclude_unset: bool = False,
+        exclude_defaults: bool = False,
+        exclude_none: bool = False,
+        encoder: Optional[Callable[[Any], Any]] = None,
+        models_as_dict: bool = True,  # does not seems to be needed by dstack or dependencies
+        **dumps_kwargs: Any,
+    ) -> str:
+        """
+        Override `json()` method so that it calls `dict()`.
+        Allows changing how models are serialized by overriding `dict()` only.
+        By default, `json()` won't call `dict()`, so changes applied in `dict()` won't take place.
+        """
+        data = self.dict(
+            by_alias=by_alias,
+            include=include,
+            exclude=exclude,
+            exclude_unset=exclude_unset,
+            exclude_defaults=exclude_defaults,
+            exclude_none=exclude_none,
+        )
+        if self.__custom_root_type__:
+            data = data["__root__"]
+        return self.__config__.json_dumps(data, default=encoder, **dumps_kwargs)
 
 
 class Duration(int):

dstack/_internal/core/models/configurations.py

@@ -4,6 +4,7 @@ from enum import Enum
 from pathlib import PurePosixPath
 from typing import Any, Dict, List, Optional, Union
 
+import orjson
 from pydantic import Field, ValidationError, conint, constr, root_validator, validator
 from typing_extensions import Annotated, Literal
 
@@ -18,6 +19,9 @@ from dstack._internal.core.models.resources import Range, ResourcesSpec
 from dstack._internal.core.models.services import AnyModel, OpenAIChatModel
 from dstack._internal.core.models.unix import UnixUser
 from dstack._internal.core.models.volumes import MountPoint, VolumeConfiguration, parse_mount_point
+from dstack._internal.utils.json_utils import (
+    pydantic_orjson_dumps_with_indent,
+)
 
 CommandsList = List[str]
 ValidPort = conint(gt=0, le=65536)
@@ -394,8 +398,9 @@ class TaskConfiguration(
 
 class ServiceConfigurationParams(CoreModel):
     port: Annotated[
+        # NOTE: it's a PortMapping for historical reasons. Only `port.container_port` is used.
         Union[ValidPort, constr(regex=r"^[0-9]+:[0-9]+$"), PortMapping],
-        Field(description="The port, that application listens on or the mapping"),
+        Field(description="The port the application listens on"),
     ]
     gateway: Annotated[
         Optional[Union[bool, str]],
@@ -573,6 +578,9 @@ class DstackConfiguration(CoreModel):
     ]
 
     class Config:
+        json_loads = orjson.loads
+        json_dumps = pydantic_orjson_dumps_with_indent
+
         @staticmethod
         def schema_extra(schema: Dict[str, Any]):
             schema["$schema"] = "http://json-schema.org/draft-07/schema#"

dstack/_internal/core/models/fleets.py

@@ -8,7 +8,7 @@ from pydantic import Field, root_validator, validator
 from typing_extensions import Annotated, Literal
 
 from dstack._internal.core.models.backends.base import BackendType
-from dstack._internal.core.models.common import CoreModel
+from dstack._internal.core.models.common import ApplyAction, CoreModel
 from dstack._internal.core.models.envs import Env
 from dstack._internal.core.models.instances import Instance, InstanceOfferWithAvailability, SSHKey
 from dstack._internal.core.models.profiles import (
@@ -324,6 +324,7 @@ class FleetPlan(CoreModel):
     offers: List[InstanceOfferWithAvailability]
     total_offers: int
     max_offer_price: Optional[float] = None
+    action: Optional[ApplyAction] = None  # default value for backward compatibility
 
     def get_effective_spec(self) -> FleetSpec:
         if self.effective_spec is not None:

dstack/_internal/core/models/profiles.py

@@ -1,12 +1,14 @@
 from enum import Enum
 from typing import Any, Dict, List, Optional, Union, overload
 
+import orjson
 from pydantic import Field, root_validator, validator
 from typing_extensions import Annotated, Literal
 
 from dstack._internal.core.models.backends.base import BackendType
 from dstack._internal.core.models.common import CoreModel, Duration
 from dstack._internal.utils.common import list_enum_values_for_annotation
+from dstack._internal.utils.json_utils import pydantic_orjson_dumps_with_indent
 from dstack._internal.utils.tags import tags_validator
 
 DEFAULT_RETRY_DURATION = 3600
@@ -74,11 +76,9 @@ def parse_off_duration(v: Optional[Union[int, str, bool]]) -> Optional[Union[str
     return parse_duration(v)
 
 
-def parse_idle_duration(v: Optional[Union[int, str, bool]]) -> Optional[Union[str, int, bool]]:
-    if v is False:
+def parse_idle_duration(v: Optional[Union[int, str]]) -> Optional[Union[str, int]]:
+    if v == "off" or v == -1:
         return -1
-    if v is True:
-        return None
     return parse_duration(v)
 
 
@@ -249,7 +249,7 @@ class ProfileParams(CoreModel):
         ),
     ] = None
     idle_duration: Annotated[
-        Optional[Union[Literal["off"], str, int, bool]],
+        Optional[Union[Literal["off"], str, int]],
         Field(
             description=(
                 "Time to wait before terminating idle instances."
@@ -343,6 +343,9 @@ class ProfilesConfig(CoreModel):
     profiles: List[Profile]
 
     class Config:
+        json_loads = orjson.loads
+        json_dumps = pydantic_orjson_dumps_with_indent
+
         schema_extra = {"$schema": "http://json-schema.org/draft-07/schema#"}
 
     def default(self) -> Optional[Profile]:

dstack/_internal/core/models/resources.py

@@ -382,14 +382,6 @@ class ResourcesSpec(CoreModel):
     gpu: Annotated[Optional[GPUSpec], Field(description="The GPU requirements")] = None
     disk: Annotated[Optional[DiskSpec], Field(description="The disk resources")] = DEFAULT_DISK
 
-    # TODO: Remove in 0.20. Added for backward compatibility.
-    @root_validator
-    def _post_validate(cls, values):
-        cpu = values.get("cpu")
-        if isinstance(cpu, CPUSpec) and cpu.arch in [None, gpuhunt.CPUArchitecture.X86]:
-            values["cpu"] = cpu.count
-        return values
-
     def pretty_format(self) -> str:
         # TODO: Remove in 0.20. Use self.cpu directly
         cpu = parse_obj_as(CPUSpec, self.cpu)
@@ -407,3 +399,18 @@ class ResourcesSpec(CoreModel):
             resources.update(disk_size=self.disk.size)
         res = pretty_resources(**resources)
         return res
+
+    def dict(self, *args, **kwargs) -> Dict:
+        # super() does not work with pydantic-duality
+        res = CoreModel.dict(self, *args, **kwargs)
+        self._update_serialized_cpu(res)
+        return res
+
+    # TODO: Remove in 0.20. Added for backward compatibility.
+    def _update_serialized_cpu(self, values: Dict):
+        cpu = values["cpu"]
+        if cpu:
+            arch = cpu.get("arch")
+            count = cpu.get("count")
+            if count and arch in [None, gpuhunt.CPUArchitecture.X86.value]:
+                values["cpu"] = count