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

dstack/_internal/server/models.py

@@ -5,7 +5,6 @@ from typing import Callable, List, Optional, Union
 from sqlalchemy import (
     BigInteger,
     Boolean,
-    Column,
     DateTime,
     Enum,
     Float,
@@ -15,7 +14,6 @@ from sqlalchemy import (
     LargeBinary,
     MetaData,
     String,
-    Table,
     Text,
     TypeDecorator,
     UniqueConstraint,
@@ -351,13 +349,17 @@ class JobModel(BaseModel):
     job_spec_data: Mapped[str] = mapped_column(Text)
     job_provisioning_data: Mapped[Optional[str]] = mapped_column(Text)
     runner_timestamp: Mapped[Optional[int]] = mapped_column(BigInteger)
+    inactivity_secs: Mapped[Optional[int]] = mapped_column(Integer)  # 0 - active, None - N/A
     # `removed` is used to ensure that the instance is killed after the job is finished
     remove_at: Mapped[Optional[datetime]] = mapped_column(NaiveDateTime)
     volumes_detached_at: Mapped[Optional[datetime]] = mapped_column(NaiveDateTime)
     # `instance_assigned` means instance assignment was done.
     # if `instance_assigned` is True and `instance` is None, no instance was assiged.
     instance_assigned: Mapped[bool] = mapped_column(Boolean, default=False)
-    instance: Mapped[Optional["InstanceModel"]] = relationship(back_populates="job")
+    instance_id: Mapped[Optional[uuid.UUID]] = mapped_column(
+        ForeignKey("instances.id", ondelete="CASCADE")
+    )
+    instance: Mapped[Optional["InstanceModel"]] = relationship(back_populates="jobs")
     used_instance_id: Mapped[Optional[uuid.UUID]] = mapped_column(UUIDType(binary=False))
     replica_num: Mapped[int] = mapped_column(Integer)
     job_runtime_data: Mapped[Optional[str]] = mapped_column(Text)
@@ -543,15 +545,19 @@ class InstanceModel(BaseModel):
 
     remote_connection_info: Mapped[Optional[str]] = mapped_column(Text)
 
-    # current job
-    job_id: Mapped[Optional[uuid.UUID]] = mapped_column(ForeignKey("jobs.id"))
-    job: Mapped[Optional["JobModel"]] = relationship(back_populates="instance", lazy="joined")
+    # NULL means `auto` (only during provisioning, when ready it's not NULL)
+    total_blocks: Mapped[Optional[int]] = mapped_column(Integer)
+    busy_blocks: Mapped[int] = mapped_column(Integer, default=0)
+
+    jobs: Mapped[list["JobModel"]] = relationship(back_populates="instance", lazy="joined")
     last_job_processed_at: Mapped[Optional[datetime]] = mapped_column(NaiveDateTime)
 
-    # volumes attached to the instance
-    volumes: Mapped[List["VolumeModel"]] = relationship(
-        secondary="volumes_attachments",
-        back_populates="instances",
+    volume_attachments: Mapped[List["VolumeAttachmentModel"]] = relationship(
+        back_populates="instance",
+        # Add delete-orphan option so that removing entries from volume_attachments
+        # automatically marks them for deletion.
+        # SQLalchemy requires delete when using delete-orphan.
+        cascade="save-update, merge, delete-orphan, delete",
     )
 
 
@@ -581,23 +587,21 @@ class VolumeModel(BaseModel):
 
     configuration: Mapped[str] = mapped_column(Text)
     volume_provisioning_data: Mapped[Optional[str]] = mapped_column(Text)
-    # FIXME: volume_attachment_data should be in "volumes_attachments"
-    # to support multi-attach volumes
+
+    attachments: Mapped[List["VolumeAttachmentModel"]] = relationship(back_populates="volume")
+
+    # Deprecated in favor of VolumeAttachmentModel.attachment_data
     volume_attachment_data: Mapped[Optional[str]] = mapped_column(Text)
 
-    # instances the volume is attached to
-    instances: Mapped[List["InstanceModel"]] = relationship(
-        secondary="volumes_attachments",
-        back_populates="volumes",
-    )
 
+class VolumeAttachmentModel(BaseModel):
+    __tablename__ = "volumes_attachments"
 
-volumes_attachments_table = Table(
-    "volumes_attachments",
-    BackendModel.metadata,
-    Column("volume_id", ForeignKey("volumes.id"), primary_key=True),
-    Column("instace_id", ForeignKey("instances.id"), primary_key=True),
-)
+    volume_id: Mapped[uuid.UUID] = mapped_column(ForeignKey("volumes.id"), primary_key=True)
+    volume: Mapped[VolumeModel] = relationship(back_populates="attachments")
+    instance_id: Mapped[uuid.UUID] = mapped_column(ForeignKey("instances.id"), primary_key=True)
+    instance: Mapped[InstanceModel] = relationship(back_populates="volume_attachments")
+    attachment_data: Mapped[Optional[str]] = mapped_column(Text)
 
 
 class PlacementGroupModel(BaseModel):

dstack/_internal/server/routers/runs.py

@@ -47,6 +47,7 @@ async def list_runs(
     """
     Returns all runs visible to user sorted by descending `submitted_at`.
     `project_name`, `repo_id`, `username`, and `only_active` can be specified as filters.
+    Setting `only_active` to `true` excludes finished runs and deleted runs.
     Specifying `repo_id` without `project_name` returns no runs.
 
     The results are paginated. To get the next page, pass `submitted_at` and `id` of

dstack/_internal/server/schemas/runner.py

@@ -34,6 +34,7 @@ class PullResponse(CoreModel):
     job_logs: List[LogEvent]
     runner_logs: List[LogEvent]
     last_updated: int
+    no_connections_secs: Optional[int] = None  # Optional for compatibility with old runners
 
 
 class SubmitBody(CoreModel):

dstack/_internal/server/services/backends/configurators/azure.py

@@ -2,6 +2,7 @@ import json
 from concurrent.futures import ThreadPoolExecutor, as_completed
 from typing import List, Optional, Tuple
 
+import azure.core.exceptions
 from azure.core.credentials import TokenCredential
 from azure.mgmt import network as network_mgmt
 from azure.mgmt import resource as resource_mgmt
@@ -154,16 +155,17 @@ class AzureConfigurator(Configurator):
         if is_core_model_instance(config.creds, AzureClientCreds):
             self._set_client_creds_tenant_id(config.creds, config.tenant_id)
         credential, _ = auth.authenticate(config.creds)
-        resource_group = self._create_resource_group(
-            credential=credential,
-            subscription_id=config.subscription_id,
-            location=MAIN_LOCATION,
-            project_name=project.name,
-        )
+        if config.resource_group is None:
+            config.resource_group = self._create_resource_group(
+                credential=credential,
+                subscription_id=config.subscription_id,
+                location=MAIN_LOCATION,
+                project_name=project.name,
+            )
         self._create_network_resources(
             credential=credential,
             subscription_id=config.subscription_id,
-            resource_group=resource_group,
+            resource_group=config.resource_group,
             locations=config.locations,
             create_default_network=config.vpc_ids is None,
         )
@@ -172,7 +174,6 @@ class AzureConfigurator(Configurator):
             type=self.TYPE.value,
             config=AzureStoredConfig(
                 **AzureConfigInfo.__response__.parse_obj(config).dict(),
-                resource_group=resource_group,
             ).json(),
             auth=DecryptedString(plaintext=AzureCreds.parse_obj(config.creds).__root__.json()),
         )
@@ -322,6 +323,7 @@ class AzureConfigurator(Configurator):
         self, config: AzureConfigInfoWithCredsPartial, credential: auth.AzureCredential
     ):
         self._check_tags_config(config)
+        self._check_resource_group(config=config, credential=credential)
         self._check_vpc_config(config=config, credential=credential)
 
     def _check_tags_config(self, config: AzureConfigInfoWithCredsPartial):
@@ -336,6 +338,18 @@ class AzureConfigurator(Configurator):
         except BackendError as e:
             raise ServerClientError(e.args[0])
 
+    def _check_resource_group(
+        self, config: AzureConfigInfoWithCredsPartial, credential: auth.AzureCredential
+    ):
+        if config.resource_group is None:
+            return
+        resource_manager = ResourceManager(
+            credential=credential,
+            subscription_id=config.subscription_id,
+        )
+        if not resource_manager.resource_group_exists(config.resource_group):
+            raise ServerClientError(f"Resource group {config.resource_group} not found")
+
     def _check_vpc_config(
         self, config: AzureConfigInfoWithCredsPartial, credential: auth.AzureCredential
     ):
@@ -406,6 +420,18 @@ class ResourceManager:
         )
         return resource_group.name
 
+    def resource_group_exists(
+        self,
+        name: str,
+    ) -> bool:
+        try:
+            self.resource_client.resource_groups.get(
+                resource_group_name=name,
+            )
+        except azure.core.exceptions.ResourceNotFoundError:
+            return False
+        return True
+
 
 class NetworkManager:
     def __init__(self, credential: TokenCredential, subscription_id: str):

dstack/_internal/server/services/config.py

@@ -124,6 +124,15 @@ class AzureConfig(CoreModel):
     type: Annotated[Literal["azure"], Field(description="The type of the backend")] = "azure"
     tenant_id: Annotated[str, Field(description="The tenant ID")]
     subscription_id: Annotated[str, Field(description="The subscription ID")]
+    resource_group: Annotated[
+        Optional[str],
+        Field(
+            description=(
+                "The resource group for resources created by `dstack`."
+                " If not specified, `dstack` will create a new resource group"
+            )
+        ),
+    ] = None
     regions: Annotated[
         Optional[List[str]],
         Field(description="The list of Azure regions (locations). Omit to use all regions"),

dstack/_internal/server/services/fleets.py

@@ -2,7 +2,7 @@ import random
 import string
 import uuid
 from datetime import datetime, timezone
-from typing import List, Optional, Tuple, Union, cast
+from typing import List, Literal, Optional, Tuple, Union, cast
 
 from sqlalchemy import and_, func, or_, select
 from sqlalchemy.ext.asyncio import AsyncSession
@@ -31,6 +31,7 @@ from dstack._internal.core.models.instances import (
     InstanceOfferWithAvailability,
     InstanceStatus,
     RemoteConnectionInfo,
+    SSHConnectionParams,
     SSHKey,
 )
 from dstack._internal.core.models.pools import Instance
@@ -256,6 +257,8 @@ async def get_plan(
             project=project,
             profile=spec.merged_profile,
             requirements=_get_fleet_requirements(spec),
+            fleet_spec=spec,
+            blocks=spec.configuration.blocks,
         )
         offers = [offer for _, offer in offers_with_backends]
     _remove_fleet_spec_sensitive_info(spec)
@@ -275,11 +278,15 @@ async def get_create_instance_offers(
     project: ProjectModel,
     profile: Profile,
     requirements: Requirements,
-    exclude_not_available=False,
+    fleet_spec: Optional[FleetSpec] = None,
     fleet_model: Optional[FleetModel] = None,
+    blocks: Union[int, Literal["auto"]] = 1,
+    exclude_not_available: bool = False,
 ) -> List[Tuple[Backend, InstanceOfferWithAvailability]]:
     multinode = False
     master_job_provisioning_data = None
+    if fleet_spec is not None:
+        multinode = fleet_spec.configuration.placement == InstanceGroupPlacement.CLUSTER
     if fleet_model is not None:
         fleet = fleet_model_to_fleet(fleet_model)
         multinode = fleet.spec.configuration.placement == InstanceGroupPlacement.CLUSTER
@@ -296,6 +303,7 @@ async def get_create_instance_offers(
         exclude_not_available=exclude_not_available,
         multinode=multinode,
         master_job_provisioning_data=master_job_provisioning_data,
+        blocks=blocks,
     )
     offers = [
         (backend, offer)
@@ -406,6 +414,7 @@ async def create_fleet_instance_model(
         instance_num=instance_num,
         placement_group_name=placement_group_name,
         reservation=reservation,
+        blocks=spec.configuration.blocks,
     )
     return instance_model
 
@@ -424,18 +433,33 @@ async def create_fleet_ssh_instance_model(
         ssh_user = ssh_params.user
         ssh_key = ssh_params.ssh_key
         port = ssh_params.port
+        proxy_jump = ssh_params.proxy_jump
         internal_ip = None
+        blocks = 1
     else:
         hostname = host.hostname
         ssh_user = host.user or ssh_params.user
         ssh_key = host.ssh_key or ssh_params.ssh_key
         port = host.port or ssh_params.port
+        proxy_jump = host.proxy_jump or ssh_params.proxy_jump
         internal_ip = host.internal_ip
+        blocks = host.blocks
 
     if ssh_user is None or ssh_key is None:
         # This should not be reachable but checked by fleet spec validation
         raise ServerClientError("ssh key or user not specified")
 
+    if proxy_jump is not None:
+        ssh_proxy = SSHConnectionParams(
+            hostname=proxy_jump.hostname,
+            port=proxy_jump.port or 22,
+            username=proxy_jump.user,
+        )
+        ssh_proxy_keys = [proxy_jump.ssh_key]
+    else:
+        ssh_proxy = None
+        ssh_proxy_keys = None
+
     instance_model = await pools_services.create_ssh_instance_model(
         project=project,
         pool=pool,
@@ -445,10 +469,13 @@ async def create_fleet_ssh_instance_model(
         host=hostname,
         ssh_user=ssh_user,
         ssh_keys=[ssh_key],
+        ssh_proxy=ssh_proxy,
+        ssh_proxy_keys=ssh_proxy_keys,
         env=env,
         internal_ip=internal_ip,
         instance_network=ssh_params.network,
         port=port or 22,
+        blocks=blocks,
     )
     return instance_model
 
@@ -544,7 +571,7 @@ async def generate_fleet_name(session: AsyncSession, project: ProjectModel) -> s
 
 
 def is_fleet_in_use(fleet_model: FleetModel, instance_nums: Optional[List[int]] = None) -> bool:
-    instances_in_use = [i for i in fleet_model.instances if i.job_id is not None and not i.deleted]
+    instances_in_use = [i for i in fleet_model.instances if i.jobs and not i.deleted]
     selected_instance_in_use = instances_in_use
     if instance_nums is not None:
         selected_instance_in_use = [i for i in instances_in_use if i.instance_num in instance_nums]
@@ -606,6 +633,8 @@ async def create_instance(
         instance_configuration=None,
         termination_policy=termination_policy,
         termination_idle_time=termination_idle_time,
+        total_blocks=1,
+        busy_blocks=0,
     )
     logger.info(
         "Added a new instance %s",

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

@@ -74,10 +74,18 @@ class GatewayClient:
         resp.raise_for_status()
         self.is_server_ready = True
 
-    async def register_replica(self, run: Run, job_submission: JobSubmission):
+    async def register_replica(
+        self,
+        run: Run,
+        job_submission: JobSubmission,
+        ssh_head_proxy: Optional[SSHConnectionParams],
+        ssh_head_proxy_private_key: Optional[str],
+    ):
         payload = {
             "job_id": job_submission.id.hex,
             "app_port": run.run_spec.configuration.port.container_port,
+            "ssh_head_proxy": ssh_head_proxy.dict() if ssh_head_proxy is not None else None,
+            "ssh_head_proxy_private_key": ssh_head_proxy_private_key,
         }
         jpd = job_submission.job_provisioning_data
         if not jpd.dockerized: