dstack 0.19.6rc1__py3-none-any.whl → 0.19.8__py3-none-any.whl

dstack/_internal/server/services/gateways/__init__.py

@@ -31,13 +31,19 @@ from dstack._internal.core.models.gateways import (
     Gateway,
     GatewayComputeConfiguration,
     GatewayConfiguration,
+    GatewaySpec,
     GatewayStatus,
     LetsEncryptGatewayCertificate,
 )
 from dstack._internal.core.services import validate_dstack_resource_name
 from dstack._internal.server import settings
 from dstack._internal.server.db import get_db
-from dstack._internal.server.models import GatewayComputeModel, GatewayModel, ProjectModel
+from dstack._internal.server.models import (
+    GatewayComputeModel,
+    GatewayModel,
+    ProjectModel,
+    UserModel,
+)
 from dstack._internal.server.services.backends import (
     check_backend_type_available,
     get_project_backend_by_type_or_error,
@@ -50,6 +56,7 @@ from dstack._internal.server.services.locking import (
     get_locker,
     string_to_lock_id,
 )
+from dstack._internal.server.services.plugins import apply_plugin_policies
 from dstack._internal.server.utils.common import gather_map_async
 from dstack._internal.utils.common import get_current_datetime, run_async
 from dstack._internal.utils.crypto import generate_rsa_key_pair_bytes
@@ -129,9 +136,17 @@ async def create_gateway_compute(
 
 async def create_gateway(
     session: AsyncSession,
+    user: UserModel,
     project: ProjectModel,
     configuration: GatewayConfiguration,
 ) -> Gateway:
+    spec = apply_plugin_policies(
+        user=user.name,
+        project=project.name,
+        # Create pseudo spec until the gateway API is updated to accept spec
+        spec=GatewaySpec(configuration=configuration),
+    )
+    configuration = spec.configuration
     _validate_gateway_configuration(configuration)
 
     backend_model, _ = await get_project_backend_with_model_by_type_or_error(
@@ -140,7 +155,7 @@ async def create_gateway(
 
     lock_namespace = f"gateway_names_{project.name}"
     if get_db().dialect_name == "sqlite":
-        # Start new transaction to see commited changes after lock
+        # Start new transaction to see committed changes after lock
         await session.commit()
     elif get_db().dialect_name == "postgresql":
         await session.execute(

dstack/_internal/server/services/instances.py

@@ -408,7 +408,6 @@ async def create_instance_model(
     requirements: Requirements,
     instance_name: str,
     instance_num: int,
-    placement_group_name: Optional[str],
     reservation: Optional[str],
     blocks: Union[Literal["auto"], int],
     tags: Optional[Dict[str, str]],
@@ -427,7 +426,6 @@ async def create_instance_model(
         user=user.name,
         ssh_keys=[project_ssh_key],
         instance_id=str(instance_id),
-        placement_group_name=placement_group_name,
         reservation=reservation,
         tags=tags,
     )

dstack/_internal/server/services/offers.py

@@ -8,12 +8,14 @@ from dstack._internal.core.backends import (
     BACKENDS_WITH_RESERVATION_SUPPORT,
 )
 from dstack._internal.core.backends.base.backend import Backend
+from dstack._internal.core.backends.base.compute import ComputeWithPlacementGroupSupport
 from dstack._internal.core.models.backends.base import BackendType
 from dstack._internal.core.models.instances import (
     InstanceOfferWithAvailability,
     InstanceType,
     Resources,
 )
+from dstack._internal.core.models.placement import PlacementGroup
 from dstack._internal.core.models.profiles import Profile
 from dstack._internal.core.models.runs import JobProvisioningData, Requirements
 from dstack._internal.core.models.volumes import Volume
@@ -31,6 +33,7 @@ async def get_offers_by_requirements(
     volumes: Optional[List[List[Volume]]] = None,
     privileged: bool = False,
     instance_mounts: bool = False,
+    placement_group: Optional[PlacementGroup] = None,
     blocks: Union[int, Literal["auto"]] = 1,
 ) -> List[Tuple[Backend, InstanceOfferWithAvailability]]:
     backends: List[Backend] = await backends_services.get_project_backends(project=project)
@@ -116,6 +119,18 @@ async def get_offers_by_requirements(
                     new_offers.append((b, new_offer))
         offers = new_offers
 
+    if placement_group is not None:
+        new_offers = []
+        for b, o in offers:
+            for backend in backends:
+                compute = backend.compute()
+                if isinstance(
+                    compute, ComputeWithPlacementGroupSupport
+                ) and compute.is_suitable_placement_group(placement_group, o):
+                    new_offers.append((b, o))
+                    break
+        offers = new_offers
+
     if profile.instance_types is not None:
         instance_types = [i.lower() for i in profile.instance_types]
         offers = [(b, o) for b, o in offers if o.instance.name.lower() in instance_types]

dstack/_internal/server/services/placement.py

@@ -1,8 +1,9 @@
+from collections.abc import Iterable
 from typing import Optional
 from uuid import UUID
 
 from git import List
-from sqlalchemy import select
+from sqlalchemy import and_, select, update
 from sqlalchemy.ext.asyncio import AsyncSession
 
 from dstack._internal.core.models.placement import (
@@ -13,15 +14,35 @@ from dstack._internal.core.models.placement import (
 from dstack._internal.server.models import PlacementGroupModel
 
 
-async def get_fleet_placement_groups(
+async def get_fleet_placement_group_models(
     session: AsyncSession,
     fleet_id: UUID,
-) -> List[PlacementGroup]:
+) -> List[PlacementGroupModel]:
     res = await session.execute(
-        select(PlacementGroupModel).where(PlacementGroupModel.fleet_id == fleet_id)
+        select(PlacementGroupModel).where(
+            and_(
+                PlacementGroupModel.fleet_id == fleet_id,
+                PlacementGroupModel.deleted == False,
+                PlacementGroupModel.fleet_deleted == False,
+            )
+        )
+    )
+    return list(res.scalars().all())
+
+
+async def schedule_fleet_placement_groups_deletion(
+    session: AsyncSession, fleet_id: UUID, except_placement_group_ids: Iterable[UUID] = ()
+) -> None:
+    await session.execute(
+        update(PlacementGroupModel)
+        .where(
+            and_(
+                PlacementGroupModel.fleet_id == fleet_id,
+                PlacementGroupModel.id.not_in(except_placement_group_ids),
+            )
+        )
+        .values(fleet_deleted=True)  # TODO: rename `fleet_deleted` -> `to_be_deleted`
     )
-    placement_groups = res.scalars().all()
-    return [placement_group_model_to_placement_group(pg) for pg in placement_groups]
 
 
 def placement_group_model_to_placement_group(

dstack/_internal/server/services/plugins.py

@@ -0,0 +1,77 @@
+import itertools
+from importlib import import_module
+
+from backports.entry_points_selectable import entry_points  # backport for Python 3.9
+
+from dstack._internal.core.errors import ServerClientError
+from dstack._internal.utils.logging import get_logger
+from dstack.plugins import ApplyPolicy, ApplySpec, Plugin
+
+logger = get_logger(__name__)
+
+
+_PLUGINS: list[Plugin] = []
+
+
+def load_plugins(enabled_plugins: list[str]):
+    _PLUGINS.clear()
+    plugins_entrypoints = entry_points(group="dstack.plugins")
+    plugins_to_load = enabled_plugins.copy()
+    for entrypoint in plugins_entrypoints:
+        if entrypoint.name not in enabled_plugins:
+            logger.info(
+                ("Found not enabled plugin %s. Plugin will not be loaded."),
+                entrypoint.name,
+            )
+            continue
+        try:
+            module_path, _, class_name = entrypoint.value.partition(":")
+            module = import_module(module_path)
+        except ImportError:
+            logger.warning(
+                (
+                    "Failed to load plugin %s when importing %s."
+                    " Ensure the module is on the import path."
+                ),
+                entrypoint.name,
+                entrypoint.value,
+            )
+            continue
+        plugin_class = getattr(module, class_name, None)
+        if plugin_class is None:
+            logger.warning(
+                ("Failed to load plugin %s: plugin class %s not found in module %s."),
+                entrypoint.name,
+                class_name,
+                module_path,
+            )
+            continue
+        if not issubclass(plugin_class, Plugin):
+            logger.warning(
+                ("Failed to load plugin %s: plugin class %s is not a subclass of Plugin."),
+                entrypoint.name,
+                class_name,
+            )
+            continue
+        plugins_to_load.remove(entrypoint.name)
+        _PLUGINS.append(plugin_class())
+        logger.info("Loaded plugin %s", entrypoint.name)
+    if plugins_to_load:
+        logger.warning("Enabled plugins not found: %s", plugins_to_load)
+
+
+def apply_plugin_policies(user: str, project: str, spec: ApplySpec) -> ApplySpec:
+    policies = _get_apply_policies()
+    for policy in policies:
+        try:
+            spec = policy.on_apply(user=user, project=project, spec=spec)
+        except ValueError as e:
+            msg = None
+            if len(e.args) > 0:
+                msg = e.args[0]
+            raise ServerClientError(msg)
+    return spec
+
+
+def _get_apply_policies() -> list[ApplyPolicy]:
+    return list(itertools.chain(*[p.get_apply_policies() for p in _PLUGINS]))

dstack/_internal/server/services/resources.py

@@ -0,0 +1,21 @@
+import gpuhunt
+from pydantic import parse_obj_as
+
+from dstack._internal.core.models.resources import CPUSpec, ResourcesSpec
+
+
+def set_resources_defaults(resources: ResourcesSpec) -> None:
+    # TODO: Remove in 0.20. Use resources.cpu directly
+    cpu = parse_obj_as(CPUSpec, resources.cpu)
+    if cpu.arch is None:
+        gpu = resources.gpu
+        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))
+        ):
+            cpu.arch = gpuhunt.CPUArchitecture.ARM
+        else:
+            cpu.arch = gpuhunt.CPUArchitecture.X86
+        resources.cpu = cpu

dstack/_internal/server/services/runs.py

@@ -79,7 +79,9 @@ from dstack._internal.server.services.jobs import (
 from dstack._internal.server.services.locking import get_locker, string_to_lock_id
 from dstack._internal.server.services.logging import fmt
 from dstack._internal.server.services.offers import get_offers_by_requirements
+from dstack._internal.server.services.plugins import apply_plugin_policies
 from dstack._internal.server.services.projects import list_project_models, list_user_project_models
+from dstack._internal.server.services.resources import set_resources_defaults
 from dstack._internal.server.services.users import get_user_model_by_name
 from dstack._internal.utils.logging import get_logger
 from dstack._internal.utils.random_names import generate_name
@@ -279,7 +281,14 @@ async def get_plan(
     run_spec: RunSpec,
     max_offers: Optional[int],
 ) -> RunPlan:
+    # Spec must be copied by parsing to calculate merged_profile
     effective_run_spec = RunSpec.parse_obj(run_spec.dict())
+    effective_run_spec = apply_plugin_policies(
+        user=user.name,
+        project=project.name,
+        spec=effective_run_spec,
+    )
+    effective_run_spec = RunSpec.parse_obj(effective_run_spec.dict())
     _validate_run_spec_and_set_defaults(effective_run_spec)
 
     profile = effective_run_spec.merged_profile
@@ -293,12 +302,14 @@ async def get_plan(
             project=project,
             run_name=effective_run_spec.run_name,
         )
-        if (
-            current_resource is not None
-            and not current_resource.status.is_finished()
-            and _can_update_run_spec(current_resource.run_spec, effective_run_spec)
-        ):
-            action = ApplyAction.UPDATE
+        if current_resource is not None:
+            # For backward compatibility (current_resource may has been submitted before
+            # some fields, e.g., CPUSpec.arch, were added)
+            set_resources_defaults(current_resource.run_spec.configuration.resources)
+            if not current_resource.status.is_finished() and _can_update_run_spec(
+                current_resource.run_spec, effective_run_spec
+            ):
+                action = ApplyAction.UPDATE
 
     jobs = await get_jobs_from_run_spec(effective_run_spec, replica_num=0)
 
@@ -370,34 +381,48 @@ async def apply_plan(
     plan: ApplyRunPlanInput,
     force: bool,
 ) -> Run:
-    _validate_run_spec_and_set_defaults(plan.run_spec)
-    if plan.run_spec.run_name is None:
+    run_spec = plan.run_spec
+    run_spec = apply_plugin_policies(
+        user=user.name,
+        project=project.name,
+        spec=run_spec,
+    )
+    # Spec must be copied by parsing to calculate merged_profile
+    run_spec = RunSpec.parse_obj(run_spec.dict())
+    _validate_run_spec_and_set_defaults(run_spec)
+    if run_spec.run_name is None:
         return await submit_run(
             session=session,
             user=user,
             project=project,
-            run_spec=plan.run_spec,
+            run_spec=run_spec,
         )
     current_resource = await get_run_by_name(
         session=session,
         project=project,
-        run_name=plan.run_spec.run_name,
+        run_name=run_spec.run_name,
     )
     if current_resource is None or current_resource.status.is_finished():
         return await submit_run(
             session=session,
             user=user,
             project=project,
-            run_spec=plan.run_spec,
+            run_spec=run_spec,
         )
+
+    # For backward compatibility (current_resource may has been submitted before
+    # some fields, e.g., CPUSpec.arch, were added)
+    set_resources_defaults(current_resource.run_spec.configuration.resources)
     try:
-        _check_can_update_run_spec(current_resource.run_spec, plan.run_spec)
+        _check_can_update_run_spec(current_resource.run_spec, run_spec)
     except ServerClientError:
         # The except is only needed to raise an appropriate error if run is active
         if not current_resource.status.is_finished():
             raise ServerClientError("Cannot override active run. Stop the run first.")
         raise
     if not force:
+        if plan.current_resource is not None:
+            set_resources_defaults(plan.current_resource.run_spec.configuration.resources)
         if (
             plan.current_resource is None
             or plan.current_resource.id != current_resource.id
@@ -409,14 +434,12 @@ async def apply_plan(
     # FIXME: potentially long write transaction
     # Avoid getting run_model after update
     await session.execute(
-        update(RunModel)
-        .where(RunModel.id == current_resource.id)
-        .values(run_spec=plan.run_spec.json())
+        update(RunModel).where(RunModel.id == current_resource.id).values(run_spec=run_spec.json())
     )
     run = await get_run_by_name(
         session=session,
         project=project,
-        run_name=plan.run_spec.run_name,
+        run_name=run_spec.run_name,
     )
     return common_utils.get_or_error(run)
 
@@ -436,7 +459,7 @@ async def submit_run(
 
     lock_namespace = f"run_names_{project.name}"
     if get_db().dialect_name == "sqlite":
-        # Start new transaction to see commited changes after lock
+        # Start new transaction to see committed changes after lock
         await session.commit()
     elif get_db().dialect_name == "postgresql":
         await session.execute(
@@ -852,6 +875,7 @@ def _validate_run_spec_and_set_defaults(run_spec: RunSpec):
         raise ServerClientError(
             f"Maximum utilization_policy.time_window is {settings.SERVER_METRICS_TTL_SECONDS}s"
         )
+    set_resources_defaults(run_spec.configuration.resources)
 
 
 _UPDATABLE_SPEC_FIELDS = ["repo_code_hash", "configuration"]

dstack/_internal/server/services/volumes.py

@@ -21,6 +21,7 @@ from dstack._internal.core.models.volumes import (
     VolumeConfiguration,
     VolumeInstance,
     VolumeProvisioningData,
+    VolumeSpec,
     VolumeStatus,
 )
 from dstack._internal.core.services import validate_dstack_resource_name
@@ -38,6 +39,7 @@ from dstack._internal.server.services.locking import (
     get_locker,
     string_to_lock_id,
 )
+from dstack._internal.server.services.plugins import apply_plugin_policies
 from dstack._internal.server.services.projects import list_project_models, list_user_project_models
 from dstack._internal.utils import common, random_names
 from dstack._internal.utils.logging import get_logger
@@ -203,11 +205,18 @@ async def create_volume(
     user: UserModel,
     configuration: VolumeConfiguration,
 ) -> Volume:
+    spec = apply_plugin_policies(
+        user=user.name,
+        project=project.name,
+        # Create pseudo spec until the volume API is updated to accept spec
+        spec=VolumeSpec(configuration=configuration),
+    )
+    configuration = spec.configuration
     _validate_volume_configuration(configuration)
 
     lock_namespace = f"volume_names_{project.name}"
     if get_db().dialect_name == "sqlite":
-        # Start new transaction to see commited changes after lock
+        # Start new transaction to see committed changes after lock
         await session.commit()
     elif get_db().dialect_name == "postgresql":
         await session.execute(

dstack/_internal/server/testing/common.py

@@ -2,7 +2,7 @@ import json
 import uuid
 from contextlib import contextmanager
 from datetime import datetime, timezone
-from typing import Dict, List, Optional, Union
+from typing import Dict, List, Literal, Optional, Union
 from uuid import UUID
 
 import gpuhunt
@@ -25,7 +25,12 @@ from dstack._internal.core.models.configurations import (
     DevEnvironmentConfiguration,
 )
 from dstack._internal.core.models.envs import Env
-from dstack._internal.core.models.fleets import FleetConfiguration, FleetSpec, FleetStatus
+from dstack._internal.core.models.fleets import (
+    FleetConfiguration,
+    FleetSpec,
+    FleetStatus,
+    InstanceGroupPlacement,
+)
 from dstack._internal.core.models.gateways import GatewayComputeConfiguration, GatewayStatus
 from dstack._internal.core.models.instances import (
     Disk,
@@ -51,7 +56,7 @@ from dstack._internal.core.models.profiles import (
 )
 from dstack._internal.core.models.repos.base import RepoType
 from dstack._internal.core.models.repos.local import LocalRunRepoData
-from dstack._internal.core.models.resources import Memory, Range, ResourcesSpec
+from dstack._internal.core.models.resources import CPUSpec, Memory, Range, ResourcesSpec
 from dstack._internal.core.models.runs import (
     JobProvisioningData,
     JobRuntimeData,
@@ -497,10 +502,12 @@ def get_fleet_spec(conf: Optional[FleetConfiguration] = None) -> FleetSpec:
 def get_fleet_configuration(
     name: str = "test-fleet",
     nodes: Range[int] = Range(min=1, max=1),
+    placement: Optional[InstanceGroupPlacement] = None,
 ) -> FleetConfiguration:
     return FleetConfiguration(
         name=name,
         nodes=nodes,
+        placement=placement,
     )
 
 
@@ -519,13 +526,13 @@ async def create_instance(
     instance_id: Optional[UUID] = None,
     job: Optional[JobModel] = None,
     instance_num: int = 0,
-    backend: BackendType = BackendType.DATACRUNCH,
+    backend: Optional[BackendType] = BackendType.DATACRUNCH,
     termination_policy: Optional[TerminationPolicy] = None,
     termination_idle_time: int = DEFAULT_FLEET_TERMINATION_IDLE_TIME,
-    region: str = "eu-west",
+    region: Optional[str] = "eu-west",
     remote_connection_info: Optional[RemoteConnectionInfo] = None,
-    offer: Optional[InstanceOfferWithAvailability] = None,
-    job_provisioning_data: Optional[JobProvisioningData] = None,
+    offer: Optional[Union[InstanceOfferWithAvailability, Literal["auto"]]] = "auto",
+    job_provisioning_data: Optional[Union[JobProvisioningData, Literal["auto"]]] = "auto",
     total_blocks: Optional[int] = 1,
     busy_blocks: int = 0,
     name: str = "test_instance",
@@ -534,7 +541,7 @@ async def create_instance(
 ) -> InstanceModel:
     if instance_id is None:
         instance_id = uuid.uuid4()
-    if job_provisioning_data is None:
+    if job_provisioning_data == "auto":
         job_provisioning_data = get_job_provisioning_data(
             dockerized=True,
             backend=backend,
@@ -543,13 +550,13 @@ async def create_instance(
             hostname="running_instance.ip",
             internal_ip=None,
         )
-    if offer is None:
+    if offer == "auto":
         offer = get_instance_offer_with_availability(backend=backend, region=region, spot=spot)
     if profile is None:
         profile = Profile(name="test_name")
 
     if requirements is None:
-        requirements = Requirements(resources=ResourcesSpec(cpu=1))
+        requirements = Requirements(resources=ResourcesSpec(cpu=CPUSpec.parse("1")))
 
     if instance_configuration is None:
         instance_configuration = get_instance_configuration()
@@ -571,8 +578,8 @@ async def create_instance(
         created_at=created_at,
         started_at=created_at,
         finished_at=finished_at,
-        job_provisioning_data=job_provisioning_data.json(),
-        offer=offer.json(),
+        job_provisioning_data=job_provisioning_data.json() if job_provisioning_data else None,
+        offer=offer.json() if offer else None,
         price=price,
         region=region,
         backend=backend,
@@ -659,20 +666,7 @@ def get_remote_connection_info(
     env: Optional[Union[Env, dict]] = None,
 ):
     if ssh_keys is None:
-        ssh_keys = [
-            SSHKey(
-                public="ssh-ed25519 AAAAC3NzaC1lZDI1NTE5AAAAIO6mJxVbNtm0zXgMLvByrhXJCmJRveSrJxLB5/OzcyCk",
-                private="""
-                    -----BEGIN OPENSSH PRIVATE KEY-----
-                    b3BlbnNzaC1rZXktdjEAAAAABG5vbmUAAAAEbm9uZQAAAAAAAAABAAAAMwAAAAtzc2gtZW
-                    QyNTUxOQAAACDupicVWzbZtM14DC7wcq4VyQpiUb3kqycSwefzs3MgpAAAAJCiWa5Volmu
-                    VQAAAAtzc2gtZWQyNTUxOQAAACDupicVWzbZtM14DC7wcq4VyQpiUb3kqycSwefzs3MgpA
-                    AAAEAncHi4AhS6XdMp5Gzd+IMse/4ekyQ54UngByf0Sp0uH+6mJxVbNtm0zXgMLvByrhXJ
-                    CmJRveSrJxLB5/OzcyCkAAAACWRlZkBkZWZwYwECAwQ=
-                    -----END OPENSSH PRIVATE KEY-----
-                """,
-            )
-        ]
+        ssh_keys = [get_ssh_key()]
     if env is None:
         env = Env()
     elif isinstance(env, dict):
@@ -686,6 +680,21 @@ def get_remote_connection_info(
     )
 
 
+def get_ssh_key() -> SSHKey:
+    return SSHKey(
+        public="ssh-ed25519 AAAAC3NzaC1lZDI1NTE5AAAAIO6mJxVbNtm0zXgMLvByrhXJCmJRveSrJxLB5/OzcyCk",
+        private="""
+                    -----BEGIN OPENSSH PRIVATE KEY-----
+                    b3BlbnNzaC1rZXktdjEAAAAABG5vbmUAAAAEbm9uZQAAAAAAAAABAAAAMwAAAAtzc2gtZW
+                    QyNTUxOQAAACDupicVWzbZtM14DC7wcq4VyQpiUb3kqycSwefzs3MgpAAAAJCiWa5Volmu
+                    VQAAAAtzc2gtZWQyNTUxOQAAACDupicVWzbZtM14DC7wcq4VyQpiUb3kqycSwefzs3MgpA
+                    AAAEAncHi4AhS6XdMp5Gzd+IMse/4ekyQ54UngByf0Sp0uH+6mJxVbNtm0zXgMLvByrhXJ
+                    CmJRveSrJxLB5/OzcyCkAAAACWRlZkBkZWZwYwECAwQ=
+                    -----END OPENSSH PRIVATE KEY-----
+                """,
+    )
+
+
 async def create_volume(
     session: AsyncSession,
     project: ProjectModel,

dstack/_internal/utils/common.py

@@ -1,4 +1,5 @@
 import asyncio
+import enum
 import itertools
 import re
 import time
@@ -83,6 +84,8 @@ def pretty_date(time: datetime) -> str:
 
 
 def pretty_resources(
+    *,
+    cpu_arch: Optional[Any] = None,
     cpus: Optional[Any] = None,
     memory: Optional[Any] = None,
     gpu_count: Optional[Any] = None,
@@ -110,25 +113,35 @@ def pretty_resources(
     """
     parts = []
     if cpus is not None:
-        parts.append(f"{cpus}xCPU")
+        cpu_arch_lower: Optional[str] = None
+        if isinstance(cpu_arch, enum.Enum):
+            cpu_arch_lower = str(cpu_arch.value).lower()
+        elif isinstance(cpu_arch, str):
+            cpu_arch_lower = cpu_arch.lower()
+        if cpu_arch_lower == "arm":
+            cpu_arch_prefix = "arm:"
+        else:
+            cpu_arch_prefix = ""
+        parts.append(f"cpu={cpu_arch_prefix}{cpus}")
     if memory is not None:
-        parts.append(f"{memory}")
+        parts.append(f"mem={memory}")
+    if disk_size:
+        parts.append(f"disk={disk_size}")
     if gpu_count:
         gpu_parts = []
+        gpu_parts.append(f"{gpu_name or 'gpu'}")
         if gpu_memory is not None:
             gpu_parts.append(f"{gpu_memory}")
+        if gpu_count is not None:
+            gpu_parts.append(f"{gpu_count}")
         if total_gpu_memory is not None:
-            gpu_parts.append(f"total {total_gpu_memory}")
+            gpu_parts.append(f"{total_gpu_memory}")
         if compute_capability is not None:
             gpu_parts.append(f"{compute_capability}")
 
-        gpu = f"{gpu_count}x{gpu_name or 'GPU'}"
-        if gpu_parts:
-            gpu += f" ({', '.join(gpu_parts)})"
+        gpu = ":".join(gpu_parts)
         parts.append(gpu)
-    if disk_size:
-        parts.append(f"{disk_size} (disk)")
-    return ", ".join(parts)
+    return " ".join(parts)
 
 
 def since(timestamp: str) -> datetime:

dstack/_internal/utils/json_schema.py

@@ -1,6 +1,9 @@
 def add_extra_schema_types(schema_property: dict, extra_types: list[dict]):
     if "allOf" in schema_property:
-        ref = schema_property.pop("allOf")[0]
+        refs = [schema_property.pop("allOf")[0]]
+    elif "anyOf" in schema_property:
+        refs = schema_property.pop("anyOf")
     else:
-        ref = {"type": schema_property.pop("type")}
-    schema_property["anyOf"] = [ref, *extra_types]
+        refs = [{"type": schema_property.pop("type")}]
+    refs.extend(extra_types)
+    schema_property["anyOf"] = refs

dstack/api/__init__.py

@@ -14,6 +14,7 @@ from dstack._internal.core.models.repos.local import LocalRepo
 from dstack._internal.core.models.repos.remote import RemoteRepo
 from dstack._internal.core.models.repos.virtual import VirtualRepo
 from dstack._internal.core.models.resources import ComputeCapability, Memory, Range
+from dstack._internal.core.models.resources import CPUSpec as CPU
 from dstack._internal.core.models.resources import DiskSpec as Disk
 from dstack._internal.core.models.resources import GPUSpec as GPU
 from dstack._internal.core.models.resources import ResourcesSpec as Resources