dstack 0.0.9__py3-none-any.whl → 0.20.7__py3-none-any.whl

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

@@ -0,0 +1,108 @@
+import asyncio
+
+from sqlalchemy import or_, select
+from sqlalchemy.ext.asyncio import AsyncSession
+from sqlalchemy.orm import joinedload
+
+from dstack._internal.core.models.runs import JobStatus
+from dstack._internal.server.db import get_db, get_session_ctx
+from dstack._internal.server.models import (
+    InstanceModel,
+    JobModel,
+    ProjectModel,
+    VolumeAttachmentModel,
+)
+from dstack._internal.server.services.jobs import (
+    process_terminating_job,
+    process_volumes_detaching,
+)
+from dstack._internal.server.services.locking import get_locker
+from dstack._internal.server.services.logging import fmt
+from dstack._internal.server.utils import sentry_utils
+from dstack._internal.utils.common import (
+    get_current_datetime,
+    get_or_error,
+)
+from dstack._internal.utils.logging import get_logger
+
+logger = get_logger(__name__)
+
+
+async def process_terminating_jobs(batch_size: int = 1):
+    tasks = []
+    for _ in range(batch_size):
+        tasks.append(_process_next_terminating_job())
+    await asyncio.gather(*tasks)
+
+
+@sentry_utils.instrument_background_task
+async def _process_next_terminating_job():
+    job_lock, job_lockset = get_locker(get_db().dialect_name).get_lockset(JobModel.__tablename__)
+    instance_lock, instance_lockset = get_locker(get_db().dialect_name).get_lockset(
+        InstanceModel.__tablename__
+    )
+    async with get_session_ctx() as session:
+        async with job_lock, instance_lock:
+            res = await session.execute(
+                select(JobModel)
+                .where(
+                    JobModel.id.not_in(job_lockset),
+                    JobModel.status == JobStatus.TERMINATING,
+                    or_(
+                        JobModel.remove_at.is_(None),
+                        JobModel.remove_at < get_current_datetime(),
+                    ),
+                )
+                .order_by(JobModel.last_processed_at.asc())
+                .limit(1)
+                .with_for_update(skip_locked=True, key_share=True)
+            )
+            job_model = res.scalar()
+            if job_model is None:
+                return
+            if job_model.used_instance_id is not None:
+                res = await session.execute(
+                    select(InstanceModel)
+                    .where(
+                        InstanceModel.id == job_model.used_instance_id,
+                        InstanceModel.id.not_in(instance_lockset),
+                    )
+                    .with_for_update(skip_locked=True, key_share=True)
+                )
+                instance_model = res.scalar()
+                if instance_model is None:
+                    # InstanceModel is locked
+                    return
+                instance_lockset.add(instance_model.id)
+            job_lockset.add(job_model.id)
+        job_model_id = job_model.id
+        instance_model_id = job_model.used_instance_id
+        try:
+            await _process_job(
+                session=session,
+                job_model=job_model,
+            )
+        finally:
+            job_lockset.difference_update([job_model_id])
+            instance_lockset.difference_update([instance_model_id])
+
+
+async def _process_job(session: AsyncSession, job_model: JobModel):
+    logger.debug("%s: terminating job", fmt(job_model))
+    res = await session.execute(
+        select(InstanceModel)
+        .where(InstanceModel.id == job_model.used_instance_id)
+        .options(
+            joinedload(InstanceModel.project).joinedload(ProjectModel.backends),
+            joinedload(InstanceModel.volume_attachments).joinedload(VolumeAttachmentModel.volume),
+            joinedload(InstanceModel.jobs).load_only(JobModel.id),
+        )
+    )
+    instance_model = res.unique().scalar()
+    if job_model.volumes_detached_at is None:
+        await process_terminating_job(session, job_model, instance_model)
+    else:
+        instance_model = get_or_error(instance_model)
+        await process_volumes_detaching(session, job_model, instance_model)
+    job_model.last_processed_at = get_current_datetime()
+    await session.commit()

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

@@ -0,0 +1,129 @@
+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_db, get_session_ctx
+from dstack._internal.server.models import (
+    FleetModel,
+    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
+from dstack._internal.server.services.volumes import switch_volume_status
+from dstack._internal.server.utils import sentry_utils
+from dstack._internal.utils.common import get_current_datetime, run_async
+from dstack._internal.utils.logging import get_logger
+
+logger = get_logger(__name__)
+
+
+@sentry_utils.instrument_background_task
+async def process_submitted_volumes():
+    lock, lockset = get_locker(get_db().dialect_name).get_lockset(VolumeModel.__tablename__)
+    async with get_session_ctx() as session:
+        async with lock:
+            res = await session.execute(
+                select(VolumeModel)
+                .where(
+                    VolumeModel.status == VolumeStatus.SUBMITTED,
+                    VolumeModel.id.not_in(lockset),
+                )
+                .order_by(VolumeModel.last_processed_at.asc())
+                .limit(1)
+                .with_for_update(skip_locked=True, key_share=True)
+            )
+            volume_model = res.scalar()
+            if volume_model is None:
+                return
+            lockset.add(volume_model.id)
+        volume_model_id = volume_model.id
+        try:
+            await _process_submitted_volume(session=session, volume_model=volume_model)
+        finally:
+            lockset.difference_update([volume_model_id])
+
+
+async def _process_submitted_volume(session: AsyncSession, volume_model: VolumeModel):
+    logger.info("Started submitted volume %s processing", volume_model.name)
+    # Refetch to load related attributes.
+    res = await session.execute(
+        select(VolumeModel)
+        .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)
+            .load_only(FleetModel.name)
+        )
+        .execution_options(populate_existing=True)
+    )
+    volume_model = res.unique().scalar_one()
+    volume = volumes_services.volume_model_to_volume(volume_model)
+    try:
+        backend = await backends_services.get_project_backend_by_type_or_error(
+            project=volume_model.project,
+            backend_type=volume.configuration.backend,
+            overrides=True,
+        )
+    except BackendNotAvailable:
+        logger.error(
+            "Failed to process volume %s. Backend %s not available.",
+            volume.name,
+            volume.configuration.backend.value,
+        )
+        volume_model.status_message = "Backend not available"
+        switch_volume_status(session, volume_model, VolumeStatus.FAILED)
+        volume_model.last_processed_at = get_current_datetime()
+        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(
+                compute.register_volume,
+                volume=volume,
+            )
+        else:
+            logger.info("Provisioning new volume %s", volume_model.name)
+            vpd = await run_async(
+                compute.create_volume,
+                volume=volume,
+            )
+    except BackendError as e:
+        logger.info("Failed to create volume %s: %s", volume_model.name, repr(e))
+        status_message = f"Backend error: {repr(e)}"
+        if len(e.args) > 0:
+            status_message = str(e.args[0])
+        volume_model.status_message = status_message
+        switch_volume_status(session, volume_model, VolumeStatus.FAILED)
+        volume_model.last_processed_at = get_current_datetime()
+        await session.commit()
+        return
+    except Exception as e:
+        logger.exception("Got exception when creating volume %s", volume_model.name)
+        volume_model.status_message = f"Unexpected error: {repr(e)}"
+        switch_volume_status(session, volume_model, VolumeStatus.FAILED)
+        volume_model.last_processed_at = get_current_datetime()
+        await session.commit()
+        return
+
+    logger.info("Added new volume %s", volume_model.name)
+
+    # Provisioned volumes marked as active since they become available almost immediately in AWS
+    # TODO: Consider checking volume state
+    volume_model.volume_provisioning_data = vpd.json()
+    switch_volume_status(session, volume_model, VolumeStatus.ACTIVE)
+    volume_model.last_processed_at = get_current_datetime()
+    await session.commit()

dstack/_internal/server/compatibility/common.py

