dstack 0.18.40rc1__py3-none-any.whl → 0.18.42__py3-none-any.whl

dstack/_internal/server/testing/common.py

@@ -5,6 +5,7 @@ from datetime import datetime, timezone
 from typing import Dict, List, Optional, Union
 from uuid import UUID
 
+import gpuhunt
 from sqlalchemy.ext.asyncio import AsyncSession
 
 from dstack._internal.core.models.backends.base import BackendType
@@ -13,14 +14,20 @@ from dstack._internal.core.models.configurations import (
     AnyRunConfiguration,
     DevEnvironmentConfiguration,
 )
+from dstack._internal.core.models.envs import Env
 from dstack._internal.core.models.fleets import FleetConfiguration, FleetSpec, FleetStatus
 from dstack._internal.core.models.gateways import GatewayStatus
 from dstack._internal.core.models.instances import (
+    Disk,
+    Gpu,
+    InstanceAvailability,
     InstanceConfiguration,
+    InstanceOfferWithAvailability,
     InstanceStatus,
     InstanceType,
     RemoteConnectionInfo,
     Resources,
+    SSHKey,
 )
 from dstack._internal.core.models.placement import (
     PlacementGroupConfiguration,
@@ -48,7 +55,7 @@ from dstack._internal.core.models.runs import (
 from dstack._internal.core.models.users import GlobalRole
 from dstack._internal.core.models.volumes import (
     Volume,
-    VolumeAttachmentData,
+    VolumeAttachment,
     VolumeConfiguration,
     VolumeProvisioningData,
     VolumeStatus,
@@ -69,6 +76,7 @@ from dstack._internal.server.models import (
     RepoModel,
     RunModel,
     UserModel,
+    VolumeAttachmentModel,
     VolumeModel,
 )
 from dstack._internal.server.services.jobs import get_job_specs_from_run_spec
@@ -311,17 +319,30 @@ async def create_job(
     return job
 
 
-def get_job_provisioning_data(dockerized: bool = False) -> JobProvisioningData:
+def get_job_provisioning_data(
+    dockerized: bool = False,
+    backend: BackendType = BackendType.AWS,
+    region: str = "us-east-1",
+    gpu_count: int = 0,
+    cpu_count: int = 1,
+    memory_gib: float = 0.5,
+    spot: bool = False,
+    hostname: str = "127.0.0.4",
+    internal_ip: Optional[str] = "127.0.0.4",
+) -> JobProvisioningData:
+    gpus = [Gpu(name="T4", memory_mib=16384, vendor=gpuhunt.AcceleratorVendor.NVIDIA)] * gpu_count
     return JobProvisioningData(
-        backend=BackendType.AWS,
+        backend=backend,
         instance_type=InstanceType(
             name="instance",
-            resources=Resources(cpus=1, memory_mib=512, spot=False, gpus=[]),
+            resources=Resources(
+                cpus=cpu_count, memory_mib=int(memory_gib * 1024), spot=spot, gpus=gpus
+            ),
         ),
         instance_id="instance_id",
-        hostname="127.0.0.4",
-        internal_ip="127.0.0.4",
-        region="us-east-1",
+        hostname=hostname,
+        internal_ip=internal_ip,
+        region=region,
         price=10.5,
         username="ubuntu",
         ssh_port=22,
@@ -337,6 +358,8 @@ def get_job_runtime_data(
     gpu: Optional[int] = None,
     memory: Optional[float] = None,
     ports: Optional[dict[int, int]] = None,
+    offer: Optional[InstanceOfferWithAvailability] = None,
+    volume_names: Optional[list[str]] = None,
 ) -> JobRuntimeData:
     return JobRuntimeData(
         network_mode=NetworkMode(network_mode),
@@ -344,6 +367,8 @@ def get_job_runtime_data(
         gpu=gpu,
         memory=Memory(memory) if memory is not None else None,
         ports=ports,
+        offer=offer,
+        volume_names=volume_names,
     )
 
 
@@ -481,56 +506,26 @@ async def create_instance(
     termination_idle_time: int = DEFAULT_POOL_TERMINATION_IDLE_TIME,
     region: str = "eu-west",
     remote_connection_info: Optional[RemoteConnectionInfo] = None,
+    offer: Optional[InstanceOfferWithAvailability] = None,
     job_provisioning_data: Optional[JobProvisioningData] = None,
+    total_blocks: Optional[int] = 1,
+    busy_blocks: int = 0,
     name: str = "test_instance",
     volumes: Optional[List[VolumeModel]] = None,
 ) -> InstanceModel:
     if instance_id is None:
         instance_id = uuid.uuid4()
     if job_provisioning_data is None:
-        job_provisioning_data_dict = {
-            "backend": backend.value,
-            "instance_type": {
-                "name": "instance",
-                "resources": {
-                    "cpus": 1,
-                    "memory_mib": 512,
-                    "gpus": [],
-                    "spot": spot,
-                    "disk": {"size_mib": 102400},
-                    "description": "",
-                },
-            },
-            "instance_id": "running_instance.id",
-            "ssh_proxy": None,
-            "hostname": "running_instance.ip",
-            "region": region,
-            "price": 0.1,
-            "username": "root",
-            "ssh_port": 22,
-            "dockerized": True,
-            "backend_data": None,
-        }
-    else:
-        job_provisioning_data_dict = job_provisioning_data.dict()
-    offer = {
-        "backend": backend.value,
-        "instance": {
-            "name": "instance",
-            "resources": {
-                "cpus": 2,
-                "memory_mib": 12000,
-                "gpus": [],
-                "spot": spot,
-                "disk": {"size_mib": 102400},
-                "description": "",
-            },
-        },
-        "region": region,
-        "price": 1,
-        "availability": "available",
-    }
-
+        job_provisioning_data = get_job_provisioning_data(
+            dockerized=True,
+            backend=backend,
+            region=region,
+            spot=spot,
+            hostname="running_instance.ip",
+            internal_ip=None,
+        )
+    if offer is None:
+        offer = get_instance_offer_with_availability(backend=backend, region=region, spot=spot)
     if profile is None:
         profile = Profile(name="test_name")
 
@@ -548,6 +543,9 @@ async def create_instance(
 
     if volumes is None:
         volumes = []
+    volume_attachments = []
+    for volume in volumes:
+        volume_attachments.append(VolumeAttachmentModel(volume=volume))
 
     im = InstanceModel(
         id=instance_id,
@@ -561,8 +559,8 @@ async def create_instance(
         created_at=created_at,
         started_at=created_at,
         finished_at=finished_at,
-        job_provisioning_data=json.dumps(job_provisioning_data_dict),
-        offer=json.dumps(offer),
+        job_provisioning_data=job_provisioning_data.json(),
+        offer=offer.json(),
         price=1,
         region=region,
         backend=backend,
@@ -572,14 +570,87 @@ async def create_instance(
         requirements=requirements.json(),
         instance_configuration=instance_configuration.json(),
         remote_connection_info=remote_connection_info.json() if remote_connection_info else None,
-        job=job,
-        volumes=volumes,
+        volume_attachments=volume_attachments,
+        total_blocks=total_blocks,
+        busy_blocks=busy_blocks,
     )
+    if job:
+        im.jobs.append(job)
     session.add(im)
     await session.commit()
     return im
 
 
+def get_instance_offer_with_availability(
+    backend: BackendType = BackendType.AWS,
+    region: str = "eu-west",
+    gpu_count: int = 0,
+    cpu_count: int = 2,
+    memory_gib: float = 12,
+    disk_gib: float = 100.0,
+    spot: bool = False,
+    blocks: int = 1,
+    total_blocks: int = 1,
+    availability_zones: Optional[List[str]] = None,
+):
+    gpus = [Gpu(name="T4", memory_mib=16384, vendor=gpuhunt.AcceleratorVendor.NVIDIA)] * gpu_count
+    return InstanceOfferWithAvailability(
+        backend=backend,
+        instance=InstanceType(
+            name="instance",
+            resources=Resources(
+                cpus=cpu_count,
+                memory_mib=int(memory_gib * 1024),
+                gpus=gpus,
+                spot=spot,
+                disk=Disk(size_mib=int(disk_gib * 1024)),
+                description="",
+            ),
+        ),
+        region=region,
+        price=1,
+        availability=InstanceAvailability.AVAILABLE,
+        availability_zones=availability_zones,
+        blocks=blocks,
+        total_blocks=total_blocks,
+    )
+
+
+def get_remote_connection_info(
+    host: str = "10.0.0.10",
+    port: int = 22,
+    ssh_user: str = "ubuntu",
+    ssh_keys: Optional[list[SSHKey]] = None,
+    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-----
+                """,
+            )
+        ]
+    if env is None:
+        env = Env()
+    elif isinstance(env, dict):
+        env = Env.parse_obj(env)
+    return RemoteConnectionInfo(
+        host=host,
+        port=port,
+        ssh_user=ssh_user,
+        ssh_keys=ssh_keys,
+        env=env,
+    )
+
+
 async def create_volume(
     session: AsyncSession,
     project: ProjectModel,
@@ -604,7 +675,7 @@ async def create_volume(
         volume_provisioning_data=volume_provisioning_data.json()
         if volume_provisioning_data
         else None,
-        instances=[],
+        attachments=[],
         deleted_at=deleted_at,
         deleted=True if deleted_at else False,
     )
@@ -626,16 +697,14 @@ def get_volume(
     deleted: bool = False,
     volume_id: Optional[str] = None,
     provisioning_data: Optional[VolumeProvisioningData] = None,
-    attachment_data: Optional[VolumeAttachmentData] = None,
-    device_name: Optional[str] = None,
+    attachments: Optional[List[VolumeAttachment]] = None,
 ) -> Volume:
     if id_ is None:
         id_ = uuid.uuid4()
     if configuration is None:
         configuration = get_volume_configuration()
-    if device_name is not None:
-        assert attachment_data is None, "attachment_data and device_name are mutually exclusive"
-        attachment_data = VolumeAttachmentData(device_name=device_name)
+    if attachments is None:
+        attachments = []
     return Volume(
         id=id_,
         name=name,
@@ -649,7 +718,7 @@ def get_volume(
         deleted=deleted,
         volume_id=volume_id,
         provisioning_data=provisioning_data,
-        attachment_data=attachment_data,
+        attachments=attachments,
     )
 
 

dstack/_internal/utils/common.py

@@ -157,24 +157,38 @@ def parse_pretty_duration(duration: str) -> int:
     return amount * multiplier
 
 
+DURATION_UNITS_DESC = [
+    ("w", 7 * 24 * 3600),
+    ("d", 24 * 3600),
+    ("h", 3600),
+    ("m", 60),
+    ("s", 1),
+]
+
+
 def format_pretty_duration(seconds: int) -> str:
     if seconds == 0:
         return "0s"
     if seconds < 0:
         raise ValueError("Seconds cannot be negative")
-    units = [
-        ("w", 7 * 24 * 3600),
-        ("d", 24 * 3600),
-        ("h", 3600),
-        ("m", 60),
-        ("s", 1),
-    ]
-    for unit, multiplier in units:
+    for unit, multiplier in DURATION_UNITS_DESC:
         if seconds % multiplier == 0:
             return f"{seconds // multiplier}{unit}"
     return f"{seconds}s"  # Fallback to seconds if no larger unit fits perfectly
 
 
+def format_duration_multiunit(seconds: int) -> str:
+    """90 -> 1m 30s, 4545 -> 1h 15m 45s, etc"""
+    if seconds < 0:
+        raise ValueError("Seconds cannot be negative")
+    result = ""
+    for unit, multiplier in DURATION_UNITS_DESC:
+        if unit_value := seconds // multiplier:
+            result += f" {unit_value}{unit}"
+            seconds -= unit_value * multiplier
+    return result.lstrip() or "0s"
+
+
 def sizeof_fmt(num, suffix="B"):
     for unit in ["", "Ki", "Mi", "Gi", "Ti", "Pi", "Ei", "Zi"]:
         if abs(num) < 1024.0:

dstack/_internal/utils/env.py

@@ -0,0 +1,14 @@
+import os
+
+
+def get_bool(name: str, default: bool = False) -> bool:
+    try:
+        value = os.environ[name]
+    except KeyError:
+        return default
+    value = value.lower()
+    if value in ["0", "false", "off"]:
+        return False
+    if value in ["1", "true", "on"]:
+        return True
+    raise ValueError(f"Invalid bool value: {name}={value}")

dstack/_internal/utils/ssh.py

@@ -159,7 +159,7 @@ def get_ssh_config(path: PathLike, host: str) -> Optional[Dict[str, str]]:
         return None
 
 
-def update_ssh_config(path: PathLike, host: str, options: Dict[str, Union[str, FilePath]]):
+def update_ssh_config(path: PathLike, host: str, options: Dict[str, Union[str, int, FilePath]]):
     Path(path).parent.mkdir(parents=True, exist_ok=True)
     with FileLock(str(path) + ".lock"):
         copy_mode = True

dstack/api/server/_fleets.py

@@ -62,16 +62,29 @@ def _get_fleet_spec_excludes(fleet_spec: FleetSpec) -> Optional[_ExcludeDict]:
     spec_excludes: _ExcludeDict = {}
     configuration_excludes: _ExcludeDict = {}
     profile_excludes: set[str] = set()
+    ssh_config_excludes: _ExcludeDict = {}
+    ssh_hosts_excludes: set[str] = set()
 
     # TODO: Can be removed in 0.19
     if fleet_spec.configuration_path is None:
         spec_excludes["configuration_path"] = True
     if fleet_spec.configuration.ssh_config is not None:
+        if fleet_spec.configuration.ssh_config.proxy_jump is None:
+            ssh_config_excludes["proxy_jump"] = True
+        if all(
+            isinstance(h, str) or h.proxy_jump is None
+            for h in fleet_spec.configuration.ssh_config.hosts
+        ):
+            ssh_hosts_excludes.add("proxy_jump")
         if all(
             isinstance(h, str) or h.internal_ip is None
             for h in fleet_spec.configuration.ssh_config.hosts
         ):
-            configuration_excludes["ssh_config"] = {"hosts": {"__all__": {"internal_ip"}}}
+            ssh_hosts_excludes.add("internal_ip")
+        if all(
+            isinstance(h, str) or h.blocks == 1 for h in fleet_spec.configuration.ssh_config.hosts
+        ):
+            ssh_hosts_excludes.add("blocks")
     # client >= 0.18.30 / server <= 0.18.29 compatibility tweak
     if fleet_spec.configuration.reservation is None:
         configuration_excludes["reservation"] = True
@@ -84,7 +97,18 @@ def _get_fleet_spec_excludes(fleet_spec: FleetSpec) -> Optional[_ExcludeDict]:
     # client >= 0.18.38 / server <= 0.18.37 compatibility tweak
     if fleet_spec.profile is not None and fleet_spec.profile.stop_duration is None:
         profile_excludes.add("stop_duration")
+    # client >= 0.18.41 / server <= 0.18.40 compatibility tweak
+    if fleet_spec.configuration.availability_zones is None:
+        configuration_excludes["availability_zones"] = True
+    if fleet_spec.profile is not None and fleet_spec.profile.availability_zones is None:
+        profile_excludes.add("availability_zones")
+    if fleet_spec.configuration.blocks == 1:
+        configuration_excludes["blocks"] = True
 
+    if ssh_hosts_excludes:
+        ssh_config_excludes["hosts"] = {"__all__": ssh_hosts_excludes}
+    if ssh_config_excludes:
+        configuration_excludes["ssh_config"] = ssh_config_excludes
     if configuration_excludes:
         spec_excludes["configuration"] = configuration_excludes
     if profile_excludes:

dstack/api/server/_runs.py

@@ -7,6 +7,7 @@ from pydantic import parse_obj_as
 from dstack._internal.core.models.common import is_core_model_instance
 from dstack._internal.core.models.configurations import (
     STRIP_PREFIX_DEFAULT,
+    DevEnvironmentConfiguration,
     ServiceConfiguration,
 )
 from dstack._internal.core.models.pools import Instance
@@ -82,7 +83,10 @@ class RunsAPIClient(APIClientGroup):
     ) -> Run:
         plan_input: ApplyRunPlanInput = ApplyRunPlanInput.__response__.parse_obj(plan)
         body = ApplyRunPlanRequest(plan=plan_input, force=force)
-        resp = self._request(f"/api/project/{project_name}/runs/apply", body=body.json())
+        resp = self._request(
+            f"/api/project/{project_name}/runs/apply",
+            body=body.json(exclude=_get_apply_plan_excludes(plan_input)),
+        )
         return parse_obj_as(Run.__response__, resp.json())
 
     def submit(self, project_name: str, run_spec: RunSpec) -> Run:
@@ -121,8 +125,15 @@ class RunsAPIClient(APIClientGroup):
         return parse_obj_as(Instance.__response__, resp.json())
 
 
+def _get_apply_plan_excludes(plan: ApplyRunPlanInput) -> Optional[dict]:
+    run_spec_excludes = _get_run_spec_excludes(plan.run_spec)
+    if run_spec_excludes is not None:
+        return {"plan": run_spec_excludes}
+    return None
+
+
 def _get_run_spec_excludes(run_spec: RunSpec) -> Optional[dict]:
-    spec_excludes: dict[str, set[str]] = {}
+    spec_excludes: dict[str, Any] = {}
     configuration_excludes: dict[str, Any] = {}
     profile_excludes: set[str] = set()
     configuration = run_spec.configuration
@@ -164,6 +175,16 @@ def _get_run_spec_excludes(run_spec: RunSpec) -> Optional[dict]:
         for v in configuration.volumes
     ):
         configuration_excludes["volumes"] = {"__all__": {"optional"}}
+    # client >= 0.18.41 / server <= 0.18.40 compatibility tweak
+    if configuration.availability_zones is None:
+        configuration_excludes["availability_zones"] = True
+    if profile is not None and profile.availability_zones is None:
+        profile_excludes.add("availability_zones")
+    if (
+        is_core_model_instance(configuration, DevEnvironmentConfiguration)
+        and configuration.inactivity_duration is None
+    ):
+        configuration_excludes["inactivity_duration"] = True
 
     if configuration_excludes:
         spec_excludes["configuration"] = configuration_excludes

dstack/api/server/_volumes.py

@@ -27,9 +27,20 @@ class VolumesAPIClient(APIClientGroup):
         configuration: VolumeConfiguration,
     ) -> Volume:
         body = CreateVolumeRequest(configuration=configuration)
-        resp = self._request(f"/api/project/{project_name}/volumes/create", body=body.json())
+        resp = self._request(
+            f"/api/project/{project_name}/volumes/create",
+            body=body.json(exclude=_get_volume_configuration_excludes(configuration)),
+        )
         return parse_obj_as(Volume.__response__, resp.json())
 
     def delete(self, project_name: str, names: List[str]) -> None:
         body = DeleteVolumesRequest(names=names)
         self._request(f"/api/project/{project_name}/volumes/delete", body=body.json())
+
+
+def _get_volume_configuration_excludes(configuration: VolumeConfiguration) -> dict:
+    configuration_excludes = {}
+    # client >= 0.18.41 / server <= 0.18.40 compatibility tweak
+    if configuration.availability_zone is None:
+        configuration_excludes["availability_zone"] = True
+    return {"configuration": configuration_excludes}

dstack/version.py

@@ -1,3 +1,3 @@
-__version__ = "0.18.40rc1"
+__version__ = "0.18.42"
 __is_release__ = True
 base_image = "0.6"

{dstack-0.18.40rc1.dist-info → dstack-0.18.42.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: dstack
-Version: 0.18.40rc1
+Version: 0.18.42
 Summary: dstack is an open-source orchestration engine for running AI workloads on any cloud or on-premises.
 Home-page: https://dstack.ai
 Author: Andrey Cheptsov