dstack 0.19.12rc1__py3-none-any.whl → 0.19.14__py3-none-any.whl

dstack/_internal/server/testing/common.py

@@ -140,6 +140,7 @@ async def create_project(
     created_at: datetime = datetime(2023, 1, 2, 3, 4, tzinfo=timezone.utc),
     ssh_private_key: str = "",
     ssh_public_key: str = "",
+    is_public: bool = False,
 ) -> ProjectModel:
     if owner is None:
         owner = await create_user(session=session, name="test_owner")
@@ -149,6 +150,7 @@ async def create_project(
         created_at=created_at,
         ssh_private_key=ssh_private_key,
         ssh_public_key=ssh_public_key,
+        is_public=is_public,
     )
     session.add(project)
     await session.commit()
@@ -263,6 +265,7 @@ async def create_run(
     run_id: Optional[UUID] = None,
     deleted: bool = False,
     priority: int = 0,
+    deployment_num: int = 0,
 ) -> RunModel:
     if run_spec is None:
         run_spec = get_run_spec(
@@ -284,6 +287,8 @@ async def create_run(
         last_processed_at=submitted_at,
         jobs=[],
         priority=priority,
+        deployment_num=deployment_num,
+        desired_replica_count=1,
     )
     session.add(run)
     await session.commit()
@@ -303,18 +308,23 @@ async def create_job(
     instance: Optional[InstanceModel] = None,
     job_num: int = 0,
     replica_num: int = 0,
+    deployment_num: Optional[int] = None,
     instance_assigned: bool = False,
     disconnected_at: Optional[datetime] = None,
 ) -> JobModel:
+    if deployment_num is None:
+        deployment_num = run.deployment_num
     run_spec = RunSpec.parse_raw(run.run_spec)
     job_spec = (await get_job_specs_from_run_spec(run_spec, replica_num=replica_num))[0]
+    job_spec.job_num = job_num
     job = JobModel(
         project_id=run.project_id,
         run_id=run.id,
         run_name=run.run_name,
         job_num=job_num,
-        job_name=run.run_name + f"-0-{replica_num}",
+        job_name=run.run_name + f"-{job_num}-{replica_num}",
         replica_num=replica_num,
+        deployment_num=deployment_num,
         submission_num=submission_num,
         submitted_at=submitted_at,
         last_processed_at=last_processed_at,

dstack/_internal/settings.py

@@ -14,6 +14,9 @@ DSTACK_USE_LATEST_FROM_BRANCH = os.getenv("DSTACK_USE_LATEST_FROM_BRANCH") is no
 
 DSTACK_BASE_IMAGE = os.getenv("DSTACK_BASE_IMAGE", "dstackai/base")
 DSTACK_BASE_IMAGE_VERSION = os.getenv("DSTACK_BASE_IMAGE_VERSION", version.base_image)
+DSTACK_BASE_IMAGE_UBUNTU_VERSION = os.getenv(
+    "DSTACK_BASE_IMAGE_UBUNTU_VERSION", version.base_image_ubuntu_version
+)
 
 
 class FeatureFlags:

dstack/_internal/utils/common.py

@@ -314,3 +314,7 @@ def make_proxy_url(server_url: str, proxy_url: str) -> str:
         path=concat_url_path(server.path, proxy.path),
     )
     return proxy.geturl()
+
+
+def list_enum_values_for_annotation(enum_class: type[enum.Enum]) -> str:
+    return ", ".join(f"`{e.value}`" for e in enum_class)

dstack/api/_public/runs.py

@@ -31,6 +31,7 @@ from dstack._internal.core.models.resources import ResourcesSpec
 from dstack._internal.core.models.runs import (
     Job,
     JobSpec,
+    JobStatus,
     RunPlan,
     RunSpec,
     RunStatus,
@@ -184,7 +185,7 @@ class Run(ABC):
         self,
         start_time: Optional[datetime] = None,
         diagnose: bool = False,
-        replica_num: int = 0,
+        replica_num: Optional[int] = None,
         job_num: int = 0,
     ) -> Iterable[bytes]:
         """
@@ -246,7 +247,7 @@ class Run(ABC):
         ssh_identity_file: Optional[PathLike] = None,
         bind_address: Optional[str] = None,
         ports_overrides: Optional[List[PortMapping]] = None,
-        replica_num: int = 0,
+        replica_num: Optional[int] = None,
         job_num: int = 0,
     ) -> bool:
         """
@@ -254,6 +255,7 @@ class Run(ABC):
 
         Args:
             ssh_identity_file: SSH keypair to access instances.
+            replica_num: replica_num or None to attach to any running replica.
 
         Raises:
             dstack.api.PortUsedError: If ports are in use or the run is attached by another process.
@@ -265,7 +267,9 @@ class Run(ABC):
 
         job = self._find_job(replica_num=replica_num, job_num=job_num)
         if job is None:
-            raise ClientError(f"Failed to find replica={replica_num} job={job_num}")
+            replica_repr = replica_num if replica_num is not None else "<any running>"
+            raise ClientError(f"Failed to find replica={replica_repr} job={job_num}")
+        replica_num = job.job_spec.replica_num
 
         name = self.name
         if replica_num != 0 or job_num != 0:
@@ -358,9 +362,14 @@ class Run(ABC):
             self._ssh_attach.detach()
             self._ssh_attach = None
 
-    def _find_job(self, replica_num: int, job_num: int) -> Optional[Job]:
+    def _find_job(self, replica_num: Optional[int], job_num: int) -> Optional[Job]:
         for j in self._run.jobs:
-            if j.job_spec.replica_num == replica_num and j.job_spec.job_num == job_num:
+            if (
+                replica_num is not None
+                and j.job_spec.replica_num == replica_num
+                or replica_num is None
+                and j.job_submissions[-1].status == JobStatus.RUNNING
+            ) and j.job_spec.job_num == job_num:
                 return j
         return None
 

dstack/api/server/_fleets.py

@@ -1,9 +1,13 @@
-from typing import Any, Dict, List, Optional, Union
+from typing import List, Union
 
 from pydantic import parse_obj_as
 
+from dstack._internal.core.compatibility.fleets import (
+    get_apply_plan_excludes,
+    get_create_fleet_excludes,
+    get_get_plan_excludes,
+)
 from dstack._internal.core.models.fleets import ApplyFleetPlanInput, Fleet, FleetPlan, FleetSpec
-from dstack._internal.core.models.instances import Instance
 from dstack._internal.server.schemas.fleets import (
     ApplyFleetPlanRequest,
     CreateFleetRequest,
@@ -34,7 +38,7 @@ class FleetsAPIClient(APIClientGroup):
         spec: FleetSpec,
     ) -> FleetPlan:
         body = GetFleetPlanRequest(spec=spec)
-        body_json = body.json(exclude=_get_get_plan_excludes(spec))
+        body_json = body.json(exclude=get_get_plan_excludes(spec))
         resp = self._request(f"/api/project/{project_name}/fleets/get_plan", body=body_json)
         return parse_obj_as(FleetPlan.__response__, resp.json())
 
@@ -46,7 +50,7 @@ class FleetsAPIClient(APIClientGroup):
     ) -> Fleet:
         plan_input = ApplyFleetPlanInput.__response__.parse_obj(plan)
         body = ApplyFleetPlanRequest(plan=plan_input, force=force)
