Prepare gateway API for state sync mechanism

src/dstack/_internal/proxy/gateway/app.py

@@ -22,6 +22,7 @@
 from dstack._internal.proxy.gateway.routers.auth import router as auth_router
 from dstack._internal.proxy.gateway.routers.config import router as config_router
 from dstack._internal.proxy.gateway.routers.registry import router as registry_router
+from dstack._internal.proxy.gateway.routers.services import router as services_router
 from dstack._internal.proxy.gateway.routers.stats import router as stats_router
 from dstack._internal.proxy.gateway.services.nginx import Nginx
 from dstack._internal.proxy.gateway.services.registry import ACCESS_LOG_PATH, apply_all
@@ -80,6 +81,7 @@ def make_app(repo: Optional[GatewayProxyRepo] = None, nginx: Optional[Nginx] = N
     app.include_router(config_router, prefix="/api/config")
     app.include_router(model_proxy_router, prefix="/api/models")
     app.include_router(registry_router, prefix="/api/registry")
+    app.include_router(services_router, prefix="/api/services")
     app.include_router(stats_router, prefix="/api/stats")
 
     @app.get("/")

src/dstack/_internal/proxy/gateway/routers/registry.py

@@ -9,6 +9,7 @@
     RegisterEntrypointRequest,
     RegisterReplicaRequest,
     RegisterServiceRequest,
+    SetServiceIdRequest,
 )
 from dstack._internal.proxy.gateway.services.nginx import Nginx
 from dstack._internal.proxy.lib.deps import get_service_connection_pool
