Allow specifying gateway in service configurations (#1972)

jvstme · web-flow · commit fc5286d33a89 · 2024-11-07T10:35:53.000Z
diff --git a/src/dstack/_internal/core/models/configurations.py b/src/dstack/_internal/core/models/configurations.py
@@ -209,6 +209,15 @@ class ServiceConfigurationParams(CoreModel):
         Union[ValidPort, constr(regex=r"^[0-9]+:[0-9]+$"), PortMapping],
         Field(description="The port, that application listens on or the mapping"),
     ]
+    gateway: Annotated[
+        Optional[Union[bool, str]],
+        Field(
+            description=(
+                "The name of the gateway. Specify boolean `false` to run without a gateway."
+                " Omit to run with the default gateway"
+            ),
+        ),
+    ] = None
     model: Annotated[
         Optional[Union[AnyModel, str]],
         Field(
@@ -267,6 +276,16 @@ def convert_replicas(cls, v: Any) -> Range[int]:
             )
         return v
 
+    @validator("gateway")
+    def validate_gateway(
+        cls, v: Optional[Union[bool, str]]
+    ) -> Optional[Union[Literal[False], str]]:
+        if v == True:
+            raise ValueError(
+                "The `gateway` property must be a string or boolean `false`, not boolean `true`"
+            )
+        return v
+
     @root_validator()
     def validate_scaling(cls, values):
         scaling = values.get("scaling")
diff --git a/src/dstack/_internal/server/services/gateways/__init__.py b/src/dstack/_internal/server/services/gateways/__init__.py
@@ -28,6 +28,7 @@
     SSHError,
 )
 from dstack._internal.core.models.backends.base import BackendType