-        body_json = body.json(exclude=_get_apply_plan_excludes(plan_input))
+        body_json = body.json(exclude=get_apply_plan_excludes(plan_input))
         resp = self._request(f"/api/project/{project_name}/fleets/apply", body=body_json)
         return parse_obj_as(Fleet.__response__, resp.json())
 
@@ -66,70 +70,6 @@ class FleetsAPIClient(APIClientGroup):
         spec: FleetSpec,
     ) -> Fleet:
         body = CreateFleetRequest(spec=spec)
-        body_json = body.json(exclude=_get_create_fleet_excludes(spec))
+        body_json = body.json(exclude=get_create_fleet_excludes(spec))
         resp = self._request(f"/api/project/{project_name}/fleets/create", body=body_json)
         return parse_obj_as(Fleet.__response__, resp.json())
-
-
-def _get_get_plan_excludes(fleet_spec: FleetSpec) -> Dict:
-    get_plan_excludes = {}
-    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 = {}
-    spec_excludes = _get_fleet_spec_excludes(plan_input.spec)
-    if spec_excludes:
-        apply_plan_excludes["spec"] = apply_plan_excludes
-    current_resource = plan_input.current_resource
-    if current_resource is not None:
-        current_resource_excludes = {}
-        apply_plan_excludes["current_resource"] = current_resource_excludes
-        if all(map(_should_exclude_instance_cpu_arch, current_resource.instances)):
-            current_resource_excludes["instances"] = {
-                "__all__": {"instance_type": {"resources": {"cpu_arch"}}}
-            }
-    return {"plan": apply_plan_excludes}
-
-
-def _should_exclude_instance_cpu_arch(instance: Instance) -> bool:
-    try:
-        return instance.instance_type.resources.cpu_arch is None
-    except AttributeError:
-        return True
-
-
-def _get_create_fleet_excludes(fleet_spec: FleetSpec) -> Dict:
-    create_fleet_excludes = {}
-    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]:
-    """
-    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()
-    profile = fleet_spec.profile
-    if profile.fleets is None:
-        profile_excludes.add("fleets")
-    if fleet_spec.configuration.tags is None:
-        configuration_excludes["tags"] = True
-    if profile.tags is None:
-        profile_excludes.add("tags")
-    if configuration_excludes:
-        spec_excludes["configuration"] = configuration_excludes
-    if profile_excludes:
-        spec_excludes["profile"] = profile_excludes
-    if spec_excludes:
-        return spec_excludes
-    return None

dstack/api/server/_gateways.py

@@ -1,7 +1,8 @@
-from typing import Dict, List
+from typing import List
 
 from pydantic import parse_obj_as
 
+from dstack._internal.core.compatibility.gateways import get_create_gateway_excludes
 from dstack._internal.core.models.gateways import Gateway, GatewayConfiguration
 from dstack._internal.server.schemas.gateways import (
     CreateGatewayRequest,
@@ -31,7 +32,7 @@ class GatewaysAPIClient(APIClientGroup):
         body = CreateGatewayRequest(configuration=configuration)
         resp = self._request(
             f"/api/project/{project_name}/gateways/create",
-            body=body.json(exclude=_get_gateway_configuration_excludes(configuration)),
+            body=body.json(exclude=get_create_gateway_excludes(configuration)),
         )
         return parse_obj_as(Gateway.__response__, resp.json())
 
@@ -51,15 +52,3 @@ class GatewaysAPIClient(APIClientGroup):
             f"/api/project/{project_name}/gateways/set_wildcard_domain", body=body.json()
         )
         return parse_obj_as(Gateway.__response__, resp.json())
-
-
-def _get_gateway_configuration_excludes(configuration: GatewayConfiguration) -> Dict:
-    """
-    Returns `configuration` 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.
-    """
-    configuration_excludes = {}
-    if configuration.tags is None:
-        configuration_excludes["tags"] = True
-    return {"configuration": configuration_excludes}

dstack/api/server/_projects.py

@@ -17,8 +17,8 @@ class ProjectsAPIClient(APIClientGroup):
         resp = self._request("/api/projects/list")
         return parse_obj_as(List[Project.__response__], resp.json())
 
-    def create(self, project_name: str) -> Project:
-        body = CreateProjectRequest(project_name=project_name)
+    def create(self, project_name: str, is_public: bool = False) -> Project:
+        body = CreateProjectRequest(project_name=project_name, is_public=is_public)
         resp = self._request("/api/projects/create", body=body.json())
         return parse_obj_as(Project.__response__, resp.json())
 

dstack/api/server/_runs.py

@@ -1,13 +1,12 @@
 from datetime import datetime
-from typing import Any, Dict, List, Optional, Union
+from typing import List, Optional, Union
 from uuid import UUID
 
 from pydantic import parse_obj_as
 
-from dstack._internal.core.models.configurations import ServiceConfiguration
+from dstack._internal.core.compatibility.runs import get_apply_plan_excludes, get_get_plan_excludes
 from dstack._internal.core.models.runs import (
     ApplyRunPlanInput,
-    JobSubmission,
     Run,
     RunPlan,
     RunSpec,
@@ -60,7 +59,7 @@ class RunsAPIClient(APIClientGroup):
         body = GetRunPlanRequest(run_spec=run_spec, max_offers=max_offers)
         resp = self._request(
             f"/api/project/{project_name}/runs/get_plan",
-            body=body.json(exclude=_get_get_plan_excludes(body)),
+            body=body.json(exclude=get_get_plan_excludes(body)),
         )
         return parse_obj_as(RunPlan.__response__, resp.json())
 
@@ -74,7 +73,7 @@ class RunsAPIClient(APIClientGroup):
         body = ApplyRunPlanRequest(plan=plan_input, force=force)
         resp = self._request(
             f"/api/project/{project_name}/runs/apply",
-            body=body.json(exclude=_get_apply_plan_excludes(plan_input)),
+            body=body.json(exclude=get_apply_plan_excludes(plan_input)),
         )
         return parse_obj_as(Run.__response__, resp.json())
 
@@ -85,114 +84,3 @@ class RunsAPIClient(APIClientGroup):
     def delete(self, project_name: str, runs_names: List[str]):
         body = DeleteRunsRequest(runs_names=runs_names)
         self._request(f"/api/project/{project_name}/runs/delete", body=body.json())
-
-
-def _get_apply_plan_excludes(plan: ApplyRunPlanInput) -> Optional[Dict]:
-    """
-    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 = {}
-    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 = {}
-        apply_plan_excludes["current_resource"] = current_resource_excludes
-        current_resource_excludes["run_spec"] = _get_run_spec_excludes(current_resource.run_spec)
-        job_submissions_excludes = {}
-        current_resource_excludes["jobs"] = {
-            "__all__": {"job_submissions": {"__all__": job_submissions_excludes}}
-        }
-        job_submissions = [js for j in current_resource.jobs for js in j.job_submissions]
-        if all(map(_should_exclude_job_submission_jpd_cpu_arch, job_submissions)):
-            job_submissions_excludes["job_provisioning_data"] = {
-                "instance_type": {"resources": {"cpu_arch"}}
-            }
-        if all(map(_should_exclude_job_submission_jrd_cpu_arch, job_submissions)):
-            job_submissions_excludes["job_runtime_data"] = {
-                "offer": {"instance": {"resources": {"cpu_arch"}}}
-            }
-        if all(js.exit_status is None for js in job_submissions):
-            job_submissions_excludes["exit_status"] = True
-        latest_job_submission = current_resource.latest_job_submission
-        if latest_job_submission is not None:
-            latest_job_submission_excludes = {}
-            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"] = {
-                    "instance_type": {"resources": {"cpu_arch"}}
-                }
-            if _should_exclude_job_submission_jrd_cpu_arch(latest_job_submission):
-                latest_job_submission_excludes["job_runtime_data"] = {
-                    "offer": {"instance": {"resources": {"cpu_arch"}}}
-                }
-            if latest_job_submission.exit_status is None:
-                latest_job_submission_excludes["exit_status"] = True
-    return {"plan": apply_plan_excludes}
-
-
-def _should_exclude_job_submission_jpd_cpu_arch(job_submission: JobSubmission) -> bool:
-    try:
-        return job_submission.job_provisioning_data.instance_type.resources.cpu_arch is None
-    except AttributeError:
-        return True
-
-
-def _should_exclude_job_submission_jrd_cpu_arch(job_submission: JobSubmission) -> bool:
-    try:
-        return job_submission.job_runtime_data.offer.instance.resources.cpu_arch is None
-    except AttributeError:
-        return True
-
-
-def _get_get_plan_excludes(request: GetRunPlanRequest) -> Optional[Dict]:
-    """
-    Excludes new fields when they are not set to keep
-    clients backward-compatibility with older servers.
-    """
-    get_plan_excludes = {}
-    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
-    if request.max_offers is None:
-        get_plan_excludes["max_offers"] = True
-    return get_plan_excludes
-
-
-def _get_run_spec_excludes(run_spec: RunSpec) -> Optional[Dict]:
-    """
-    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()
-    configuration = run_spec.configuration
-    profile = run_spec.profile
-
-    if configuration.fleets is None:
-        configuration_excludes["fleets"] = True
-    if profile is not None and profile.fleets is None:
-        profile_excludes.add("fleets")
-    if configuration.tags is None:
-        configuration_excludes["tags"] = True
-    if profile is not None and profile.tags is None:
-        profile_excludes.add("tags")
-    if isinstance(configuration, ServiceConfiguration) and not configuration.rate_limits:
-        configuration_excludes["rate_limits"] = True
-    if configuration.shell is None:
-        configuration_excludes["shell"] = True
-    if configuration.priority is None:
-        configuration_excludes["priority"] = True
-
-    if configuration_excludes:
-        spec_excludes["configuration"] = configuration_excludes
-    if profile_excludes:
-        spec_excludes["profile"] = profile_excludes
-    if spec_excludes:
-        return spec_excludes
-    return None

dstack/api/server/_volumes.py

@@ -1,7 +1,8 @@
-from typing import Dict, List
+from typing import List
 
 from pydantic import parse_obj_as
 
+from dstack._internal.core.compatibility.volumes import get_create_volume_excludes
 from dstack._internal.core.models.volumes import Volume, VolumeConfiguration
 from dstack._internal.server.schemas.volumes import (
     CreateVolumeRequest,
@@ -29,22 +30,10 @@ class VolumesAPIClient(APIClientGroup):
         body = CreateVolumeRequest(configuration=configuration)
         resp = self._request(
             f"/api/project/{project_name}/volumes/create",
-            body=body.json(exclude=_get_volume_configuration_excludes(configuration)),
+            body=body.json(exclude=get_create_volume_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:
-    """
-    Returns `configuration` 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.
-    """
-    configuration_excludes = {}
-    if configuration.tags is None:
-        configuration_excludes["tags"] = True
-    return {"configuration": configuration_excludes}

dstack/plugins/builtin/rest_plugin/_plugin.py

@@ -1,10 +1,14 @@
 import json
 import os
-from typing import Type
+from typing import Dict, Optional, Type
 
 import requests
 from pydantic import ValidationError
 
+from dstack._internal.core.compatibility.fleets import get_fleet_spec_excludes
+from dstack._internal.core.compatibility.gateways import get_gateway_spec_excludes
+from dstack._internal.core.compatibility.runs import get_run_spec_excludes
+from dstack._internal.core.compatibility.volumes import get_volume_spec_excludes
 from dstack._internal.core.errors import ServerClientError
 from dstack._internal.core.models.fleets import FleetSpec
 from dstack._internal.core.models.gateways import GatewaySpec
@@ -44,12 +48,17 @@ class CustomApplyPolicy(ApplyPolicy):
             logger.error(f"Plugin service rejected apply request: {response.error}")
             raise ServerClientError(f"Apply request rejected: {response.error}")
 
-    def _call_plugin_service(self, spec_request: SpecApplyRequest, endpoint: str) -> ApplySpec:
+    def _call_plugin_service(
+        self,
+        spec_request: SpecApplyRequest,
+        endpoint: str,
+        excludes: Optional[Dict],
+    ) -> ApplySpec:
         response = None
         try:
             response = requests.post(
                 f"{self._plugin_service_uri}{endpoint}",
-                json=spec_request.dict(),
+                json=spec_request.dict(exclude={"spec": excludes}),
                 headers={"accept": "application/json", "Content-Type": "application/json"},
                 timeout=PLUGIN_REQUEST_TIMEOUT_SEC,
             )
@@ -75,10 +84,11 @@ class CustomApplyPolicy(ApplyPolicy):
         user: str,
         project: str,
         spec: ApplySpec,
+        excludes: Optional[Dict] = None,
     ) -> ApplySpec:
         try:
             spec_request = request_cls(user=user, project=project, spec=spec)
-            spec_json = self._call_plugin_service(spec_request, endpoint)
+            spec_json = self._call_plugin_service(spec_request, endpoint, excludes)
             response = response_cls(**spec_json)
             self._check_request_rejected(response)
             return response.spec
@@ -88,7 +98,13 @@ class CustomApplyPolicy(ApplyPolicy):
 
     def on_run_apply(self, user: str, project: str, spec: RunSpec) -> RunSpec:
         return self._on_apply(
-            RunSpecRequest, RunSpecResponse, "/apply_policies/on_run_apply", user, project, spec
+            RunSpecRequest,
+            RunSpecResponse,
+            "/apply_policies/on_run_apply",
+            user,
+            project,
+            spec,
+            excludes=get_run_spec_excludes(spec),
         )
 
     def on_fleet_apply(self, user: str, project: str, spec: FleetSpec) -> FleetSpec:
@@ -99,6 +115,7 @@ class CustomApplyPolicy(ApplyPolicy):
             user,
             project,
             spec,
+            excludes=get_fleet_spec_excludes(spec),
         )
 
     def on_volume_apply(self, user: str, project: str, spec: VolumeSpec) -> VolumeSpec:
@@ -109,6 +126,7 @@ class CustomApplyPolicy(ApplyPolicy):
             user,
             project,
             spec,
+            excludes=get_volume_spec_excludes(spec),
         )
 
     def on_gateway_apply(self, user: str, project: str, spec: GatewaySpec) -> GatewaySpec:
@@ -119,6 +137,7 @@ class CustomApplyPolicy(ApplyPolicy):
             user,
             project,
             spec,
+            excludes=get_gateway_spec_excludes(spec),
         )
 
 

dstack/version.py

@@ -1,3 +1,3 @@
-__version__ = "0.19.12rc1"
+__version__ = "0.19.14"
 __is_release__ = True
-base_image = "0.8"
+base_image = "0.10" base_image_ubuntu_version = "22.04"

{dstack-0.19.12rc1.dist-info → dstack-0.19.14.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: dstack
-Version: 0.19.12rc1
+Version: 0.19.14
 Summary: dstack is an open-source orchestration engine for running AI workloads on any cloud or on-premises.
 Project-URL: Homepage, https://dstack.ai
 Project-URL: Source, https://github.com/dstackai/dstack