@@ -0,0 +1,20 @@
+from typing import Optional
+
+from packaging.version import Version
+
+from dstack._internal.core.models.instances import (
+    InstanceAvailability,
+    InstanceOfferWithAvailability,
+)
+
+
+def patch_offers_list(
+    offers: list[InstanceOfferWithAvailability], client_version: Optional[Version]
+) -> None:
+    if client_version is None:
+        return
+    # CLIs prior to 0.20.4 incorrectly display the `no_balance` availability in the run/fleet plan
+    if client_version < Version("0.20.4"):
+        for offer in offers:
+            if offer.availability == InstanceAvailability.NO_BALANCE:
+                offer.availability = InstanceAvailability.NOT_AVAILABLE

dstack/_internal/server/compatibility/gpus.py

@@ -0,0 +1,22 @@
+from typing import Optional
+
+from packaging.version import Version
+
+from dstack._internal.core.models.instances import InstanceAvailability
+from dstack._internal.server.schemas.gpus import ListGpusResponse
+
+
+def patch_list_gpus_response(
+    response: ListGpusResponse, client_version: Optional[Version]
+) -> None:
+    if client_version is None:
+        return
+    # CLIs prior to 0.20.4 incorrectly display the `no_balance` availability in `dstack offer --group-by gpu`
+    if client_version < Version("0.20.4"):
+        for gpu in response.gpus:
+            if InstanceAvailability.NO_BALANCE in gpu.availability:
+                gpu.availability = [
+                    a for a in gpu.availability if a != InstanceAvailability.NO_BALANCE
+                ]
+                if InstanceAvailability.NOT_AVAILABLE not in gpu.availability:
+                    gpu.availability.append(InstanceAvailability.NOT_AVAILABLE)

dstack/_internal/server/db.py

@@ -0,0 +1,127 @@
+from contextlib import asynccontextmanager
+from typing import Optional
+
+from alembic import command, config
+from sqlalchemy import AsyncAdaptedQueuePool, event
+from sqlalchemy.engine.interfaces import DBAPIConnection
+from sqlalchemy.ext.asyncio import (
+    AsyncEngine,
+    AsyncSession,
+    async_sessionmaker,
+    create_async_engine,
+)
+from sqlalchemy.pool import ConnectionPoolEntry
+
+from dstack._internal.server import settings
+from dstack._internal.server.services.locking import advisory_lock_ctx
+from dstack._internal.server.settings import DATABASE_URL
+
+
+class Database:
+    def __init__(self, url: str, engine: Optional[AsyncEngine] = None):
+        self.url = url
+        if engine is not None:
+            self.engine = engine
+        else:
+            self.engine = create_async_engine(
+                self.url,
+                echo=settings.SQL_ECHO_ENABLED,
+                poolclass=AsyncAdaptedQueuePool,
+                pool_size=settings.DB_POOL_SIZE,
+                max_overflow=settings.DB_MAX_OVERFLOW,
+            )
+        self.session_maker = async_sessionmaker(
+            bind=self.engine,  # type: ignore[assignment]
+            expire_on_commit=False,
+            class_=AsyncSession,
+        )
+
+        if self.dialect_name == "sqlite":
+
+            @event.listens_for(self.engine.sync_engine, "connect")
+            def set_sqlite_pragma(dbapi_connection: DBAPIConnection, _: ConnectionPoolEntry):
+                cursor = dbapi_connection.cursor()
+                cursor.execute("PRAGMA journal_mode=WAL;")
+                cursor.execute("PRAGMA foreign_keys=ON;")
+                cursor.execute("PRAGMA synchronous=NORMAL;")
+                cursor.execute("PRAGMA busy_timeout=30000;")
+                cursor.close()
+
+    @property
+    def dialect_name(self) -> str:
+        return self.engine.dialect.name
+
+    def get_session(self) -> AsyncSession:
+        return self.session_maker()
+
+
+def get_new_db() -> Database:
+    """
+    Creates a new Database with a new Engine.
+    Use this when you need to access the DB in a new thread instead of calling Database directly
+    since it's easier to monkey-patch.
+    """
+    return Database(url=DATABASE_URL)
+
+
+_db = get_new_db()
+
+
+def get_db() -> Database:
+    return _db
+
+
+def override_db(new_db: Database):
+    global _db
+    _db = new_db
+
+
+async def migrate():
+    async with _db.engine.connect() as connection:
+        async with advisory_lock_ctx(
+            bind=connection,
+            dialect_name=_db.dialect_name,
+            resource="migrations",
+        ):
+            await connection.run_sync(_run_alembic_upgrade)
+
+
+async def get_session():
+    async with _db.get_session() as session:
+        yield session
+        await session.commit()
+
+
+get_session_ctx = asynccontextmanager(get_session)
+
+
+def session_decorator(func):
+    async def new_func(*args, **kwargs):
+        async with get_session_ctx() as s:
+            return await func(*args, session=s, **kwargs)
+
+    return new_func
+
+
+def is_db_sqlite() -> bool:
+    return get_db().dialect_name == "sqlite"
+
+
+def is_db_postgres() -> bool:
+    return get_db().dialect_name == "postgresql"
+
+
+async def sqlite_commit(session: AsyncSession):
+    """
+    Commit an sqlite transaction.
+    Should be used before taking locks in active sessions to see committed changes.
+    """
+    if is_db_sqlite():
+        await session.commit()
+
+
+def _run_alembic_upgrade(connection):
+    alembic_cfg = config.Config()
+    alembic_cfg.set_main_option("script_location", settings.ALEMBIC_MIGRATIONS_LOCATION)
+    alembic_cfg.attributes["connection"] = connection
+    command.upgrade(alembic_cfg, "head")

dstack/_internal/server/deps.py

@@ -0,0 +1,19 @@
+from fastapi import Depends
+from sqlalchemy.ext.asyncio import AsyncSession
+
+from dstack._internal.server.db import get_session
+from dstack._internal.server.models import ProjectModel
+from dstack._internal.server.services.projects import get_project_model_by_name
+from dstack._internal.server.utils.routers import error_not_found
+
+
+class Project:
+    async def __call__(
+        self,
+        project_name: str,
+        session: AsyncSession = Depends(get_session),
+    ) -> ProjectModel:
+        project = await get_project_model_by_name(session=session, project_name=project_name)
+        if project is None:
+            raise error_not_found()
+        return project

dstack/_internal/server/main.py

@@ -0,0 +1,4 @@
+from dstack._internal.server.app import create_app, register_routes
+
+app = create_app()
+register_routes(app)

dstack/_internal/server/migrations/env.py

@@ -0,0 +1,112 @@
+import asyncio
+from logging.config import fileConfig
+
+import alembic_postgresql_enum  # noqa: F401
+from alembic import context
+from sqlalchemy import Connection, MetaData, text
+
+from dstack._internal.server.db import get_db
+from dstack._internal.server.models import BaseModel, EnumAsString
+
+config = context.config
+
+if config.config_file_name is not None and config.attributes.get("configure_logging", True):
+    fileConfig(config.config_file_name)
+
+target_metadata = BaseModel.metadata
+
+
+def set_target_metadata(metadata: MetaData):
+    global target_metadata
+    target_metadata = metadata
+
+
+def render_item(type_, obj, autogen_context):
+    """Apply custom rendering for selected items."""
+    if type_ == "type" and isinstance(obj, EnumAsString):
+        return f"sa.String(length={obj.length})"
+    # default rendering for other objects
+    return False
+
+
+def run_migrations_offline():
+    """Run migrations in 'offline' mode.
+    This configures the context with just a URL
+    and not an Engine, though an Engine is acceptable
+    here as well.  By skipping the Engine creation
+    we don't even need a DBAPI to be available.
+    Calls to context.execute() here emit the given string to the
+    script output.
+    """
+    context.configure(
+        url=get_db().url,
+        target_metadata=target_metadata,
+        literal_binds=True,
+        dialect_opts={"paramstyle": "named"},
+        render_item=render_item,
+    )
+    with context.begin_transaction():
+        context.run_migrations()
+
+
+# Programmatic API use (connection sharing) With Asyncio
+# https://alembic.sqlalchemy.org/en/latest/cookbook.html#programmatic-api-use-connection-sharing-with-asyncio
+def run_migrations_online():
+    """Run migrations in 'online' mode.
+    In this scenario we need to create an Engine
+    and associate a connection with the context.
+    """
+    connection = config.attributes.get("connection", None)
+    if connection is None:
+        asyncio.run(run_async_migrations())
+    else:
+        run_migrations(connection)
+
+
+def run_migrations(connection: Connection):
+    # Temporarily disable foreign keys,
+    # so that sqlite batch table migrations are performed without data loss:
+    # https://alembic.sqlalchemy.org/en/latest/batch.html#dealing-with-referencing-foreign-keys
+    if connection.dialect.name == "sqlite":
+        connection.execute(text("PRAGMA foreign_keys=OFF;"))
+    elif connection.dialect.name == "postgresql":
+        # lock_timeout is needed so that migrations that acquire locks
+        # do not wait for locks forever, blocking live queries.
+        # Better to fail and retry a deployment.
+        connection.execute(text("SET lock_timeout='10s';"))
+    connection.commit()
+    context.configure(
+        connection=connection,
+        target_metadata=target_metadata,
+        compare_type=True,
+        render_as_batch=True,
+        render_item=render_item,
+        # Running each migration in a separate transaction.
+        # Running all migrations in one transaction may lead to deadlocks in HA deployments
+        # because lock ordering is not respected across all migrations.
+        transaction_per_migration=True,
+    )
+    with context.begin_transaction():
+        context.run_migrations()
+    if connection.dialect.name == "sqlite":
+        connection.execute(text("PRAGMA foreign_keys=ON;"))
+    connection.commit()
+
+
+async def run_async_migrations():
+    engine = get_db().engine
+    async with engine.connect() as connection:
+        await connection.run_sync(run_migrations)
+    await engine.dispose()
+
+
+def main():
+    if context.is_offline_mode():
+        run_migrations_offline()
+    else:
+        run_migrations_online()
+
+
+# invoked via alembic command
+if __name__ == "env_py":
+    main()

