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

dstack/api/server/__init__.py

@@ -6,7 +6,12 @@ from typing import Dict, List, Optional, Type
 import requests
 
 from dstack import version
-from dstack._internal.core.errors import ClientError, ServerClientError, URLNotFoundError
+from dstack._internal.core.errors import (
+    ClientError,
+    MethodNotAllowedError,
+    ServerClientError,
+    URLNotFoundError,
+)
 from dstack._internal.utils.logging import get_logger
 from dstack.api.server._backends import BackendsAPIClient
 from dstack.api.server._fleets import FleetsAPIClient
@@ -156,6 +161,8 @@ class APIClient:
                 )
             if resp.status_code == 404:
                 raise URLNotFoundError(f"Status code 404 when requesting {resp.request.url}")
+            if resp.status_code == 405:
+                raise MethodNotAllowedError(f"Status code 405 when requesting {resp.request.url}")
             if 400 <= resp.status_code < 600:
                 raise ClientError(
                     f"Unexpected error: status code {resp.status_code}"

dstack/api/server/_fleets.py

@@ -3,6 +3,7 @@ from typing import Any, Dict, List, Optional, Union
 from pydantic import parse_obj_as
 
 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,
@@ -83,9 +84,24 @@ def _get_apply_plan_excludes(plan_input: ApplyFleetPlanInput) -> Dict:
     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)

dstack/api/server/_runs.py

@@ -7,6 +7,7 @@ from pydantic import parse_obj_as
 from dstack._internal.core.models.configurations import ServiceConfiguration
 from dstack._internal.core.models.runs import (
     ApplyRunPlanInput,
+    JobSubmission,
     Run,
     RunPlan,
     RunSpec,
@@ -96,13 +97,53 @@ 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:
         apply_plan_excludes["run_spec"] = run_spec_excludes
-    if plan.current_resource is not None:
-        apply_plan_excludes["current_resource"] = {
-            "run_spec": _get_run_spec_excludes(plan.current_resource.run_spec)
+    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"}}}
+            }
+        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"}}}
+                }
     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

dstack/plugins/__init__.py

@@ -0,0 +1,8 @@
+# ruff: noqa: F401
+from dstack._internal.core.models.fleets import FleetSpec
+from dstack._internal.core.models.gateways import GatewaySpec
+from dstack._internal.core.models.runs import RunSpec
+from dstack._internal.core.models.volumes import VolumeSpec
+from dstack.plugins._base import ApplyPolicy, Plugin
+from dstack.plugins._models import ApplySpec
+from dstack.plugins._utils import get_plugin_logger

dstack/plugins/_base.py

@@ -0,0 +1,72 @@
+from dstack._internal.core.models.fleets import FleetSpec
+from dstack._internal.core.models.gateways import GatewaySpec
+from dstack._internal.core.models.runs import RunSpec
+from dstack._internal.core.models.volumes import VolumeSpec
+from dstack.plugins._models import ApplySpec
+
+
+class ApplyPolicy:
+    """
+    A base apply policy class to modify specs on `dstack apply`.
+    Subclass it and return the subclass instance in `Plugin.get_apply_policies()`.
+    """
+
+    def on_apply(self, user: str, project: str, spec: ApplySpec) -> ApplySpec:
+        """
+        Modify `spec` before it's applied.
+        Raise `ValueError` for `spec` to be rejected as invalid.
+
+        This method can be called twice:
+          * first when a user gets a plan
+          * second when a user applies a plan
+
+        In both cases, the original spec is passed, so the method does not
+        need to check if it modified the spec before.
+
+        It's safe to modify and return `spec` without copying.
+        """
+        if isinstance(spec, RunSpec):
+            return self.on_run_apply(user=user, project=project, spec=spec)
+        if isinstance(spec, FleetSpec):
+            return self.on_fleet_apply(user=user, project=project, spec=spec)
+        if isinstance(spec, VolumeSpec):
+            return self.on_volume_apply(user=user, project=project, spec=spec)
+        if isinstance(spec, GatewaySpec):
+            return self.on_gateway_apply(user=user, project=project, spec=spec)
+        raise ValueError(f"Unknown spec type {type(spec)}")
+
+    def on_run_apply(self, user: str, project: str, spec: RunSpec) -> RunSpec:
+        """
+        Called by the default `on_apply()` implementation for runs.
+        """
+        return spec
+
+    def on_fleet_apply(self, user: str, project: str, spec: FleetSpec) -> FleetSpec:
+        """
+        Called by the default `on_apply()` implementation for fleets.
+        """
+        return spec
+
+    def on_volume_apply(self, user: str, project: str, spec: VolumeSpec) -> VolumeSpec:
+        """
+        Called by the default `on_apply()` implementation for volumes.
+        """
+        return spec
+
+    def on_gateway_apply(self, user: str, project: str, spec: GatewaySpec) -> GatewaySpec:
+        """
+        Called by the default `on_apply()` implementation for gateways.
+        """
+        return spec
+
+
+class Plugin:
+    """
+    A base plugin class.
+    Plugins must subclass it, implement public methods,
+    and register the subclass as an entrypoint of the package
+    (https://packaging.python.org/en/latest/specifications/entry-points/).
+    """
+
+    def get_apply_policies(self) -> list[ApplyPolicy]:
+        return []

dstack/plugins/_models.py

@@ -0,0 +1,8 @@
+from typing import TypeVar
+
+from dstack._internal.core.models.fleets import FleetSpec
+from dstack._internal.core.models.gateways import GatewaySpec
+from dstack._internal.core.models.runs import RunSpec
+from dstack._internal.core.models.volumes import VolumeSpec
+
+ApplySpec = TypeVar("ApplySpec", RunSpec, FleetSpec, VolumeSpec, GatewaySpec)

dstack/plugins/_utils.py

@@ -0,0 +1,19 @@
+import logging
+
+from dstack._internal.utils.logging import get_logger
+
+
+def get_plugin_logger(name: str) -> logging.Logger:
+    """
+    Use this function to set up loggers in plugins.
+
+    Put at the top of the plugin modules:
+
+    ```
+    from dstack.plugins import get_plugin_logger
+
+    logger = get_plugin_logger(__name__)
+    ```
+
+    """
+    return get_logger(f"dstack.plugins.{name}")

dstack/version.py

@@ -1,3 +1,3 @@
-__version__ = "0.19.6rc1"
+__version__ = "0.19.8"
 __is_release__ = True
 base_image = "0.7"

{dstack-0.19.6rc1.dist-info → dstack-0.19.8.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: dstack
-Version: 0.19.6rc1
+Version: 0.19.8
 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
@@ -21,7 +21,7 @@ Requires-Dist: cryptography
 Requires-Dist: cursor
 Requires-Dist: filelock
 Requires-Dist: gitpython
-Requires-Dist: gpuhunt<0.2.0,>=0.1.3
+Requires-Dist: gpuhunt==0.1.5
 Requires-Dist: jsonschema
 Requires-Dist: packaging
 Requires-Dist: paramiko>=3.2.0
@@ -49,9 +49,11 @@ Requires-Dist: asyncpg; extra == 'all'
 Requires-Dist: azure-identity>=1.12.0; extra == 'all'
 Requires-Dist: azure-mgmt-authorization>=3.0.0; extra == 'all'
 Requires-Dist: azure-mgmt-compute>=29.1.0; extra == 'all'
+Requires-Dist: azure-mgmt-msi>=7.0.0; extra == 'all'
 Requires-Dist: azure-mgmt-network<28.0.0,>=23.0.0; extra == 'all'
 Requires-Dist: azure-mgmt-resource>=22.0.0; extra == 'all'
 Requires-Dist: azure-mgmt-subscription>=3.1.1; extra == 'all'
+Requires-Dist: backports-entry-points-selectable; extra == 'all'
 Requires-Dist: boto3; extra == 'all'
 Requires-Dist: botocore; extra == 'all'
 Requires-Dist: datacrunch; extra == 'all'
@@ -87,6 +89,7 @@ Requires-Dist: alembic-postgresql-enum; extra == 'aws'
 Requires-Dist: alembic>=1.10.2; extra == 'aws'
 Requires-Dist: apscheduler<4; extra == 'aws'
 Requires-Dist: asyncpg; extra == 'aws'
+Requires-Dist: backports-entry-points-selectable; extra == 'aws'
 Requires-Dist: boto3; extra == 'aws'
 Requires-Dist: botocore; extra == 'aws'
 Requires-Dist: docker>=6.0.0; extra == 'aws'
@@ -114,9 +117,11 @@ Requires-Dist: asyncpg; extra == 'azure'
 Requires-Dist: azure-identity>=1.12.0; extra == 'azure'
 Requires-Dist: azure-mgmt-authorization>=3.0.0; extra == 'azure'
 Requires-Dist: azure-mgmt-compute>=29.1.0; extra == 'azure'
+Requires-Dist: azure-mgmt-msi>=7.0.0; extra == 'azure'
 Requires-Dist: azure-mgmt-network<28.0.0,>=23.0.0; extra == 'azure'
 Requires-Dist: azure-mgmt-resource>=22.0.0; extra == 'azure'
 Requires-Dist: azure-mgmt-subscription>=3.1.1; extra == 'azure'
+Requires-Dist: backports-entry-points-selectable; extra == 'azure'
 Requires-Dist: docker>=6.0.0; extra == 'azure'
 Requires-Dist: fastapi; extra == 'azure'
 Requires-Dist: grpcio>=1.50; extra == 'azure'
@@ -139,6 +144,7 @@ Requires-Dist: alembic-postgresql-enum; extra == 'datacrunch'
 Requires-Dist: alembic>=1.10.2; extra == 'datacrunch'
 Requires-Dist: apscheduler<4; extra == 'datacrunch'
 Requires-Dist: asyncpg; extra == 'datacrunch'
+Requires-Dist: backports-entry-points-selectable; extra == 'datacrunch'
 Requires-Dist: datacrunch; extra == 'datacrunch'
 Requires-Dist: docker>=6.0.0; extra == 'datacrunch'
 Requires-Dist: fastapi; extra == 'datacrunch'
@@ -170,6 +176,7 @@ Requires-Dist: alembic-postgresql-enum; extra == 'gcp'
 Requires-Dist: alembic>=1.10.2; extra == 'gcp'
 Requires-Dist: apscheduler<4; extra == 'gcp'
 Requires-Dist: asyncpg; extra == 'gcp'
+Requires-Dist: backports-entry-points-selectable; extra == 'gcp'
 Requires-Dist: docker>=6.0.0; extra == 'gcp'
 Requires-Dist: fastapi; extra == 'gcp'
 Requires-Dist: google-api-python-client>=2.80.0; extra == 'gcp'
@@ -199,6 +206,7 @@ Requires-Dist: alembic-postgresql-enum; extra == 'kubernetes'
 Requires-Dist: alembic>=1.10.2; extra == 'kubernetes'
 Requires-Dist: apscheduler<4; extra == 'kubernetes'
 Requires-Dist: asyncpg; extra == 'kubernetes'
+Requires-Dist: backports-entry-points-selectable; extra == 'kubernetes'
 Requires-Dist: docker>=6.0.0; extra == 'kubernetes'
 Requires-Dist: fastapi; extra == 'kubernetes'
 Requires-Dist: grpcio>=1.50; extra == 'kubernetes'
@@ -222,6 +230,7 @@ Requires-Dist: alembic-postgresql-enum; extra == 'lambda'
 Requires-Dist: alembic>=1.10.2; extra == 'lambda'
 Requires-Dist: apscheduler<4; extra == 'lambda'
 Requires-Dist: asyncpg; extra == 'lambda'
+Requires-Dist: backports-entry-points-selectable; extra == 'lambda'
 Requires-Dist: boto3; extra == 'lambda'
 Requires-Dist: botocore; extra == 'lambda'
 Requires-Dist: docker>=6.0.0; extra == 'lambda'
@@ -246,6 +255,7 @@ Requires-Dist: alembic-postgresql-enum; extra == 'nebius'
 Requires-Dist: alembic>=1.10.2; extra == 'nebius'
 Requires-Dist: apscheduler<4; extra == 'nebius'
 Requires-Dist: asyncpg; extra == 'nebius'
+Requires-Dist: backports-entry-points-selectable; extra == 'nebius'
 Requires-Dist: docker>=6.0.0; extra == 'nebius'
 Requires-Dist: fastapi; extra == 'nebius'
 Requires-Dist: grpcio>=1.50; extra == 'nebius'
@@ -269,6 +279,7 @@ Requires-Dist: alembic-postgresql-enum; extra == 'oci'
 Requires-Dist: alembic>=1.10.2; extra == 'oci'
 Requires-Dist: apscheduler<4; extra == 'oci'
 Requires-Dist: asyncpg; extra == 'oci'
+Requires-Dist: backports-entry-points-selectable; extra == 'oci'
 Requires-Dist: docker>=6.0.0; extra == 'oci'
 Requires-Dist: fastapi; extra == 'oci'
 Requires-Dist: grpcio>=1.50; extra == 'oci'
@@ -292,6 +303,7 @@ Requires-Dist: alembic-postgresql-enum; extra == 'server'
 Requires-Dist: alembic>=1.10.2; extra == 'server'
 Requires-Dist: apscheduler<4; extra == 'server'
 Requires-Dist: asyncpg; extra == 'server'
+Requires-Dist: backports-entry-points-selectable; extra == 'server'
 Requires-Dist: docker>=6.0.0; extra == 'server'
 Requires-Dist: fastapi; extra == 'server'
 Requires-Dist: grpcio>=1.50; extra == 'server'