Make gateway deletion async

r4victor · r4victor · commit 94de210f1df5 · 2026-02-20T13:59:15.000+05:00
diff --git a/src/dstack/_internal/server/background/pipeline_tasks/gateways.py b/src/dstack/_internal/server/background/pipeline_tasks/gateways.py
@@ -4,9 +4,10 @@
 from datetime import timedelta
 from typing import Optional, Sequence
 
-from sqlalchemy import or_, select, update
+from sqlalchemy import delete, or_, select, update
 from sqlalchemy.orm import joinedload, load_only
 
+from dstack._internal.core.backends.base.compute import ComputeWithGatewaySupport
 from dstack._internal.core.errors import BackendError, BackendNotAvailable
 from dstack._internal.core.models.gateways import GatewayStatus
 from dstack._internal.server.background.pipeline_tasks.base import (
@@ -24,14 +25,16 @@
     GatewayComputeModel,
     GatewayModel,
     ProjectModel,
+    UserModel,
 )
 from dstack._internal.server.services import backends as backends_services
+from dstack._internal.server.services import events
 from dstack._internal.server.services import gateways as gateways_services
 from dstack._internal.server.services.gateways import emit_gateway_status_change_event
 from dstack._internal.server.services.gateways.pool import gateway_connections_pool
 from dstack._internal.server.services.locking import get_locker
 from dstack._internal.server.services.logging import fmt
-from dstack._internal.utils.common import get_current_datetime
+from dstack._internal.utils.common import get_current_datetime, run_async
 from dstack._internal.utils.logging import get_logger
 
 logger = get_logger(__name__)
@@ -40,6 +43,7 @@
 @dataclass
 class GatewayPipelineItem(PipelineItem):
     status: GatewayStatus
+    to_be_deleted: bool
 
 
 class GatewayPipeline(Pipeline[GatewayPipelineItem]):
@@ -121,8 +125,11 @@ async def fetch(self, limit: int) -> list[GatewayPipelineItem]:
                 res = await session.execute(
                     select(GatewayModel)
                     .where(
-                        GatewayModel.status.in_(
-                            [GatewayStatus.SUBMITTED, GatewayStatus.PROVISIONING]
+                        or_(
+                            GatewayModel.status.in_(
+                                [GatewayStatus.SUBMITTED, GatewayStatus.PROVISIONING]
+                            ),
+                            GatewayModel.to_be_deleted == True,
                         ),
                         or_(
                             GatewayModel.last_processed_at <= now - self._min_processing_interval,
@@ -143,9 +150,10 @@ async def fetch(self, limit: int) -> list[GatewayPipelineItem]:
                     .options(
                         load_only(
                             GatewayModel.id,
-                            GatewayModel.status,
                             GatewayModel.lock_token,
                             GatewayModel.lock_expires_at,
+                            GatewayModel.status,
+                            GatewayModel.to_be_deleted,
                         )
                     )
                 )
@@ -166,6 +174,7 @@ async def fetch(self, limit: int) -> list[GatewayPipelineItem]:
                             lock_token=lock_token,
                             prev_lock_expired=prev_lock_expired,
                             status=gateway_model.status,
+                            to_be_deleted=gateway_model.to_be_deleted,
                         )
                     )
                 await session.commit()
@@ -184,7 +193,9 @@ def __init__(
         )
 
     async def process(self, item: GatewayPipelineItem):
-        if item.status == GatewayStatus.SUBMITTED:
+        if item.to_be_deleted:
+            await _process_to_be_deleted_item(item)
+        elif item.status == GatewayStatus.SUBMITTED:
             await _process_submitted_item(item)
         elif item.status == GatewayStatus.PROVISIONING:
             await _process_provisioning_item(item)
@@ -235,6 +246,7 @@ async def _process_submitted_item(item: GatewayPipelineItem):
                 item.__tablename__,
                 item.id,
             )
+            # TODO: Clean up gateway_compute_model.
             return
         emit_gateway_status_change_event(
             session=session,
@@ -407,3 +419,132 @@ async def _process_provisioning_gateway(gateway_model: GatewayModel) -> _Provisi
     return _ProvisioningResult(
         gateway_update_map={"status": GatewayStatus.RUNNING},
     )
+
+
+async def _process_to_be_deleted_item(item: GatewayPipelineItem):
+    async with get_session_ctx() as session:
+        res = await session.execute(
+            select(GatewayModel)
+            .where(
+                GatewayModel.id == item.id,
+                GatewayModel.lock_token == item.lock_token,
+            )
+            .options(joinedload(GatewayModel.project).joinedload(ProjectModel.backends))
+            .options(joinedload(GatewayModel.gateway_compute))
+            .options(
+                joinedload(GatewayModel.deleted_by_user).load_only(UserModel.id, UserModel.name)
+            )
+        )
+        gateway_model = res.unique().scalar_one_or_none()
+        if gateway_model is None:
+            logger.warning(
+                "Failed to process %s item %s: lock_token mismatch."
+                " The item is expected to be processed and updated on another fetch iteration.",
+                item.__tablename__,
+                item.id,
+            )
+            return
+
+    result = await _process_to_be_deleted_gateway(gateway_model)
+    async with get_session_ctx() as session:
+        if result.delete_gateway:
+            res = await session.execute(
+                delete(GatewayModel)
+                .where(
+                    GatewayModel.id == gateway_model.id,
+                    GatewayModel.lock_token == gateway_model.lock_token,
+                )
+                .returning(GatewayModel.id)
+            )
+            deleted_ids = list(res.scalars().all())
+            if len(deleted_ids) == 0:
+                logger.warning(
+                    "Failed to delete %s item %s after processing: lock_token changed."
+                    " The item is expected to be processed and deleted on another fetch iteration.",
+                    item.__tablename__,
+                    item.id,
+                )
+                return
+            actor = events.SystemActor()
+            if gateway_model.deleted_by_user is not None:
+                actor = events.UserActor.from_user(gateway_model.deleted_by_user)
+            events.emit(
+                session,
+                "Gateway deleted",
+                actor=actor,
+                targets=[events.Target.from_model(gateway_model)],
+            )
+        else:
+            res = await session.execute(
+                update(GatewayModel)
+                .where(
+                    GatewayModel.id == gateway_model.id,
+                    GatewayModel.lock_token == gateway_model.lock_token,
+                )
+                .values(**get_processed_update_map())
+                .returning(GatewayModel.id)
+            )
+            updated_ids = list(res.scalars().all())
+            if len(updated_ids) == 0:
+                logger.warning(
+                    "Failed to update %s item %s after processing: lock_token changed."
+                    " The item is expected to be processed and updated on another fetch iteration.",
+                    item.__tablename__,
+                    item.id,
+                )
+                return
+
+        if result.gateway_compute_update_map:
+            res = await session.execute(
+                update(GatewayComputeModel)
+                .where(GatewayComputeModel.id == gateway_model.gateway_compute_id)
+                .values(**result.gateway_compute_update_map)
+                .returning(GatewayComputeModel.id)
+            )
+            updated_ids = list(res.scalars().all())
+            if len(updated_ids) == 0:
+                logger.error(
+                    "Failed to update compute model %s for gateway %s."
+                    " This is unexpected and may happen only if the compute model was manually deleted.",
+                    gateway_model.id,
+                    item.id,
+                )
+                return
+
+
+@dataclass
+class _DeletedResult:
+    delete_gateway: bool
+    gateway_compute_update_map: UpdateMap = field(default_factory=dict)
+
+
+async def _process_to_be_deleted_gateway(gateway_model: GatewayModel) -> _DeletedResult:
+    backend = await backends_services.get_project_backend_by_type_or_error(
+        project=gateway_model.project, backend_type=gateway_model.backend.type
+    )
+    compute = backend.compute()
+    assert isinstance(compute, ComputeWithGatewaySupport)
+    gateway_compute_configuration = gateways_services.get_gateway_compute_configuration(
+        gateway_model
+    )
+    if gateway_model.gateway_compute is not None and gateway_compute_configuration is not None:
+        logger.info("Deleting gateway compute for %s...", gateway_model.name)
+        try:
+            await run_async(
+                compute.terminate_gateway,
+                gateway_model.gateway_compute.instance_id,
+                gateway_compute_configuration,
+                gateway_model.gateway_compute.backend_data,
+            )
+        except Exception:
+            logger.exception(
+                "Error when deleting gateway compute for %s",
+                gateway_model.name,
+            )
+            return _DeletedResult(delete_gateway=False)
+        logger.info("Deleted gateway compute for %s", gateway_model.name)
+    result = _DeletedResult(delete_gateway=True)
+    if gateway_model.gateway_compute is not None:
+        await gateway_connections_pool.remove(gateway_model.gateway_compute.ip_address)
+        result.gateway_compute_update_map = {"active": False, "deleted": True}
+    return result
diff --git a/src/dstack/_internal/server/migrations/versions/2026/02_20_0840_d911914ecf17_add_gatewaymodel_deleted_columns.py b/src/dstack/_internal/server/migrations/versions/2026/02_20_0840_d911914ecf17_add_gatewaymodel_deleted_columns.py
@@ -0,0 +1,53 @@
+"""Add GatewayModel deleted columns
+
+Revision ID: d911914ecf17
+Revises: 66c2fdda33c2
+Create Date: 2026-02-20 08:40:11.616560+00:00
+
+"""
+
+import sqlalchemy as sa
+import sqlalchemy_utils
+from alembic import op
+
+# revision identifiers, used by Alembic.
+revision = "d911914ecf17"
+down_revision = "66c2fdda33c2"
+branch_labels = None
+depends_on = None
+
+
+def upgrade() -> None:
+    # ### commands auto generated by Alembic - please adjust! ###
+    with op.batch_alter_table("gateways", schema=None) as batch_op:
+        batch_op.add_column(
+            sa.Column("to_be_deleted", sa.Boolean(), server_default=sa.text("0"), nullable=False)
+        )
+        batch_op.add_column(
+            sa.Column(
+                "deleted_by_user_id",
+                sqlalchemy_utils.types.uuid.UUIDType(binary=False),
+                nullable=True,
+            )
+        )
+        batch_op.create_foreign_key(
+            batch_op.f("fk_gateways_deleted_by_user_id_users"),
+            "users",
+            ["deleted_by_user_id"],
+            ["id"],
+            ondelete="SET NULL",
+        )
+
+    # ### end Alembic commands ###
+
+
+def downgrade() -> None:
+    # ### commands auto generated by Alembic - please adjust! ###
+    with op.batch_alter_table("gateways", schema=None) as batch_op:
+        batch_op.drop_constraint(
+            batch_op.f("fk_gateways_deleted_by_user_id_users"), type_="foreignkey"
+        )
+        batch_op.drop_column("deleted_by_user_id")
+        batch_op.drop_column("to_be_deleted")
+
+    # ### end Alembic commands ###
diff --git a/src/dstack/_internal/server/models.py b/src/dstack/_internal/server/models.py
@@ -521,6 +521,12 @@ class GatewayModel(PipelineModelMixin, BaseModel):
 
     runs: Mapped[List["RunModel"]] = relationship(back_populates="gateway")
 
+    to_be_deleted: Mapped[bool] = mapped_column(Boolean, server_default=false())
+    deleted_by_user_id: Mapped[Optional[uuid.UUID]] = mapped_column(
+        ForeignKey("users.id", ondelete="SET NULL")
+    )
+    deleted_by_user: Mapped[Optional[UserModel]] = relationship(foreign_keys=[deleted_by_user_id])
+
     __table_args__ = (UniqueConstraint("project_id", "name", name="uq_gateways_project_id_name"),)
 
 
diff --git a/src/dstack/_internal/server/services/gateways/__init__.py b/src/dstack/_internal/server/services/gateways/__init__.py
@@ -10,7 +10,6 @@
 import httpx
 from sqlalchemy import func, select, update
 from sqlalchemy.ext.asyncio import AsyncSession
-from sqlalchemy.orm import selectinload
 
 import dstack._internal.utils.random_names as random_names
 from dstack._internal.core.backends.base.compute import (
@@ -50,7 +49,6 @@
 from dstack._internal.server.services import events
 from dstack._internal.server.services.backends import (
     check_backend_type_available,
-    get_project_backend_by_type_or_error,
     get_project_backend_with_model_by_type_or_error,
 )
 from dstack._internal.server.services.gateways.connection import GatewayConnection
@@ -295,53 +293,24 @@ async def delete_gateways(
         res = await session.execute(
             select(GatewayModel)
             .where(
+                GatewayModel.id.in_(gateways_ids),
                 GatewayModel.project_id == project.id,
-                GatewayModel.name.in_(gateways_names),
+                GatewayModel.lock_expires_at.is_(None),
             )
-            .options(selectinload(GatewayModel.gateway_compute))
             .execution_options(populate_existing=True)
             .order_by(GatewayModel.id)  # take locks in order
-            .with_for_update(key_share=True)
+            .with_for_update(key_share=True, nowait=True)
         )
         gateway_models = res.scalars().all()
-        for gateway_model in gateway_models:
-            backend = await get_project_backend_by_type_or_error(
-                project=project, backend_type=gateway_model.backend.type
-            )
-            compute = backend.compute()
-            assert isinstance(compute, ComputeWithGatewaySupport)
-            gateway_compute_configuration = get_gateway_compute_configuration(gateway_model)
-            if (
-                gateway_model.gateway_compute is not None
-                and gateway_compute_configuration is not None
-            ):
-                logger.info("Deleting gateway compute for %s...", gateway_model.name)
-                try:
-                    await run_async(
-                        compute.terminate_gateway,
-                        gateway_model.gateway_compute.instance_id,
-                        gateway_compute_configuration,
-                        gateway_model.gateway_compute.backend_data,
-                    )
-                except Exception:
-                    logger.exception(
-                        "Error when deleting gateway compute for %s",
-                        gateway_model.name,
-                    )
-                    continue
-                logger.info("Deleted gateway compute for %s", gateway_model.name)
-            if gateway_model.gateway_compute is not None:
-                await gateway_connections_pool.remove(gateway_model.gateway_compute.ip_address)
-                gateway_model.gateway_compute.active = False
-                gateway_model.gateway_compute.deleted = True
-                session.add(gateway_model.gateway_compute)
-            await session.delete(gateway_model)
-            events.emit(
-                session,
-                "Gateway deleted",
-                actor=events.UserActor.from_user(user),
-                targets=[events.Target.from_model(gateway_model)],
+        if len(gateway_models) != len(gateways_ids):
+            # TODO: Make the delete endpoint fully async without lock – put the request in queue and process in background.
+            raise ServerClientError(
+                "Failed to delete gateways: gateways are being processed currently. Try again later."
             )
+        for gateway_model in gateway_models:
+            if not gateway_model.to_be_deleted:
+                gateway_model.to_be_deleted = True
+                gateway_model.deleted_by_user_id = user.id
         await session.commit()
 
 