+from dstack._internal.core.models.common import is_core_model_instance
 from dstack._internal.core.models.configurations import SERVICE_HTTPS_DEFAULT, ServiceConfiguration
 from dstack._internal.core.models.gateways import (
     Gateway,
@@ -354,27 +355,38 @@ async def generate_gateway_name(session: AsyncSession, project: ProjectModel) ->
             return name
 
 
-async def register_service(session: AsyncSession, run_model: RunModel):
-    gateway = run_model.project.default_gateway
+async def register_service(session: AsyncSession, run_model: RunModel, run_spec: RunSpec):
+    assert is_core_model_instance(run_spec.configuration, ServiceConfiguration)
+
+    if isinstance(run_spec.configuration.gateway, str):
+        gateway = await get_project_gateway_model_by_name(
+            session=session, project=run_model.project, name=run_spec.configuration.gateway
+        )
+        if gateway is None:
+            raise ResourceNotExistsError(
+                f"Gateway {run_spec.configuration.gateway} does not exist"
+            )
+    elif run_spec.configuration.gateway == False:
+        gateway = None
+    else:
+        gateway = run_model.project.default_gateway
 
     if gateway is not None:
-        service_spec = await _register_service_in_gateway(session, run_model, gateway)
+        service_spec = await _register_service_in_gateway(session, run_model, run_spec, gateway)
         run_model.gateway = gateway
     elif not settings.FORBID_SERVICES_WITHOUT_GATEWAY:
-        service_spec = _register_service_in_server(run_model)
+        service_spec = _register_service_in_server(run_model, run_spec)
     else:
         raise ResourceNotExistsError(
             "This dstack-server installation forbids services without a gateway."
-            " Please configure a default gateway."
+            " Please configure a gateway."
         )
     run_model.service_spec = service_spec.json()
 
 
 async def _register_service_in_gateway(
-    session: AsyncSession, run_model: RunModel, gateway: GatewayModel
+    session: AsyncSession, run_model: RunModel, run_spec: RunSpec, gateway: GatewayModel
 ) -> ServiceSpec:
-    run_spec: RunSpec = RunSpec.__response__.parse_raw(run_model.run_spec)
-
     if gateway.gateway_compute is None:
         raise ServerClientError("Gateway has no instance associated with it")
 
@@ -427,8 +439,7 @@ async def _register_service_in_gateway(
     return service_spec
 
 
-def _register_service_in_server(run_model: RunModel) -> ServiceSpec:
-    run_spec: RunSpec = RunSpec.__response__.parse_raw(run_model.run_spec)
+def _register_service_in_server(run_model: RunModel, run_spec: RunSpec) -> ServiceSpec:
     if run_spec.configuration.https != SERVICE_HTTPS_DEFAULT:
         # Note: if the user sets `https: <default-value>`, it will be ignored silently
         # TODO: in 0.19, make `https` Optional to be able to tell if it was set or omitted
diff --git a/src/dstack/_internal/server/services/runs.py b/src/dstack/_internal/server/services/runs.py
@@ -431,7 +431,7 @@ async def submit_run(
         replicas = 1
         if run_spec.configuration.type == "service":
             replicas = run_spec.configuration.replicas.min
-            await gateways.register_service(session, run_model)
+            await gateways.register_service(session, run_model, run_spec)
 
         for replica_num in range(replicas):
             jobs = await get_jobs_from_run_spec(run_spec, replica_num=replica_num)
diff --git a/src/tests/_internal/server/routers/test_runs.py b/src/tests/_internal/server/routers/test_runs.py
@@ -1,8 +1,8 @@
 import copy
 import json
 from datetime import datetime, timezone
-from typing import Dict, List, Optional
-from unittest.mock import Mock, patch
+from typing import Dict, Generator, List, Optional, Tuple, Union
+from unittest.mock import AsyncMock, Mock, patch
 from uuid import UUID
 
 import pytest
@@ -14,6 +14,7 @@
 from dstack._internal.core.models.backends.base import BackendType
 from dstack._internal.core.models.common import ApplyAction
 from dstack._internal.core.models.configurations import ServiceConfiguration
+from dstack._internal.core.models.gateways import GatewayStatus
 from dstack._internal.core.models.instances import (
     InstanceAvailability,
     InstanceOfferWithAvailability,
@@ -43,6 +44,9 @@
 from dstack._internal.server.services.projects import add_project_member
 from dstack._internal.server.services.runs import run_model_to_run
 from dstack._internal.server.testing.common import (
+    create_backend,
+    create_gateway,
+    create_gateway_compute,
     create_job,
     create_project,
     create_repo,
@@ -358,6 +362,32 @@ def get_dev_env_run_dict(
     }
 
 
+def get_service_run_spec(
+    repo_id: str,
+    run_name: Optional[str] = None,
+    gateway: Optional[Union[bool, str]] = None,
+) -> dict:
+    return {
+        "configuration": {
+            "type": "service",
+            "commands": ["python -m http.server"],
+            "port": 8000,
+            "gateway": gateway,
+            "model": "test-model",
+        },
+        "configuration_path": "dstack.yaml",
+        "profile": {
+            "name": "string",
+        },
+        "repo_code_hash": None,
+        "repo_data": {"repo_dir": "/repo", "repo_type": "local"},
+        "repo_id": repo_id,
+        "run_name": run_name,
+        "ssh_key_pub": "ssh_key",
+        "working_dir": ".",
+    }
+
+
 class TestListRuns:
     @pytest.mark.asyncio
     @pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True)
@@ -1481,3 +1511,144 @@ async def test_backend_does_not_support_create_instance(
                 ]
             }
             assert result == expected
+
+
+@pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True)
+class TestSubmitService:
+    @pytest.fixture(autouse=True)
+    def mock_gateway_connections(self) -> Generator[None, None, None]:
+        with patch(
+            "dstack._internal.server.services.gateways.gateway_connections_pool.get_or_add"
+        ) as get_conn_mock:
+            get_conn_mock.return_value.client = Mock()
+            get_conn_mock.return_value.client.return_value = AsyncMock()
+            yield
+
+    @pytest.mark.asyncio
+    @pytest.mark.parametrize(
+        (
+            "existing_gateways",
+            "specified_gateway_in_run_conf",
+            "expected_service_url",
+            "expected_model_url",
+        ),
+        [
+            pytest.param(
+                [("default-gateway", True), ("non-default-gateway", False)],
+                None,
+                "https://test-service.default-gateway.example",
+                "https://gateway.default-gateway.example",
+                id="submits-to-default-gateway",
+            ),
+            pytest.param(
+                [("default-gateway", True), ("non-default-gateway", False)],
+                "non-default-gateway",
+                "https://test-service.non-default-gateway.example",
+                "https://gateway.non-default-gateway.example",
+                id="submits-to-specified-gateway",
+            ),
+            pytest.param(
+                [("non-default-gateway", False)],
+                None,
+                "/proxy/services/test-project/test-service/",
+                "/proxy/models/test-project/",
+                id="submits-in-server-when-no-default-gateway",
+            ),
+            pytest.param(
+                [("default-gateway", True)],
+                False,
+                "/proxy/services/test-project/test-service/",
+                "/proxy/models/test-project/",
+                id="submits-in-server-when-specified",
+            ),
+        ],
+    )
+    async def test_submit_to_correct_proxy(
+        self,
+        test_db,
+        session: AsyncSession,
+        client: AsyncClient,
+        existing_gateways: List[Tuple[str, bool]],
+        specified_gateway_in_run_conf: str,
+        expected_service_url: str,
+        expected_model_url: str,
+    ) -> 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)
+        for gateway_name, is_default in existing_gateways:
+            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,
+                name=gateway_name,
+                wildcard_domain=f"{gateway_name}.example",
+            )
+            if is_default:
+                project.default_gateway_id = gateway.id
+                await session.commit()
+        run_spec = get_service_run_spec(
+            repo_id=repo.name,
+            run_name="test-service",
+            gateway=specified_gateway_in_run_conf,
+        )
+        response = await client.post(
+            f"/api/project/{project.name}/runs/submit",
+            headers=get_auth_headers(user.token),
+            json={"run_spec": run_spec},
+        )
+        assert response.status_code == 200
+        assert response.json()["service"]["url"] == expected_service_url
+        assert response.json()["service"]["model"]["base_url"] == expected_model_url
+
+    @pytest.mark.asyncio
+    async def test_return_error_if_specified_gateway_not_exists(
+        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)
+        await add_project_member(
+            session=session, project=project, user=user, project_role=ProjectRole.USER
+        )
+        repo = await create_repo(session=session, project_id=project.id)
+        run_spec = get_service_run_spec(repo_id=repo.name, gateway="nonexistent")
+        response = await client.post(
+            f"/api/project/{project.name}/runs/submit",
+            headers=get_auth_headers(user.token),
+            json={"run_spec": run_spec},
+        )
+        assert response.status_code == 400
+        assert response.json() == {
+            "detail": [
+                {"msg": "Gateway nonexistent does not exist", "code": "resource_not_exists"}
+            ]
+        }
+
+    @pytest.mark.asyncio
+    async def test_return_error_if_specified_gateway_is_true(
+        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)
+        await add_project_member(
+            session=session, project=project, user=user, project_role=ProjectRole.USER
+        )
+        repo = await create_repo(session=session, project_id=project.id)
+        run_spec = get_service_run_spec(repo_id=repo.name, gateway=True)
+        response = await client.post(
+            f"/api/project/{project.name}/runs/submit",
+            headers=get_auth_headers(user.token),
+            json={"run_spec": run_spec},
+        )
+        assert response.status_code == 422
+        assert "must be a string or boolean `false`, not boolean `true`" in response.text
