diff --git a/src/dstack/_internal/proxy/gateway/app.py b/src/dstack/_internal/proxy/gateway/app.py index f43ebada0a..6627e9e3f3 100644 --- a/src/dstack/_internal/proxy/gateway/app.py +++ b/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("/") diff --git a/src/dstack/_internal/proxy/gateway/routers/registry.py b/src/dstack/_internal/proxy/gateway/routers/registry.py index 2a5f23f7a0..5bf69675c8 100644 --- a/src/dstack/_internal/proxy/gateway/routers/registry.py +++ b/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, diff --git a/src/dstack/_internal/proxy/gateway/routers/services.py b/src/dstack/_internal/proxy/gateway/routers/services.py new file mode 100644 index 0000000000..b1ad11fcdf --- /dev/null +++ b/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) diff --git a/src/dstack/_internal/proxy/gateway/schemas/registry.py b/src/dstack/_internal/proxy/gateway/schemas/registry.py index beb6242f05..89e60c65cf 100644 --- a/src/dstack/_internal/proxy/gateway/schemas/registry.py +++ b/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 diff --git a/src/dstack/_internal/proxy/gateway/schemas/services.py b/src/dstack/_internal/proxy/gateway/schemas/services.py new file mode 100644 index 0000000000..cbba05b518 --- /dev/null +++ b/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] diff --git a/src/dstack/_internal/proxy/gateway/services/registry.py b/src/dstack/_internal/proxy/gateway/services/registry.py index 3b7968d62b..3f296e6ca5 100644 --- a/src/dstack/_internal/proxy/gateway/services/registry.py +++ b/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, diff --git a/src/dstack/_internal/proxy/gateway/services/services.py b/src/dstack/_internal/proxy/gateway/services/services.py new file mode 100644 index 0000000000..5543d92160 --- /dev/null +++ b/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 + ] + ) diff --git a/src/dstack/_internal/proxy/lib/models.py b/src/dstack/_internal/proxy/lib/models.py index 16e9fd4d6e..df025e6232 100644 --- a/src/dstack/_internal/proxy/lib/models.py +++ b/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: diff --git a/src/dstack/_internal/proxy/lib/testing/common.py b/src/dstack/_internal/proxy/lib/testing/common.py index 7cd5722e31..5481dcfe8b 100644 --- a/src/dstack/_internal/proxy/lib/testing/common.py +++ b/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, diff --git a/src/dstack/_internal/server/services/gateways/client.py b/src/dstack/_internal/server/services/gateways/client.py index 7dabe87536..01bebea1aa 100644 --- a/src/dstack/_internal/server/services/gateways/client.py +++ b/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"), diff --git a/src/dstack/_internal/server/services/proxy/repo.py b/src/dstack/_internal/server/services/proxy/repo.py index 3a3ab14069..b1986522cf 100644 --- a/src/dstack/_internal/server/services/proxy/repo.py +++ b/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, diff --git a/src/dstack/_internal/server/services/services/__init__.py b/src/dstack/_internal/server/services/services/__init__.py index df3985b674..0e2738832b 100644 --- a/src/dstack/_internal/server/services/services/__init__.py +++ b/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, diff --git a/src/tests/_internal/proxy/gateway/routers/test_registry.py b/src/tests/_internal/proxy/gateway/routers/test_registry.py index 239413cfda..b7a5b2e2e8 100644 --- a/src/tests/_internal/proxy/gateway/routers/test_registry.py +++ b/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: diff --git a/src/tests/_internal/proxy/gateway/routers/test_services.py b/src/tests/_internal/proxy/gateway/routers/test_services.py new file mode 100644 index 0000000000..90c9238308 --- /dev/null +++ b/src/tests/_internal/proxy/gateway/routers/test_services.py @@ -0,0 +1,57 @@ +import uuid + +import httpx +import pytest + +from dstack._internal.proxy.gateway.app import make_app +from dstack._internal.proxy.gateway.repo.repo import GatewayProxyRepo +from dstack._internal.proxy.lib.models import Replica +from dstack._internal.proxy.lib.testing.common import make_project, make_service + + +def make_client(repo: GatewayProxyRepo) -> httpx.AsyncClient: + app = make_app(repo) + return httpx.AsyncClient(transport=httpx.ASGITransport(app=app), base_url="http://test/") + + +@pytest.mark.asyncio +class TestListServices: + async def test_empty(self): + repo = GatewayProxyRepo() + client = make_client(repo) + resp = await client.get("/api/services/list") + assert resp.status_code == 200 + assert resp.json() == {"services": []} + + async def test_list(self): + repo = GatewayProxyRepo() + service_id = uuid.uuid4().hex + replica_id = uuid.uuid4().hex + service = make_service( + "test-proj", "srv-1", domain="srv-1.gtw.test", run_id=service_id + ).with_replicas( + [ + Replica( + id=replica_id, + app_port=80, + ssh_destination="ubuntu@server", + ssh_port=22, + ssh_proxy=None, + ) + ] + ) + await repo.set_project(make_project("test-proj")) + await repo.set_service(service) + client = make_client(repo) + resp = await client.get("/api/services/list") + assert resp.status_code == 200 + assert resp.json() == { + "services": [ + { + "id": service_id, + "project_name": "test-proj", + "run_name": "srv-1", + "replicas": [{"id": replica_id}], + } + ] + }