dstack/_internal/server/migrations/script.py.mako

@@ -0,0 +1,28 @@
+"""${message}
+
+Revision ID: ${up_revision}
+Revises: ${down_revision | comma,n}
+Create Date: ${create_date}
+
+"""
+
+from alembic import op
+import sqlalchemy as sa
+import sqlalchemy_utils
+
+import dstack._internal.server.models
+${imports if imports else ""}
+
+# revision identifiers, used by Alembic.
+revision = ${repr(up_revision)}
+down_revision = ${repr(down_revision)}
+branch_labels = ${repr(branch_labels)}
+depends_on = ${repr(depends_on)}
+
+
+def upgrade() -> None:
+    ${upgrades if upgrades else "pass"}
+
+
+def downgrade() -> None:
+    ${downgrades if downgrades else "pass"}

dstack/_internal/server/migrations/versions/006512f572b4_add_projects_original_name.py

@@ -0,0 +1,38 @@
+"""Add projects.original_name
+
+Revision ID: 006512f572b4
+Revises: 06e977bc61c7
+Create Date: 2025-11-27 15:11:21.249079
+
+"""
+
+import sqlalchemy as sa
+from alembic import op
+
+# revision identifiers, used by Alembic.
+revision = "006512f572b4"
+down_revision = "06e977bc61c7"
+branch_labels = None
+depends_on = None
+
+
+def upgrade() -> None:
+    # ### commands auto generated by Alembic - please adjust! ###
+    # For postgres, this was moved from a previous migration to avoid deadlocks.
+    # Keep SQLite in previous migration since it doesn't have deadlock issues and
+    # does not support if_not_exists.
+    if op.get_context().dialect.name != "sqlite":
+        with op.batch_alter_table("projects", schema=None) as batch_op:
+            batch_op.add_column(
+                sa.Column("original_name", sa.String(length=50), nullable=True),
+                if_not_exists=True,
+            )
+    # ### end Alembic commands ###
+
+
+def downgrade() -> None:
+    # ### commands auto generated by Alembic - please adjust! ###
+    if op.get_context().dialect.name != "sqlite":
+        with op.batch_alter_table("projects", schema=None) as batch_op:
+            batch_op.drop_column("original_name")
+    # ### end Alembic commands ###