dstack 0.18.43__py3-none-any.whl → 0.19.0rc1__py3-none-any.whl

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

@@ -6,7 +6,8 @@ from sqlalchemy import select
 from sqlalchemy.ext.asyncio import AsyncSession
 from sqlalchemy.orm import joinedload, lazyload, selectinload
 
-from dstack._internal.core.backends.base import Backend
+from dstack._internal.core.backends.base.backend import Backend
+from dstack._internal.core.backends.base.compute import ComputeWithVolumeSupport
 from dstack._internal.core.errors import BackendError, ServerClientError
 from dstack._internal.core.models.common import NetworkMode
 from dstack._internal.core.models.fleets import (
@@ -17,10 +18,8 @@ from dstack._internal.core.models.fleets import (
 )
 from dstack._internal.core.models.instances import InstanceOfferWithAvailability, InstanceStatus
 from dstack._internal.core.models.profiles import (
-    DEFAULT_POOL_NAME,
     DEFAULT_RUN_TERMINATION_IDLE_TIME,
     CreationPolicy,
-    Profile,
     TerminationPolicy,
 )
 from dstack._internal.core.models.resources import Memory
@@ -35,12 +34,12 @@ from dstack._internal.core.models.runs import (
 )
 from dstack._internal.core.models.volumes import Volume
 from dstack._internal.core.services.profiles import get_termination
+from dstack._internal.server import settings
 from dstack._internal.server.db import get_db, get_session_ctx
 from dstack._internal.server.models import (
     FleetModel,
     InstanceModel,
     JobModel,
-    PoolModel,
     ProjectModel,
     RunModel,
     VolumeAttachmentModel,
@@ -50,6 +49,12 @@ from dstack._internal.server.services.backends import get_project_backend_by_typ
 from dstack._internal.server.services.fleets import (
     fleet_model_to_fleet,
 )
+from dstack._internal.server.services.instances import (
+    filter_pool_instances,
+    get_instance_offer,
+    get_instance_provisioning_data,
+    get_shared_pool_instances_with_offers,
+)
 from dstack._internal.server.services.jobs import (
     check_can_attach_job_volumes,
     find_job,
@@ -61,12 +66,6 @@ from dstack._internal.server.services.jobs import (
 from dstack._internal.server.services.locking import get_locker
 from dstack._internal.server.services.logging import fmt
 from dstack._internal.server.services.offers import get_offers_by_requirements
-from dstack._internal.server.services.pools import (
-    filter_pool_instances,
-    get_instance_offer,
-    get_instance_provisioning_data,
-    get_shared_pool_instances_with_offers,
-)
 from dstack._internal.server.services.runs import (
     check_run_spec_requires_instance_mounts,
     run_model_to_run,
@@ -172,29 +171,27 @@ async def _process_submitted_job(session: AsyncSession, job_model: JobModel):
     except ServerClientError as e:
         logger.warning("%s: failed to prepare run volumes: %s", fmt(job_model), repr(e))
         job_model.status = JobStatus.TERMINATING
-        # TODO: Replace with JobTerminationReason.VOLUME_ERROR in 0.19
-        job_model.termination_reason = JobTerminationReason.TERMINATED_BY_SERVER
+        job_model.termination_reason = JobTerminationReason.VOLUME_ERROR
         job_model.termination_reason_message = e.msg
         job_model.last_processed_at = common_utils.get_current_datetime()
         await session.commit()
         return
 
-    pool = await _get_pool(session=session, project=project, profile=profile)
-
     # Submitted jobs processing happens in two steps (transactions).
     # First, the jobs gets an instance assigned (or no instance).
     # Then, the job runs on the assigned instance or a new instance is provisioned.
     # This is needed to avoid holding instances lock for a long time.
     if not job_model.instance_assigned:
-        # Try assigning instances from the pool.
+        # Try assigning an existing instance
         res = await session.execute(
             select(InstanceModel)
             .where(
-                InstanceModel.pool_id == pool.id,
+                InstanceModel.project_id == project.id,
                 InstanceModel.deleted == False,
                 InstanceModel.total_blocks > InstanceModel.busy_blocks,
             )
             .options(lazyload(InstanceModel.jobs))
+            .order_by(InstanceModel.id)  # take locks in order
             .with_for_update()
         )
         pool_instances = list(res.unique().scalars().all())
@@ -287,7 +284,6 @@ async def _process_submitted_job(session: AsyncSession, job_model: JobModel):
         )
         instance = _create_instance_model_for_job(
             project=project,
-            pool=pool,
             fleet_model=fleet_model,
             run_spec=run_spec,
             job_model=job_model,
@@ -319,6 +315,7 @@ async def _process_submitted_job(session: AsyncSession, job_model: JobModel):
         select(VolumeModel)
         .where(VolumeModel.id.in_(volumes_ids))
         .options(selectinload(VolumeModel.user))
+        .order_by(VolumeModel.id)  # take locks in order
         .with_for_update()
     )
     async with get_locker().lock_ctx(VolumeModel.__tablename__, volumes_ids):
@@ -334,19 +331,6 @@ async def _process_submitted_job(session: AsyncSession, job_model: JobModel):
         await session.commit()
 
 
-async def _get_pool(session: AsyncSession, project: ProjectModel, profile: Profile) -> PoolModel:
-    res = await session.execute(
-        select(PoolModel)
-        .where(
-            PoolModel.project_id == project.id,
-            PoolModel.name == (profile.pool_name or DEFAULT_POOL_NAME),
-            PoolModel.deleted == False,
-        )
-        .options(lazyload(PoolModel.instances))
-    )
-    return res.scalar_one()
-
-
 async def _assign_job_to_pool_instance(
     session: AsyncSession,
     pool_instances: List[InstanceModel],
@@ -450,7 +434,7 @@ async def _run_job_on_new_instance(
     )
     # Limit number of offers tried to prevent long-running processing
     # in case all offers fail.
-    for backend, offer in offers[:15]:
+    for backend, offer in offers[: settings.MAX_OFFERS_TRIED]:
         logger.debug(
             "%s: trying %s in %s/%s for $%0.4f per hour",
             fmt(job_model),
@@ -545,7 +529,6 @@ async def _get_next_instance_num(session: AsyncSession, fleet_model: FleetModel)
 
 def _create_instance_model_for_job(
     project: ProjectModel,
-    pool: PoolModel,
     fleet_model: FleetModel,
     run_spec: RunSpec,
     job_model: JobModel,
@@ -568,7 +551,6 @@ def _create_instance_model_for_job(
         name=f"{fleet_model.name}-{instance_num}",
         instance_num=instance_num,
         project=project,
-        pool=pool,
         created_at=common_utils.get_current_datetime(),
         started_at=common_utils.get_current_datetime(),
         status=InstanceStatus.PROVISIONING,
@@ -674,8 +656,7 @@ async def _attach_volumes(
             except (ServerClientError, BackendError) as e:
                 logger.warning("%s: failed to attached volume: %s", fmt(job_model), repr(e))
                 job_model.status = JobStatus.TERMINATING
-                # TODO: Replace with JobTerminationReason.VOLUME_ERROR in 0.19
-                job_model.termination_reason = JobTerminationReason.TERMINATED_BY_SERVER
+                job_model.termination_reason = JobTerminationReason.VOLUME_ERROR
                 job_model.termination_reason_message = "Failed to attach volume"
             except Exception:
                 logger.exception(
@@ -683,8 +664,7 @@ async def _attach_volumes(
                     fmt(job_model),
                 )
                 job_model.status = JobStatus.TERMINATING
-                # TODO: Replace with JobTerminationReason.VOLUME_ERROR in 0.19
-                job_model.termination_reason = JobTerminationReason.TERMINATED_BY_SERVER
+                job_model.termination_reason = JobTerminationReason.VOLUME_ERROR
                 job_model.termination_reason_message = "Failed to attach volume"
             finally:
                 job_model.job_runtime_data = job_runtime_data.json()
@@ -697,13 +677,15 @@ async def _attach_volume(
     instance: InstanceModel,
     instance_id: str,
 ):
+    compute = backend.compute()
+    assert isinstance(compute, ComputeWithVolumeSupport)
     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")
     attachment_data = await common_utils.run_async(
-        backend.compute().attach_volume,
+        compute.attach_volume,
         volume=volume,
         instance_id=instance_id,
     )

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

@@ -2,6 +2,7 @@ from sqlalchemy import select
 from sqlalchemy.ext.asyncio import AsyncSession
 from sqlalchemy.orm import joinedload
 
+from dstack._internal.core.backends.base.compute import ComputeWithVolumeSupport
 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
@@ -81,17 +82,19 @@ async def _process_submitted_volume(session: AsyncSession, volume_model: VolumeM
         await session.commit()
         return
 
+    compute = backend.compute()
+    assert isinstance(compute, ComputeWithVolumeSupport)
     try:
         if volume.configuration.volume_id is not None:
             logger.info("Registering external volume %s", volume_model.name)
             vpd = await run_async(
-                backend.compute().register_volume,
+                compute.register_volume,
                 volume=volume,
             )
         else:
             logger.info("Provisioning new volume %s", volume_model.name)
             vpd = await run_async(
-                backend.compute().create_volume,
+                compute.create_volume,
                 volume=volume,
             )
     except BackendError as e:

dstack/_internal/server/migrations/versions/60e444118b6d_add_jobprometheusmetrics.py

@@ -0,0 +1,40 @@
+"""Add JobPrometheusMetrics
+
+Revision ID: 60e444118b6d
+Revises: a751ef183f27
+Create Date: 2025-02-21 10:59:26.339353
+
+"""
+
+import sqlalchemy as sa
+import sqlalchemy_utils
+from alembic import op
+
+import dstack._internal.server.models
+
+# revision identifiers, used by Alembic.
+revision = "60e444118b6d"
+down_revision = "a751ef183f27"
+branch_labels = None
+depends_on = None
+
+
+def upgrade() -> None:
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.create_table(
+        "job_prometheus_metrics",
+        sa.Column("job_id", sqlalchemy_utils.types.uuid.UUIDType(binary=False), nullable=False),
+        sa.Column("collected_at", dstack._internal.server.models.NaiveDateTime(), nullable=False),
+        sa.Column("text", sa.Text(), nullable=False),
+        sa.ForeignKeyConstraint(
+            ["job_id"], ["jobs.id"], name=op.f("fk_job_prometheus_metrics_job_id_jobs")
+        ),
+        sa.PrimaryKeyConstraint("job_id", name=op.f("pk_job_prometheus_metrics")),
+    )
+    # ### end Alembic commands ###
+
+
+def downgrade() -> None:
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.drop_table("job_prometheus_metrics")
+    # ### end Alembic commands ###

dstack/_internal/server/migrations/versions/7bc2586e8b9e_make_instancemodel_pool_id_optional.py

@@ -0,0 +1,36 @@
+"""Make InstanceModel.pool_id optional
+
+Revision ID: 7bc2586e8b9e
+Revises: bc8ca4a505c6
+Create Date: 2025-03-13 11:13:39.748303
+
+"""
+
+import sqlalchemy_utils
+from alembic import op
+
+# revision identifiers, used by Alembic.
+revision = "7bc2586e8b9e"
+down_revision = "bc8ca4a505c6"
+branch_labels = None
+depends_on = None
+
+
+def upgrade() -> None:
+    # ### commands auto generated by Alembic - please adjust! ###
+    with op.batch_alter_table("instances", schema=None) as batch_op:
+        batch_op.alter_column(
+            "pool_id", existing_type=sqlalchemy_utils.UUIDType(binary=False), nullable=True
+        )
+
+    # ### end Alembic commands ###
+
+
+def downgrade() -> None:
+    # ### commands auto generated by Alembic - please adjust! ###
+    with op.batch_alter_table("instances", schema=None) as batch_op:
+        batch_op.alter_column(
+            "pool_id", existing_type=sqlalchemy_utils.UUIDType(binary=False), nullable=False
+        )
+
+    # ### end Alembic commands ###

dstack/_internal/server/migrations/versions/98d1b92988bc_add_jobterminationreason_terminated_due_.py

@@ -0,0 +1,140 @@
+"""Add JobTerminationReason.TERMINATED_DUE_TO_UTILIZATION_POLICY
+
+Revision ID: 98d1b92988bc
+Revises: 60e444118b6d
+Create Date: 2025-02-28 15:12:37.649876
+
+"""
+
+import sqlalchemy as sa
+from alembic import op
+from alembic_postgresql_enum import TableReference
+
+# revision identifiers, used by Alembic.
+revision = "98d1b92988bc"
+down_revision = "60e444118b6d"
+branch_labels = None
+depends_on = None
+
+
+def upgrade() -> None:
+    # SQLite
+    with op.batch_alter_table("jobs", schema=None) as batch_op:
+        batch_op.alter_column(
+            "termination_reason",
+            existing_type=sa.VARCHAR(length=34),
+            type_=sa.Enum(
+                "FAILED_TO_START_DUE_TO_NO_CAPACITY",
+                "INTERRUPTED_BY_NO_CAPACITY",
+                "WAITING_INSTANCE_LIMIT_EXCEEDED",
+                "WAITING_RUNNER_LIMIT_EXCEEDED",
+                "TERMINATED_BY_USER",
+                "VOLUME_ERROR",
+                "GATEWAY_ERROR",
+                "SCALED_DOWN",
+                "DONE_BY_RUNNER",
+                "ABORTED_BY_USER",
+                "TERMINATED_BY_SERVER",
+                "INACTIVITY_DURATION_EXCEEDED",
+                "TERMINATED_DUE_TO_UTILIZATION_POLICY",
+                "CONTAINER_EXITED_WITH_ERROR",
+                "PORTS_BINDING_FAILED",
+                "CREATING_CONTAINER_ERROR",
+                "EXECUTOR_ERROR",
+                "MAX_DURATION_EXCEEDED",
+                name="jobterminationreason",
+            ),
+            existing_nullable=True,
+        )
+    # PostgreSQL
+    op.sync_enum_values(
+        enum_schema="public",
+        enum_name="jobterminationreason",
+        new_values=[
+            "FAILED_TO_START_DUE_TO_NO_CAPACITY",
+            "INTERRUPTED_BY_NO_CAPACITY",
+            "WAITING_INSTANCE_LIMIT_EXCEEDED",
+            "WAITING_RUNNER_LIMIT_EXCEEDED",
+            "TERMINATED_BY_USER",
+            "VOLUME_ERROR",
+            "GATEWAY_ERROR",
+            "SCALED_DOWN",
+            "DONE_BY_RUNNER",
+            "ABORTED_BY_USER",
+            "TERMINATED_BY_SERVER",
+            "INACTIVITY_DURATION_EXCEEDED",
+            "TERMINATED_DUE_TO_UTILIZATION_POLICY",
+            "CONTAINER_EXITED_WITH_ERROR",
+            "PORTS_BINDING_FAILED",
+            "CREATING_CONTAINER_ERROR",
+            "EXECUTOR_ERROR",
+            "MAX_DURATION_EXCEEDED",
+        ],
+        affected_columns=[
+            TableReference(
+                table_schema="public", table_name="jobs", column_name="termination_reason"
+            )
+        ],
+        enum_values_to_rename=[],
+    )
+
+
+def downgrade() -> None:
+    # SQLite
+    with op.batch_alter_table("jobs", schema=None) as batch_op:
+        batch_op.alter_column(
+            "termination_reason",
+            existing_type=sa.Enum(
+                "FAILED_TO_START_DUE_TO_NO_CAPACITY",
+                "INTERRUPTED_BY_NO_CAPACITY",
+                "WAITING_INSTANCE_LIMIT_EXCEEDED",
+                "WAITING_RUNNER_LIMIT_EXCEEDED",
+                "TERMINATED_BY_USER",
+                "VOLUME_ERROR",
+                "GATEWAY_ERROR",
+                "SCALED_DOWN",
+                "DONE_BY_RUNNER",
+                "ABORTED_BY_USER",
+                "TERMINATED_BY_SERVER",
+                "INACTIVITY_DURATION_EXCEEDED",
+                "TERMINATED_DUE_TO_UTILIZATION_POLICY",
+                "CONTAINER_EXITED_WITH_ERROR",
+                "PORTS_BINDING_FAILED",
+                "CREATING_CONTAINER_ERROR",
+                "EXECUTOR_ERROR",
+                "MAX_DURATION_EXCEEDED",
+                name="jobterminationreason",
+            ),
+            type_=sa.VARCHAR(length=34),
+            existing_nullable=True,
+        )
+    # PostgreSQL
+    op.sync_enum_values(
+        enum_schema="public",
+        enum_name="jobterminationreason",
+        new_values=[
+            "FAILED_TO_START_DUE_TO_NO_CAPACITY",
+            "INTERRUPTED_BY_NO_CAPACITY",
+            "WAITING_INSTANCE_LIMIT_EXCEEDED",
+            "WAITING_RUNNER_LIMIT_EXCEEDED",
+            "TERMINATED_BY_USER",
+            "VOLUME_ERROR",
+            "GATEWAY_ERROR",
+            "SCALED_DOWN",
+            "DONE_BY_RUNNER",
+            "ABORTED_BY_USER",
+            "TERMINATED_BY_SERVER",
+            "INACTIVITY_DURATION_EXCEEDED",
+            "CONTAINER_EXITED_WITH_ERROR",
+            "PORTS_BINDING_FAILED",
+            "CREATING_CONTAINER_ERROR",
+            "EXECUTOR_ERROR",
+            "MAX_DURATION_EXCEEDED",
+        ],
+        affected_columns=[
+            TableReference(
+                table_schema="public", table_name="jobs", column_name="termination_reason"
+            )
+        ],
+        enum_values_to_rename=[],
+    )

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

@@ -0,0 +1,171 @@
+"""Store BackendType as string
+
+Revision ID: bc8ca4a505c6
+Revises: 98d1b92988bc
+Create Date: 2025-03-10 14:49:06.837118
+
+"""
+
+import sqlalchemy as sa
+from alembic import op
+from sqlalchemy.dialects import postgresql
+
+# revision identifiers, used by Alembic.
+revision = "bc8ca4a505c6"
+down_revision = "98d1b92988bc"
+branch_labels = None
+depends_on = None
+
+
+def upgrade() -> None:
+    # ### commands auto generated by Alembic - please adjust! ###
+    with op.batch_alter_table("backends", schema=None) as batch_op:
+        batch_op.alter_column(
+            "type",
+            existing_type=postgresql.ENUM(
+                "AWS",
+                "AZURE",
+                "CUDO",
+                "DATACRUNCH",
+                "DSTACK",
+                "GCP",
+                "KUBERNETES",
+                "LAMBDA",
+                "LOCAL",
+                "REMOTE",
+                "NEBIUS",
+                "OCI",
+                "RUNPOD",
+                "TENSORDOCK",
+                "VASTAI",
+                "VULTR",
+                name="backendtype",
+            ),
+            type_=sa.String(length=100),
+            existing_nullable=False,
+        )
+
+    with op.batch_alter_table("instances", schema=None) as batch_op:
+        batch_op.alter_column(
+            "backend",
+            existing_type=postgresql.ENUM(
+                "AWS",
+                "AZURE",
+                "CUDO",
+                "DATACRUNCH",
+                "DSTACK",
+                "GCP",
+                "KUBERNETES",
+                "LAMBDA",
+                "LOCAL",
+                "REMOTE",
+                "NEBIUS",
+                "OCI",
+                "RUNPOD",
+                "TENSORDOCK",
+                "VASTAI",
+                "VULTR",
+                name="backendtype",
+            ),
+            type_=sa.String(length=100),
+            existing_nullable=True,
+        )
+
+    sa.Enum(
+        "AWS",
+        "AZURE",
+        "CUDO",
+        "DATACRUNCH",
+        "DSTACK",
+        "GCP",
+        "KUBERNETES",
+        "LAMBDA",
+        "LOCAL",
+        "REMOTE",
+        "NEBIUS",
+        "OCI",
+        "RUNPOD",
+        "TENSORDOCK",
+        "VASTAI",
+        "VULTR",
+        name="backendtype",
+    ).drop(op.get_bind())
+    # ### end Alembic commands ###
+
+
+def downgrade() -> None:
+    # ### commands auto generated by Alembic - please adjust! ###
+    sa.Enum(
+        "AWS",
+        "AZURE",
+        "CUDO",
+        "DATACRUNCH",
+        "DSTACK",
+        "GCP",
+        "KUBERNETES",
+        "LAMBDA",
+        "LOCAL",
+        "REMOTE",
+        "NEBIUS",
+        "OCI",
+        "RUNPOD",
+        "TENSORDOCK",
+        "VASTAI",
+        "VULTR",
+        name="backendtype",
+    ).create(op.get_bind())
+    with op.batch_alter_table("instances", schema=None) as batch_op:
+        batch_op.alter_column(
+            "backend",
+            existing_type=sa.String(length=100),
+            type_=postgresql.ENUM(
+                "AWS",
+                "AZURE",
+                "CUDO",
+                "DATACRUNCH",
+                "DSTACK",
+                "GCP",
+                "KUBERNETES",
+                "LAMBDA",
+                "LOCAL",
+                "REMOTE",
+                "NEBIUS",
+                "OCI",
+                "RUNPOD",
+                "TENSORDOCK",
+                "VASTAI",
+                "VULTR",
+                name="backendtype",
+            ),
+            existing_nullable=True,
+            postgresql_using="backend::VARCHAR::backendtype",
+        )
+
+    with op.batch_alter_table("backends", schema=None) as batch_op:
+        batch_op.alter_column(
+            "type",
+            existing_type=sa.String(length=100),
+            type_=postgresql.ENUM(
+                "AWS",
+                "AZURE",
+                "CUDO",
+                "DATACRUNCH",
+                "DSTACK",
+                "GCP",
+                "KUBERNETES",
+                "LAMBDA",
+                "LOCAL",
+                "REMOTE",
+                "NEBIUS",
+                "OCI",
+                "RUNPOD",
+                "TENSORDOCK",
+                "VASTAI",
+                "VULTR",
+                name="backendtype",
+            ),
+            existing_nullable=False,
+            postgresql_using="type::VARCHAR::backendtype",
+        )
+
+    # ### end Alembic commands ###