From c7cbf35d3aec4087aea63785cc5ba6aa6b799ad6 Mon Sep 17 00:00:00 2001 From: Jvst Me Date: Thu, 7 May 2026 19:30:21 +0200 Subject: [PATCH] Introduce `GatewayModel.forbid_new_services` Forbidding new services can be useful for migration purposes. For now, there is no user interface for this setting, since we primarily need it for `dstack` Sky. --- ...c2_add_gatewaymodel_forbid_new_services.py | 36 +++++++++ src/dstack/_internal/server/models.py | 5 ++ .../server/services/services/__init__.py | 3 + src/dstack/_internal/server/testing/common.py | 2 + .../_internal/server/routers/test_runs.py | 81 +++++++++++++++++++ 5 files changed, 127 insertions(+) create mode 100644 src/dstack/_internal/server/migrations/versions/2026/05_07_1721_205690dfeec2_add_gatewaymodel_forbid_new_services.py diff --git a/src/dstack/_internal/server/migrations/versions/2026/05_07_1721_205690dfeec2_add_gatewaymodel_forbid_new_services.py b/src/dstack/_internal/server/migrations/versions/2026/05_07_1721_205690dfeec2_add_gatewaymodel_forbid_new_services.py new file mode 100644 index 000000000..d164fb9fd --- /dev/null +++ b/src/dstack/_internal/server/migrations/versions/2026/05_07_1721_205690dfeec2_add_gatewaymodel_forbid_new_services.py @@ -0,0 +1,36 @@ +"""Add GatewayModel.forbid_new_services + +Revision ID: 205690dfeec2 +Revises: db3679abd063 +Create Date: 2026-05-07 17:21:23.415019+00:00 + +""" + +import sqlalchemy as sa +from alembic import op + +# revision identifiers, used by Alembic. +revision = "205690dfeec2" +down_revision = "db3679abd063" +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( + "forbid_new_services", sa.Boolean(), server_default=sa.false(), nullable=False + ) + ) + + # ### 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_column("forbid_new_services") + + # ### end Alembic commands ### diff --git a/src/dstack/_internal/server/models.py b/src/dstack/_internal/server/models.py index d6a3bb940..2ab5805b6 100644 --- a/src/dstack/_internal/server/models.py +++ b/src/dstack/_internal/server/models.py @@ -614,6 +614,11 @@ class GatewayModel(PipelineModelMixin, BaseModel): status_message: Mapped[Optional[str]] = mapped_column(Text) last_processed_at: Mapped[datetime] = mapped_column(NaiveDateTime) to_be_deleted: Mapped[bool] = mapped_column(Boolean, server_default=false()) + forbid_new_services: Mapped[bool] = mapped_column(Boolean, server_default=false()) + """ + `forbid_new_services` is useful when migrating off the gateway or doing maintenance. + For now, it can only be set by server admins via an SQL query. + """ project_id: Mapped[uuid.UUID] = mapped_column(ForeignKey("projects.id", ondelete="CASCADE")) project: Mapped["ProjectModel"] = relationship(foreign_keys=[project_id]) diff --git a/src/dstack/_internal/server/services/services/__init__.py b/src/dstack/_internal/server/services/services/__init__.py index c5c0867be..955c7a486 100644 --- a/src/dstack/_internal/server/services/services/__init__.py +++ b/src/dstack/_internal/server/services/services/__init__.py @@ -105,6 +105,9 @@ async def _register_service_in_gateway( if gateway.status != GatewayStatus.RUNNING: raise ServerClientError("Gateway status is not running") + if gateway.forbid_new_services: + raise ServerClientError("Gateway does not accept new services") + gateway_configuration = get_gateway_configuration(gateway) has_replica_group_router = any( diff --git a/src/dstack/_internal/server/testing/common.py b/src/dstack/_internal/server/testing/common.py index 65d81946a..a1deb4fad 100644 --- a/src/dstack/_internal/server/testing/common.py +++ b/src/dstack/_internal/server/testing/common.py @@ -639,6 +639,7 @@ async def create_gateway( gateway_compute_id: Optional[UUID] = None, status: Optional[GatewayStatus] = GatewayStatus.SUBMITTED, last_processed_at: datetime = datetime(2023, 1, 2, 3, 4, tzinfo=timezone.utc), + forbid_new_services: bool = False, ) -> GatewayModel: gateway = GatewayModel( project_id=project_id, @@ -649,6 +650,7 @@ async def create_gateway( gateway_compute_id=gateway_compute_id, status=status, last_processed_at=last_processed_at, + forbid_new_services=forbid_new_services, ) session.add(gateway) await session.commit() diff --git a/src/tests/_internal/server/routers/test_runs.py b/src/tests/_internal/server/routers/test_runs.py index 71d265b1d..e8689805e 100644 --- a/src/tests/_internal/server/routers/test_runs.py +++ b/src/tests/_internal/server/routers/test_runs.py @@ -3617,3 +3617,84 @@ async def test_unregister_dangling_service( ) # Verify that register_service was called twice (first failed, then succeeded) assert client_mock.register_service.call_count == 2 + + @pytest.mark.asyncio + async def test_return_error_if_default_gateway_forbids_new_services( + self, + test_db, + session: AsyncSession, + client: AsyncClient, + ) -> None: + user = await create_user(session=session, global_role=GlobalRole.USER) + project = await create_project(session=session, owner=user, name="test-project") + await add_project_member( + session=session, project=project, user=user, project_role=ProjectRole.USER + ) + repo = await create_repo(session=session, project_id=project.id) + backend = await create_backend(session=session, project_id=project.id) + gateway_compute = await create_gateway_compute(session=session, backend_id=backend.id) + gateway = await create_gateway( + session=session, + project_id=project.id, + backend_id=backend.id, + gateway_compute_id=gateway_compute.id, + status=GatewayStatus.RUNNING, + wildcard_domain="example.com", + forbid_new_services=True, + ) + project.default_gateway_id = gateway.id + await session.commit() + + response = await client.post( + "/api/project/test-project/runs/submit", + headers=get_auth_headers(user.token), + json={"run_spec": get_service_run_spec(repo_id=repo.name, run_name="test-service")}, + ) + + assert response.status_code == 400 + assert response.json() == { + "detail": [{"msg": "Gateway does not accept new services", "code": "error"}] + } + + @pytest.mark.asyncio + async def test_return_error_if_explicitly_specified_gateway_forbids_new_services( + self, + test_db, + session: AsyncSession, + client: AsyncClient, + ) -> None: + user = await create_user(session=session, global_role=GlobalRole.USER) + project = await create_project(session=session, owner=user, name="test-project") + await add_project_member( + session=session, project=project, user=user, project_role=ProjectRole.USER + ) + repo = await create_repo(session=session, project_id=project.id) + backend = await create_backend(session=session, project_id=project.id) + gateway_compute = await create_gateway_compute(session=session, backend_id=backend.id) + await create_gateway( + session=session, + project_id=project.id, + backend_id=backend.id, + gateway_compute_id=gateway_compute.id, + status=GatewayStatus.RUNNING, + name="restricted-gateway", + wildcard_domain="example.com", + forbid_new_services=True, + ) + + response = await client.post( + "/api/project/test-project/runs/submit", + headers=get_auth_headers(user.token), + json={ + "run_spec": get_service_run_spec( + repo_id=repo.name, + run_name="test-service", + gateway="restricted-gateway", + ) + }, + ) + + assert response.status_code == 400 + assert response.json() == { + "detail": [{"msg": "Gateway does not accept new services", "code": "error"}] + }