dstack 0.18.41__py3-none-any.whl → 0.18.42__py3-none-any.whl

dstack/_internal/cli/utils/volume.py

@@ -22,6 +22,8 @@ def get_volumes_table(
     if verbose:
         table.add_column("REGION")
     table.add_column("STATUS")
+    if verbose:
+        table.add_column("ATTACHED")
     table.add_column("CREATED")
     if verbose:
         table.add_column("ERROR")
@@ -37,11 +39,18 @@ def get_volumes_table(
                 and volume.provisioning_data.availability_zone is not None
             ):
                 region += f" ({volume.provisioning_data.availability_zone})"
+        attached = "-"
+        if volume.attachments is not None:
+            attached = ", ".join(
+                {va.instance.fleet_name for va in volume.attachments if va.instance.fleet_name}
+            )
+            attached = attached or "-"
         row = {
             "NAME": volume.name,
             "BACKEND": backend,
             "REGION": region,
             "STATUS": volume.status,
+            "ATTACHED": attached,
             "CREATED": format_date(volume.created_at),
             "ERROR": volume.status_message,
         }

dstack/_internal/core/backends/aws/compute.py

@@ -635,11 +635,12 @@ class AWSCompute(Compute):
         ec2_client = self.session.client("ec2", region_name=volume.configuration.region)
 
         logger.debug("Detaching EBS volume %s from instance %s", volume.volume_id, instance_id)
+        attachment_data = get_or_error(volume.get_attachment_data_for_instance(instance_id))
         try:
             ec2_client.detach_volume(
                 VolumeId=volume.volume_id,
                 InstanceId=instance_id,
-                Device=get_or_error(volume.attachment_data).device_name,
+                Device=attachment_data.device_name,
                 Force=force,
             )
         except botocore.exceptions.ClientError as e:

dstack/_internal/core/backends/gcp/compute.py

@@ -666,6 +666,7 @@ class GCPCompute(Compute):
             instance_id,
         )
         zone = get_or_error(volume.provisioning_data).availability_zone
+        attachment_data = get_or_error(volume.get_attachment_data_for_instance(instance_id))
         # This method has no information if the instance is a TPU or a VM,
         # so we first try to see if there is a TPU with such name
         try:
@@ -694,7 +695,7 @@ class GCPCompute(Compute):
                 project=self.config.project_id,
                 zone=get_or_error(volume.provisioning_data).availability_zone,
                 instance=instance_id,
-                device_name=get_or_error(volume.attachment_data).device_name,
+                device_name=attachment_data.device_name,
             )
             gcp_resources.wait_for_extended_operation(operation, "persistent disk detachment")
         logger.debug(

dstack/_internal/core/models/runs.py

@@ -150,9 +150,9 @@ class JobTerminationReason(str, Enum):
 class Requirements(CoreModel):
     # TODO: Make requirements' fields required
     resources: ResourcesSpec
-    max_price: Optional[float]
-    spot: Optional[bool]
-    reservation: Optional[str]
+    max_price: Optional[float] = None
+    spot: Optional[bool] = None
+    reservation: Optional[str] = None
 
     def pretty_format(self, resources_only: bool = False):
         res = self.resources.pretty_format()

dstack/_internal/core/models/volumes.py

@@ -71,6 +71,18 @@ class VolumeAttachmentData(CoreModel):
     device_name: Optional[str] = None
 
 
+class VolumeInstance(CoreModel):
+    name: str
+    fleet_name: Optional[str] = None
+    instance_num: int
+    instance_id: Optional[str] = None
+
+
+class VolumeAttachment(CoreModel):
+    instance: VolumeInstance
+    attachment_data: Optional[VolumeAttachmentData] = None
+
+
 class Volume(CoreModel):
     id: uuid.UUID
     name: str
@@ -86,8 +98,19 @@ class Volume(CoreModel):
     deleted: bool
     volume_id: Optional[str] = None  # id of the volume in the cloud
     provisioning_data: Optional[VolumeProvisioningData] = None
+    attachments: Optional[List[VolumeAttachment]] = None
+    # attachment_data is deprecated in favor of attachments.
+    # It's only set for volumes that were attached before attachments.
     attachment_data: Optional[VolumeAttachmentData] = None
 
+    def get_attachment_data_for_instance(self, instance_id: str) -> Optional[VolumeAttachmentData]:
+        if self.attachments is not None:
+            for attachment in self.attachments:
+                if attachment.instance.instance_id == instance_id:
+                    return attachment.attachment_data
+        # volume was attached before attachments were introduced
+        return self.attachment_data
+
 
 class VolumePlan(CoreModel):
     project_name: str

dstack/_internal/server/background/tasks/process_instances.py

@@ -507,9 +507,9 @@ async def _create_instance(session: AsyncSession, instance: InstanceModel) -> No
         project=instance.project,
         profile=profile,
         requirements=requirements,
-        exclude_not_available=True,
         fleet_model=instance.fleet,
         blocks="auto" if instance.total_blocks is None else instance.total_blocks,
+        exclude_not_available=True,
     )
 
     if not offers and should_retry:
@@ -915,9 +915,8 @@ def _get_instance_offer_for_instance(
         instance_offer.availability_zones = [
             z
             for z in instance_offer.availability_zones
-            if instance_offer.availability_zones == master_job_provisioning_data.availability_zone
+            if z == master_job_provisioning_data.availability_zone
         ]
-
     return instance_offer
 
 

dstack/_internal/server/background/tasks/process_running_jobs.py

@@ -205,6 +205,7 @@ async def _process_running_job(session: AsyncSession, job_model: JobModel):
                     None,
                     run,
                     job_model,
+                    job_provisioning_data,
                     volumes,
                     secrets,
                     job.job_spec.registry_auth,
@@ -376,6 +377,7 @@ def _process_provisioning_with_shim(
     ports: Dict[int, int],
     run: Run,
     job_model: JobModel,
+    job_provisioning_data: JobProvisioningData,
     volumes: List[Volume],
     secrets: Dict[str, str],
     registry_auth: Optional[RegistryAuth],
@@ -459,6 +461,7 @@ def _process_provisioning_with_shim(
             host_ssh_user=ssh_user,
             host_ssh_keys=[ssh_key] if ssh_key else [],
             container_ssh_keys=public_keys,
+            instance_id=job_provisioning_data.instance_id,
         )
     else:
         submitted = shim_client.submit(
@@ -475,6 +478,7 @@ def _process_provisioning_with_shim(
             mounts=volume_mounts,
             volumes=volumes,
             instance_mounts=instance_mounts,
+            instance_id=job_provisioning_data.instance_id,
         )
         if not submitted:
             # This can happen when we lost connection to the runner (e.g., network issues), marked

dstack/_internal/server/background/tasks/process_submitted_jobs.py

@@ -43,6 +43,7 @@ from dstack._internal.server.models import (
     PoolModel,
     ProjectModel,
     RunModel,
+    VolumeAttachmentModel,
     VolumeModel,
 )
 from dstack._internal.server.services.backends import get_project_backend_by_type_or_error
@@ -236,7 +237,7 @@ async def _process_submitted_job(session: AsyncSession, job_model: JobModel):
         res = await session.execute(
             select(InstanceModel)
             .where(InstanceModel.id == job_model.instance.id)
-            .options(selectinload(InstanceModel.volumes))
+            .options(selectinload(InstanceModel.volume_attachments))
             .execution_options(populate_existing=True)
         )
         instance = res.unique().scalar_one()
@@ -390,11 +391,11 @@ async def _assign_job_to_pool_instance(
 
     instances_with_offers.sort(key=lambda instance_with_offer: instance_with_offer[0].price or 0)
     instance, offer = instances_with_offers[0]
-    # Reload InstanceModel with volumes
+    # Reload InstanceModel with volume attachments
     res = await session.execute(
         select(InstanceModel)
         .where(InstanceModel.id == instance.id)
-        .options(joinedload(InstanceModel.volumes))
+        .options(joinedload(InstanceModel.volume_attachments))
     )
     instance = res.unique().scalar_one()
     instance.status = InstanceStatus.BUSY
@@ -580,7 +581,7 @@ def _create_instance_model_for_job(
         backend=offer.backend,
         price=offer.price,
         region=offer.region,
-        volumes=[],
+        volume_attachments=[],
         total_blocks=1,
         busy_blocks=1,
     )
@@ -696,14 +697,18 @@ async def _attach_volume(
     instance: InstanceModel,
     instance_id: str,
 ):
+    volume = volume_model_to_volume(volume_model)
+    # Refresh only to check if the volume wasn't deleted before the lock
     await session.refresh(volume_model)
     if volume_model.deleted:
         raise ServerClientError("Cannot attach a deleted volume")
-    volume = volume_model_to_volume(volume_model)
     attachment_data = await common_utils.run_async(
         backend.compute().attach_volume,
         volume=volume,
         instance_id=instance_id,
     )
-    volume_model.volume_attachment_data = attachment_data.json()
-    instance.volumes.append(volume_model)
+    volume_attachment_model = VolumeAttachmentModel(
+        volume=volume_model,
+        attachment_data=attachment_data.json(),
+    )
+    instance.volume_attachments.append(volume_attachment_model)

dstack/_internal/server/background/tasks/process_terminating_jobs.py

@@ -6,7 +6,13 @@ from sqlalchemy.orm import joinedload, lazyload
 
 from dstack._internal.core.models.runs import JobStatus
 from dstack._internal.server.db import get_session_ctx
-from dstack._internal.server.models import InstanceModel, JobModel, ProjectModel, VolumeModel
+from dstack._internal.server.models import (
+    InstanceModel,
+    JobModel,
+    ProjectModel,
+    VolumeAttachmentModel,
+    VolumeModel,
+)
 from dstack._internal.server.services.jobs import (
     process_terminating_job,
     process_volumes_detaching,
@@ -80,7 +86,12 @@ async def _process_job(session: AsyncSession, job_model: JobModel):
         .where(InstanceModel.id == job_model.used_instance_id)
         .options(
             joinedload(InstanceModel.project).joinedload(ProjectModel.backends),
-            joinedload(InstanceModel.volumes).joinedload(VolumeModel.user),
+            joinedload(InstanceModel.volume_attachments)
+            .joinedload(VolumeAttachmentModel.volume)
+            .joinedload(VolumeModel.user),
+            joinedload(InstanceModel.volume_attachments)
+            .joinedload(VolumeAttachmentModel.volume)
+            .joinedload(VolumeModel.attachments),
         )
     )
     instance_model = res.unique().scalar()

dstack/_internal/server/background/tasks/process_volumes.py

@@ -5,7 +5,12 @@ from sqlalchemy.orm import joinedload
 from dstack._internal.core.errors import BackendError, BackendNotAvailable
 from dstack._internal.core.models.volumes import VolumeStatus
 from dstack._internal.server.db import get_session_ctx
-from dstack._internal.server.models import ProjectModel, VolumeModel
+from dstack._internal.server.models import (
+    InstanceModel,
+    ProjectModel,
+    VolumeAttachmentModel,
+    VolumeModel,
+)
 from dstack._internal.server.services import backends as backends_services
 from dstack._internal.server.services import volumes as volumes_services
 from dstack._internal.server.services.locking import get_locker
@@ -49,6 +54,11 @@ async def _process_submitted_volume(session: AsyncSession, volume_model: VolumeM
         .where(VolumeModel.id == volume_model.id)
         .options(joinedload(VolumeModel.project).joinedload(ProjectModel.backends))
         .options(joinedload(VolumeModel.user))
+        .options(
+            joinedload(VolumeModel.attachments)
+            .joinedload(VolumeAttachmentModel.instance)
+            .joinedload(InstanceModel.fleet)
+        )
         .execution_options(populate_existing=True)
     )
     volume_model = res.unique().scalar_one()

dstack/_internal/server/migrations/versions/a751ef183f27_move_attachment_data_to_volumes_.py

@@ -0,0 +1,34 @@
+"""Move attachment_data to volumes_attachments
+
+Revision ID: a751ef183f27
+Revises: 1e76fb0dde87
+Create Date: 2025-02-12 13:19:57.569591
+
+"""
+
+import sqlalchemy as sa
+from alembic import op
+
+# revision identifiers, used by Alembic.
+revision = "a751ef183f27"
+down_revision = "1e76fb0dde87"
+branch_labels = None
+depends_on = None
+
+
+def upgrade() -> None:
+    # ### commands auto generated by Alembic - please adjust! ###
+    with op.batch_alter_table("volumes_attachments", schema=None) as batch_op:
+        batch_op.alter_column("instace_id", new_column_name="instance_id")
+        batch_op.add_column(sa.Column("attachment_data", sa.Text(), nullable=True))
+
+    # ### end Alembic commands ###
+
+
+def downgrade() -> None:
+    # ### commands auto generated by Alembic - please adjust! ###
+    with op.batch_alter_table("volumes_attachments", schema=None) as batch_op:
+        batch_op.drop_column("attachment_data")
+        batch_op.alter_column("instance_id", new_column_name="instace_id")
+
+    # ### end Alembic commands ###

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,
@@ -554,10 +552,12 @@ class InstanceModel(BaseModel):
     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",
     )
 
 
@@ -587,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/services/fleets.py

@@ -257,6 +257,7 @@ 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]
@@ -277,12 +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

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

@@ -242,7 +242,7 @@ async def process_terminating_job(
                 session=session, project=instance_model.project, names=jrd.volume_names
             )
         else:
-            volume_models = instance_model.volumes
+            volume_models = [va.volume for va in instance_model.volume_attachments]
         if len(volume_models) > 0:
             logger.info("Detaching volumes: %s", [v.name for v in volume_models])
             all_volumes_detached = await _detach_volumes_from_job_instance(
@@ -306,7 +306,7 @@ async def process_volumes_detaching(
             session=session, project=instance_model.project, names=jrd.volume_names
         )
     else:
-        volume_models = instance_model.volumes
+        volume_models = [va.volume for va in instance_model.volume_attachments]
     logger.info("Detaching volumes: %s", [v.name for v in volume_models])
     all_volumes_detached = await _detach_volumes_from_job_instance(
         project=instance_model.project,
@@ -439,8 +439,8 @@ async def _detach_volumes_from_job_instance(
     if job_model.volumes_detached_at is None:
         job_model.volumes_detached_at = common.get_current_datetime()
     detached_volumes_ids = {v.id for v in detached_volumes}
-    instance_model.volumes = [
-        v for v in instance_model.volumes if v.id not in detached_volumes_ids
+    instance_model.volume_attachments = [
+        va for va in instance_model.volume_attachments if va.volume_id not in detached_volumes_ids
     ]
     return all_detached
 

dstack/_internal/server/services/offers.py

@@ -50,35 +50,35 @@ async def get_offers_by_requirements(
     if volumes:
         mount_point_volumes = volumes[0]
         volumes_backend_types = [v.configuration.backend for v in mount_point_volumes]
-        if not backend_types:
+        if backend_types is None:
             backend_types = volumes_backend_types
         backend_types = [b for b in backend_types if b in volumes_backend_types]
         volumes_regions = [v.configuration.region for v in mount_point_volumes]
-        if not regions:
+        if regions is None:
             regions = volumes_regions
         regions = [r for r in regions if r in volumes_regions]
 
     if multinode:
-        if not backend_types:
+        if backend_types is None:
             backend_types = BACKENDS_WITH_MULTINODE_SUPPORT
         backend_types = [b for b in backend_types if b in BACKENDS_WITH_MULTINODE_SUPPORT]
 
     if privileged or instance_mounts:
-        if not backend_types:
+        if backend_types is None:
             backend_types = BACKENDS_WITH_CREATE_INSTANCE_SUPPORT
         backend_types = [b for b in backend_types if b in BACKENDS_WITH_CREATE_INSTANCE_SUPPORT]
 
     if profile.reservation is not None:
-        if not backend_types:
+        if backend_types is None:
             backend_types = BACKENDS_WITH_RESERVATION_SUPPORT
         backend_types = [b for b in backend_types if b in BACKENDS_WITH_RESERVATION_SUPPORT]
 
     # For multi-node, restrict backend and region.
     # The default behavior is to provision all nodes in the same backend and region.
     if master_job_provisioning_data is not None:
-        if not backend_types:
+        if backend_types is None:
             backend_types = [master_job_provisioning_data.get_base_backend()]
-        if not regions:
+        if regions is None:
             regions = [master_job_provisioning_data.region]
         backend_types = [
             b for b in backend_types if b == master_job_provisioning_data.get_base_backend()

dstack/_internal/server/services/pools.py

@@ -462,19 +462,19 @@ def filter_pool_instances(
         zones = [z for z in zones if z in volume_zones]
 
     if multinode:
-        if not backend_types:
+        if backend_types is None:
             backend_types = BACKENDS_WITH_MULTINODE_SUPPORT
         backend_types = [b for b in backend_types if b in BACKENDS_WITH_MULTINODE_SUPPORT]
 
     # For multi-node, restrict backend and region.
     # The default behavior is to provision all nodes in the same backend and region.
     if master_job_provisioning_data is not None:
-        if not backend_types:
+        if backend_types is None:
             backend_types = [master_job_provisioning_data.get_base_backend()]
         backend_types = [
             b for b in backend_types if b == master_job_provisioning_data.get_base_backend()
         ]
-        if not regions:
+        if regions is None:
             regions = [master_job_provisioning_data.region]
         regions = [r for r in regions if r == master_job_provisioning_data.region]
 

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

@@ -239,6 +239,7 @@ class ShimClient:
         host_ssh_user: str,
         host_ssh_keys: list[str],
         container_ssh_keys: list[str],
+        instance_id: str,
     ) -> None:
         if not self.is_api_v2_supported():
             raise ShimAPIVersionError()
@@ -255,7 +256,7 @@ class ShimClient:
             memory=_memory_to_bytes(memory),  # None = 0 = "all available"
             shm_size=_memory_to_bytes(shm_size),  # None = 0 = "use default value"
             network_mode=network_mode,
-            volumes=[_volume_to_shim_volume_info(v) for v in volumes],
+            volumes=[_volume_to_shim_volume_info(v, instance_id) for v in volumes],
             volume_mounts=volume_mounts,
             instance_mounts=instance_mounts,
             host_ssh_user=host_ssh_user,
@@ -303,6 +304,7 @@ class ShimClient:
         mounts: List[VolumeMountPoint],
         volumes: List[Volume],
         instance_mounts: List[InstanceMountPoint],
+        instance_id: str,
     ) -> bool:
         """
         Returns `True` if submitted and `False` if the shim already has a job (`409 Conflict`).
@@ -320,7 +322,7 @@ class ShimClient:
             ssh_user=ssh_user,
             ssh_key=ssh_key,
             mounts=mounts,
-            volumes=[_volume_to_shim_volume_info(v) for v in volumes],
+            volumes=[_volume_to_shim_volume_info(v, instance_id) for v in volumes],
             instance_mounts=instance_mounts,
         )
         resp = self._request("POST", "/api/submit", body)
@@ -398,10 +400,11 @@ def health_response_to_health_status(data: HealthcheckResponse) -> HealthStatus:
         )
 
 
-def _volume_to_shim_volume_info(volume: Volume) -> ShimVolumeInfo:
+def _volume_to_shim_volume_info(volume: Volume, instance_id: str) -> ShimVolumeInfo:
     device_name = None
-    if volume.attachment_data is not None:
-        device_name = volume.attachment_data.device_name
+    attachment_data = volume.get_attachment_data_for_instance(instance_id)
+    if attachment_data is not None:
+        device_name = attachment_data.device_name
     return ShimVolumeInfo(
         backend=volume.configuration.backend.value,
         name=volume.name,