@@ -27,6 +28,7 @@ async def register_service(
 ) -> OkResponse:
     await registry_services.register_service(
         project_name=project_name.lower(),
+        run_id=body.id,
         run_name=body.run_name.lower(),
         domain=body.domain.lower(),
         https=body.https,
@@ -62,6 +64,23 @@ async def unregister_service(
     return OkResponse()
 
 
+@router.post("/services/{run_name}/set_id")
+async def set_service_id(
+    project_name: str,
+    run_name: str,
+    body: SetServiceIdRequest,
+    repo: Annotated[GatewayProxyRepo, Depends(get_gateway_proxy_repo)],
+) -> OkResponse:
+    """Populate a missing ID for a service registered before 0.21.0"""
+    await registry_services.set_service_id(
+        project_name=project_name.lower(),
+        run_name=run_name.lower(),
+        run_id=body.id,
+        repo=repo,
+    )
+    return OkResponse()
+
+
 @router.post("/services/{run_name}/replicas/register")
 async def register_replica(
     project_name: str,

src/dstack/_internal/proxy/gateway/routers/services.py

@@ -0,0 +1,17 @@
+from typing import Annotated
+
+from fastapi import APIRouter, Depends
+
+from dstack._internal.proxy.gateway.deps import get_gateway_proxy_repo
+from dstack._internal.proxy.gateway.repo.repo import GatewayProxyRepo
+from dstack._internal.proxy.gateway.schemas.services import ServiceListResponse
+from dstack._internal.proxy.gateway.services.services import list_services
+
+router = APIRouter()
+
+
+@router.get("/list")
+async def list_all_services(
+    repo: Annotated[GatewayProxyRepo, Depends(get_gateway_proxy_repo)],
+) -> ServiceListResponse:
+    return await list_services(repo)

src/dstack/_internal/proxy/gateway/schemas/registry.py

@@ -37,6 +37,8 @@ class Options(BaseModel):
 
 
 class RegisterServiceRequest(BaseModel):
+    id: Optional[str] = None
+    """Only optional for compatibility with pre-0.21.0 callers"""
     run_name: str
     domain: str
     https: bool
@@ -49,6 +51,10 @@ class RegisterServiceRequest(BaseModel):
     router: Optional[AnyServiceRouterConfig] = None
 
 
+class SetServiceIdRequest(BaseModel):
+    id: str
+
+
 class RegisterReplicaRequest(BaseModel):
     job_id: str
     app_port: int

src/dstack/_internal/proxy/gateway/schemas/services.py

@@ -0,0 +1,19 @@
+from dstack._internal.core.models.common import CoreModel
+
+
+class ServiceListReplicaItem(CoreModel):
+    id: str
+
+
+class ServiceListItem(CoreModel):
+    """The model is minimal to allow for frequent polling by the server"""
+
+    id: str | None
+    """Can temporarily be `None` for services registered before 0.21.0"""
+    project_name: str
+    run_name: str
+    replicas: list[ServiceListReplicaItem]
+
+
+class ServiceListResponse(CoreModel):
+    services: list[ServiceListItem]

src/dstack/_internal/proxy/gateway/services/registry.py

@@ -36,6 +36,7 @@
 
 async def register_service(
     project_name: str,
+    run_id: Optional[str],
     run_name: str,
     domain: str,
     https: bool,
@@ -52,6 +53,7 @@ async def register_service(
 ) -> None:
     cors_enabled = model is not None and model.type == "chat" and model.format == "openai"
     service = models.Service(
+        id=run_id,
         project_name=project_name,
         run_name=run_name,
         domain=domain,
@@ -131,6 +133,25 @@ async def unregister_service(
     logger.info("Service %s is unregistered now", service.fmt())
 
 
+async def set_service_id(
+    project_name: str,
+    run_name: str,
+    run_id: str,
+    repo: GatewayProxyRepo,
+) -> None:
+    async with lock:
+        service = await repo.get_service(project_name, run_name)
+        if service is None:
+            raise ProxyError(f"Service {project_name}/{run_name} does not exist, cannot set ID")
+        if service.id is not None:
+            raise ProxyError(f"Service {project_name}/{run_name} already has an ID")
+
+        service = service.with_id(run_id)
+        await repo.set_service(service)
+
+    logger.info("Service %s id is set to %s", service.fmt(), run_id)
+
+
 async def register_replica(
     project_name: str,
     run_name: str,

src/dstack/_internal/proxy/gateway/services/services.py

@@ -0,0 +1,21 @@
+from dstack._internal.proxy.gateway.repo.repo import GatewayProxyRepo
+from dstack._internal.proxy.gateway.schemas.services import (
+    ServiceListItem,
+    ServiceListReplicaItem,
+    ServiceListResponse,
+)
+
+
+async def list_services(repo: GatewayProxyRepo) -> ServiceListResponse:
+    services = await repo.list_services()
+    return ServiceListResponse(
+        services=[
+            ServiceListItem(
+                id=service.id,
+                project_name=service.project_name,
+                run_name=service.run_name,
+                replicas=[ServiceListReplicaItem(id=replica.id) for replica in service.replicas],
+            )
+            for service in services
+        ]
+    )

src/dstack/_internal/proxy/lib/models.py

@@ -51,6 +51,8 @@ class RateLimit(ImmutableModel):
 
 
 class Service(ImmutableModel):
+    id: Optional[str] = None
+    """Can temporarily be `None` for services registered before 0.21.0"""
     project_name: str
     run_name: str
     domain: Optional[str] = None  # only used on gateways
@@ -79,6 +81,9 @@ def https_safe(self) -> bool:
     def with_replicas(self, new_replicas: Iterable[Replica]) -> "Service":
         return Service(**{**self.model_dump(), "replicas": tuple(new_replicas)})
 
+    def with_id(self, new_id: str) -> "Service":
+        return Service(**{**self.model_dump(), "id": new_id})
+
     def find_replica(self, replica_id: str) -> Optional[Replica]:
         for replica in self.replicas:
             if replica.id == replica_id:

src/dstack/_internal/proxy/lib/testing/common.py

@@ -1,3 +1,4 @@
+import uuid
 from typing import AsyncGenerator, Optional
 
 from dstack._internal.proxy.lib.auth import BaseProxyAuthProvider
@@ -30,8 +31,12 @@ def make_service(
     https: Optional[bool] = None,
     auth: bool = False,
     strip_prefix: bool = True,
+    run_id: Optional[str] = None,
 ) -> Service:
+    if run_id is None:
+        run_id = uuid.uuid4().hex
     return Service(
+        id=run_id,
         project_name=project_name,
         run_name=run_name,
         domain=domain,

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

@@ -7,10 +7,12 @@
 
 from dstack._internal.core.consts import DSTACK_RUNNER_SSH_PORT
 from dstack._internal.core.errors import GatewayError
+from dstack._internal.core.models.common import validate_json_extra_ignore
 from dstack._internal.core.models.configurations import RateLimit
 from dstack._internal.core.models.instances import SSHConnectionParams
 from dstack._internal.core.models.routers import AnyServiceRouterConfig
 from dstack._internal.core.models.runs import JobSpec, JobSubmission, Run, get_service_port
+from dstack._internal.proxy.gateway.schemas.services import ServiceListItem, ServiceListResponse
 from dstack._internal.proxy.gateway.schemas.stats import ServiceStats
 from dstack._internal.server import settings
 
@@ -37,6 +39,7 @@ def __init__(self, uds: Optional[str] = None, port: Optional[int] = None):
     async def register_service(
         self,
         project: str,
+        run_id: uuid.UUID,
         run_name: str,
         domain: str,
         service_https: bool,
@@ -54,6 +57,7 @@ async def register_service(
             await self.register_openai_entrypoint(project, entrypoint, gateway_https)
 
         payload = {
+            "id": run_id.hex,
             "run_name": run_name,
             "domain": domain,
             "https": service_https,
@@ -150,6 +154,16 @@ async def unregister_replica(self, project: str, run_name: str, job_id: uuid.UUI
         resp.raise_for_status()
         self.is_server_ready = True
 
+    async def set_service_id(self, project: str, run_name: str, run_id: uuid.UUID) -> None:
+        resp = await self._client.post(
+            self._url(f"/api/registry/{project}/services/{run_name}/set_id"),
+            json={"id": run_id.hex},
+        )
+        if resp.status_code == 400:
+            raise gateway_error(resp.json())
+        resp.raise_for_status()
+        self.is_server_ready = True
+
     async def register_openai_entrypoint(self, project: str, domain: str, https: bool):
         resp = await self._client.post(
             self._url(f"/api/registry/{project}/entrypoints/register"),
@@ -163,6 +177,15 @@ async def register_openai_entrypoint(self, project: str, domain: str, https: boo
         resp.raise_for_status()
         self.is_server_ready = True
 
+    async def list_services(self) -> list[ServiceListItem]:
+        resp = await self._client.get(self._url("/api/services/list"))
+        if resp.status_code == 400:
+            raise gateway_error(resp.json())
+        resp.raise_for_status()
+        resp_parsed = validate_json_extra_ignore(ServiceListResponse, resp.content)
+        self.is_server_ready = True
+        return resp_parsed.services
+
     async def submit_gateway_config(self) -> None:
         resp = await self._client.post(
             self._url("/api/config"),

src/dstack/_internal/server/services/proxy/repo.py

@@ -131,6 +131,7 @@ async def get_service(self, project_name: str, run_name: str) -> Optional[Servic
             )
             replicas.append(replica)
         return Service(
+            id=run.id.hex,
             project_name=project_name,
             run_name=run.run_name,
             domain=None,

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

@@ -187,6 +187,7 @@ async def _register_service_in_gateway(
                 do_register = partial(
                     client.register_service,
                     project=run_model.project.name,
+                    run_id=run_model.id,
                     run_name=run_model.run_name,
                     domain=domain,
                     service_https=configure_service_https,

src/tests/_internal/proxy/gateway/routers/test_registry.py

@@ -1,4 +1,5 @@
 import re
+import uuid
 from datetime import datetime
 from pathlib import Path
 from typing import Optional
@@ -13,6 +14,7 @@
 from dstack._internal.proxy.gateway.services.nginx import Nginx
 from dstack._internal.proxy.gateway.testing.common import Mocks
 from dstack._internal.proxy.lib.models import ChatModel, OpenAIChatModelFormat
+from dstack._internal.proxy.lib.testing.common import make_project, make_service
 
 
 def make_client(
@@ -32,6 +34,7 @@ def register_service_payload(
     rate_limits: Optional[list[dict]] = None,
 ) -> dict:
     return {
+        "id": uuid.uuid4().hex,
         "run_name": run_name,
         "domain": domain,
         "https": https,
@@ -115,6 +118,17 @@ async def test_register(self, tmp_path: Path, system_mocks: Mocks) -> None:
         assert "upstream" not in conf
         assert "return 503;" in conf
 
+    async def test_legacy_register_without_id(self, tmp_path: Path, system_mocks: Mocks) -> None:
+        repo = GatewayProxyRepo()
+        client = make_client(tmp_path, repo=repo)
+        payload = register_service_payload(run_name="test-run", domain="test-run.gtw.test")
+        del payload["id"]
+        resp = await client.post("/api/registry/test-proj/services/register", json=payload)
+        assert resp.status_code == 200
+        service = await repo.get_service("test-proj", "test-run")
+        assert service is not None
+        assert service.id is None
+
     async def test_register_with_https(self, tmp_path: Path, system_mocks: Mocks) -> None:
         client = make_client(tmp_path)
         resp = await client.post(
@@ -377,6 +391,55 @@ async def test_register_connection_error(self, tmp_path: Path, system_mocks: Moc
         assert conf_after == conf_before
 
 
+@pytest.mark.asyncio
+class TestSetServiceId:
+    async def test_set_id(self, tmp_path: Path, system_mocks: Mocks) -> None:
+        repo = GatewayProxyRepo()
+        client = make_client(tmp_path, repo=repo)
+        # simulate a service registered before IDs were introduced
+        await repo.set_project(make_project("test-proj"))
+        await repo.set_service(
+            make_service("test-proj", "test-run", domain="test-run.gtw.test").model_copy(
+                update={"id": None}
+            )
+        )
+        new_id = uuid.uuid4().hex
+        resp = await client.post(
+            "/api/registry/test-proj/services/test-run/set_id",
+            json={"id": new_id},
+        )
+        assert resp.status_code == 200
+        assert resp.json() == {"status": "ok"}
+        service = await repo.get_service("test-proj", "test-run")
+        assert service is not None
+        assert service.id == new_id
+
+    async def test_set_id_no_service_error(self, tmp_path: Path, system_mocks: Mocks) -> None:
+        client = make_client(tmp_path)
+        resp = await client.post(
+            "/api/registry/test-proj/services/test-run/set_id",
+            json={"id": uuid.uuid4().hex},
+        )
+        assert resp.status_code == 400
+        assert resp.json() == {
+            "detail": "Service test-proj/test-run does not exist, cannot set ID"
+        }
+
+    async def test_set_id_already_set_error(self, tmp_path: Path, system_mocks: Mocks) -> None:
+        client = make_client(tmp_path)
+        resp = await client.post(
+            "/api/registry/test-proj/services/register",
+            json=register_service_payload(run_name="test-run", domain="test-run.gtw.test"),
+        )
+        assert resp.status_code == 200
+        resp = await client.post(
+            "/api/registry/test-proj/services/test-run/set_id",
+            json={"id": uuid.uuid4().hex},
+        )
+        assert resp.status_code == 400
+        assert resp.json() == {"detail": "Service test-proj/test-run already has an ID"}
+
+
 @pytest.mark.asyncio
 class TestUnregisterService:
     async def test_unregister(self, tmp_path: Path, system_mocks: Mocks) -> None: