Restore sync gateways delete API

Author: r4victor

src/dstack/_internal/server/services/gateways/__init__.py

@@ -10,6 +10,7 @@
 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 (
@@ -49,6 +50,7 @@
 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
@@ -61,6 +63,7 @@
 from dstack._internal.server.services.pipelines import PipelineHinterProtocol
 from dstack._internal.server.services.plugins import apply_plugin_policies
 from dstack._internal.server.utils.common import gather_map_async
+from dstack._internal.settings import FeatureFlags
 from dstack._internal.utils.common import get_current_datetime, run_async
 from dstack._internal.utils.crypto import generate_rsa_key_pair_bytes
 from dstack._internal.utils.logging import get_logger
@@ -278,6 +281,32 @@ async def delete_gateways(
     project: ProjectModel,
     gateways_names: List[str],
     user: UserModel,
+):
+    # Keep both delete code paths while pipeline processing is behind a feature flag:
+    # - pipeline path marks gateways for async deletion by GatewayPipeline
+    # - sync path deletes gateway resources inline for non-pipeline processing
+    # TODO: Drop sync path after pipeline processing is enabled by default.
+    if FeatureFlags.PIPELINE_PROCESSING_ENABLED:
+        await _delete_gateways_pipeline(
+            session=session,
+            project=project,
+            gateways_names=gateways_names,
+            user=user,
+        )
+    else:
+        await _delete_gateways_sync(
+            session=session,
+            project=project,
+            gateways_names=gateways_names,
+            user=user,
+        )
+
+
+async def _delete_gateways_pipeline(
+    session: AsyncSession,
+    project: ProjectModel,
+    gateways_names: List[str],
+    user: UserModel,
 ):
     res = await session.execute(
         select(GatewayModel).where(
@@ -323,6 +352,79 @@ async def delete_gateways(
         await session.commit()
 
 
+async def _delete_gateways_sync(
+    session: AsyncSession,
+    project: ProjectModel,
+    gateways_names: List[str],
+    user: UserModel,
+):
+    res = await session.execute(
+        select(GatewayModel).where(
+            GatewayModel.project_id == project.id,
+            GatewayModel.name.in_(gateways_names),
+        )
+    )
+    gateway_models = res.scalars().all()
+    gateways_ids = sorted([g.id for g in gateway_models])
+    await session.commit()
+    logger.info("Deleting gateways: %s", [g.name for g in gateway_models])
+    async with get_locker(get_db().dialect_name).lock_ctx(
+        GatewayModel.__tablename__, gateways_ids
+    ):
+        # Refetch after lock
+        res = await session.execute(
+            select(GatewayModel)
+            .where(
+                GatewayModel.project_id == project.id,
+                GatewayModel.name.in_(gateways_names),
+            )
+            .options(selectinload(GatewayModel.gateway_compute))
+            .execution_options(populate_existing=True)
+            .order_by(GatewayModel.id)  # take locks in order
+            .with_for_update(key_share=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)],
+            )
+        await session.commit()
+
+
 async def set_gateway_wildcard_domain(
     session: AsyncSession,
     project: ProjectModel,

src/dstack/_internal/server/services/pipelines.py

@@ -11,9 +11,20 @@ def hint_fetch(self, model_name: str) -> None:
         pass
 
 
+class _NoopPipelineHinter:
+    def hint_fetch(self, model_name: str) -> None:
+        pass
+
+
+_noop_pipeline_hinter = _NoopPipelineHinter()
+
+
 def get_pipeline_hinter(request: Request) -> PipelineHinterProtocol:
     """
     Returns pipeline hinter that allows hinting replica's pipelines that there are new items for processing.
     This can reduce processing latency if the processing happens rarely.
     """
-    return request.app.state.pipeline_manager.hinter
+    pipeline_manager = getattr(request.app.state, "pipeline_manager", None)
+    if pipeline_manager is None:
+        return _noop_pipeline_hinter
+    return pipeline_manager.hinter

src/tests/_internal/server/routers/test_gateways.py

@@ -1,13 +1,15 @@
-from unittest.mock import patch
+from unittest.mock import Mock, patch
 
 import pytest
 from httpx import AsyncClient
 from sqlalchemy.ext.asyncio import AsyncSession
 
+from dstack._internal.core.errors import DstackError
 from dstack._internal.core.models.backends.base import BackendType
 from dstack._internal.core.models.users import GlobalRole, ProjectRole
 from dstack._internal.server.services.projects import add_project_member
 from dstack._internal.server.testing.common import (
+    ComputeMockSpec,
     clear_events,
     create_backend,
     create_gateway,
@@ -18,6 +20,15 @@
     list_events,
 )
 from dstack._internal.server.testing.matchers import SomeUUID4Str
+from dstack._internal.settings import FeatureFlags
+
+
+@pytest.fixture
+def patch_pipeline_processing_flag(monkeypatch: pytest.MonkeyPatch):
+    def _apply(enabled: bool):
+        monkeypatch.setattr(FeatureFlags, "PIPELINE_PROCESSING_ENABLED", enabled)
+
+    return _apply
 
 
 class TestListAndGetGateways:
@@ -453,6 +464,12 @@ async def test_only_admin_can_delete(
         )
         assert response.status_code == 403
 
+
+class TestDeleteGatewayPipelineEnabled:
+    @pytest.fixture(autouse=True)
+    def _pipeline_processing_enabled(self, patch_pipeline_processing_flag):
+        patch_pipeline_processing_flag(True)
+
     @pytest.mark.asyncio
     @pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True)
     async def test_marks_gateways_to_be_deleted(
@@ -519,6 +536,110 @@ async def test_marks_gateways_to_be_deleted(
         assert all(e.actor_user_id == user.id for e in events)
 
 
+class TestDeleteGatewayPipelineDisabled:
+    @pytest.fixture(autouse=True)
+    def _pipeline_processing_disabled(self, patch_pipeline_processing_flag):
+        patch_pipeline_processing_flag(False)
+
+    @pytest.mark.asyncio
+    @pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True)
+    async def test_deletes_gateways_synchronously(
+        self, test_db, session: AsyncSession, client: AsyncClient
+    ):
+        user = await create_user(session, global_role=GlobalRole.USER)
+        project = await create_project(session)
+        await add_project_member(
+            session=session, project=project, user=user, project_role=ProjectRole.ADMIN
+        )
+        backend_aws = await create_backend(session, project.id)
+        backend_gcp = await create_backend(session, project.id, backend_type=BackendType.GCP)
+        gateway_compute_aws = await create_gateway_compute(
+            session=session,
+            backend_id=backend_aws.id,
+        )
+        gateway_aws = await create_gateway(
+            session=session,
+            project_id=project.id,
+            backend_id=backend_aws.id,
+            name="gateway-aws",
+            gateway_compute_id=gateway_compute_aws.id,
+        )
+        gateway_compute_gcp = await create_gateway_compute(
+            session=session,
+            backend_id=backend_gcp.id,
+        )
+        gateway_gcp = await create_gateway(
+            session=session,
+            project_id=project.id,
+            backend_id=backend_gcp.id,
+            name="gateway-gcp",
+            gateway_compute_id=gateway_compute_gcp.id,
+        )
+        with patch(
+            "dstack._internal.server.services.gateways.get_project_backend_by_type_or_error"
+        ) as m:
+            aws = Mock()
+            aws.compute.return_value = Mock(spec=ComputeMockSpec)
+            aws.compute.return_value.terminate_gateway.return_value = None  # success
+            gcp = Mock()
+            gcp.compute.return_value = Mock(spec=ComputeMockSpec)
+            gcp.compute.return_value.terminate_gateway.side_effect = DstackError()  # fail
+
+            def get_backend(project, backend_type):
+                return {BackendType.AWS: aws, BackendType.GCP: gcp}[backend_type]
+
+            m.side_effect = get_backend
+
+            response = await client.post(
+                f"/api/project/{project.name}/gateways/delete",
+                json={"names": [gateway_aws.name, gateway_gcp.name]},
+                headers=get_auth_headers(user.token),
+            )
+            aws.compute.return_value.terminate_gateway.assert_called_once()
+            gcp.compute.return_value.terminate_gateway.assert_called_once()
+            assert response.status_code == 200
+
+        response = await client.post(
+            f"/api/project/{project.name}/gateways/list",
+            headers=get_auth_headers(user.token),
+        )
+        assert response.status_code == 200
+        assert response.json() == [
+            {
+                "id": str(gateway_gcp.id),
+                "backend": backend_gcp.type.value,
+                "created_at": response.json()[0]["created_at"],
+                "default": False,
+                "status": "submitted",
+                "status_message": None,
+                "instance_id": gateway_compute_gcp.instance_id,
+                "ip_address": gateway_compute_gcp.ip_address,
+                "hostname": gateway_compute_gcp.ip_address,
+                "name": gateway_gcp.name,
+                "region": gateway_gcp.region,
+                "wildcard_domain": gateway_gcp.wildcard_domain,
+                "configuration": {
+                    "type": "gateway",
+                    "name": gateway_gcp.name,
+                    "backend": backend_gcp.type.value,
+                    "region": gateway_gcp.region,
+                    "instance_type": None,
+                    "router": None,
+                    "domain": gateway_gcp.wildcard_domain,
+                    "default": False,
+                    "public_ip": True,
+                    "certificate": {"type": "lets-encrypt"},
+                    "tags": None,
+                },
+            }
+        ]
+
+        events = await list_events(session)
+        assert len(events) == 1
+        assert events[0].message == "Gateway deleted"
+        assert events[0].targets[0].entity_name == "gateway-aws"
+
+
 class TestUpdateGateway:
     @pytest.mark.asyncio
     @pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True)