Gateway-in-server early prototype

requirements_dev.txt

@@ -3,6 +3,7 @@ pre-commit
 httpx>=0.23
 pytest~=7.2
 pytest-asyncio>=0.21
+pytest-httpbin==2.1.0
 freezegun>=1.2.0
 ruff==0.5.3  # Should match .pre-commit-config.yaml
 testcontainers # testcontainers<4 may not work with asyncpg

src/dstack/_internal/gateway/deps.py

@@ -0,0 +1,35 @@
+from abc import ABC, abstractmethod
+from typing import AsyncGenerator
+
+from fastapi import Depends, Request
+from typing_extensions import Annotated
+
+from dstack._internal.gateway.repos.base import BaseGatewayRepo
+
+
+class BaseGatewayDependencyInjector(ABC):
+    """
+    The gateway uses different implementations of this injector in different
+    environments: in-serer and on a remote host. An instance with the injector interface
+    stored in FastAPI's app.state.gateway_dependency_injector configures the gateway to
+    use a specific set of dependencies, e.g. a specific repo implementation.
+    """
+
+    @abstractmethod
+    async def get_repo(self) -> AsyncGenerator[BaseGatewayRepo, None]:
+        if False:
+            yield  # show type checkers this is a generator
+
+
+async def get_injector(request: Request) -> BaseGatewayDependencyInjector:
+    injector = request.app.state.gateway_dependency_injector
+    if not isinstance(injector, BaseGatewayDependencyInjector):
+        raise RuntimeError(f"Wrong BaseGatewayDependencyInjector type {type(injector)}")
+    return injector
+
+
+async def get_gateway_repo(
+    injector: Annotated[BaseGatewayDependencyInjector, Depends(get_injector)],
+) -> AsyncGenerator[BaseGatewayRepo, None]:
+    async for repo in injector.get_repo():
+        yield repo

src/dstack/_internal/gateway/repos/base.py

@@ -0,0 +1,44 @@
+from abc import ABC, abstractmethod
+from typing import List, Optional
+
+from pydantic import BaseModel
+
+from dstack._internal.core.models.instances import SSHConnectionParams
+
+
+class Replica(BaseModel):
+    id: str
+    ssh_destination: str
+    ssh_port: int
+    ssh_proxy: Optional[SSHConnectionParams]
+
+
+class Service(BaseModel):
+    id: str
+    run_name: str
+    auth: bool
+    app_port: int
+    replicas: List[Replica]
+
+
+class Project(BaseModel):
+    name: str
+    ssh_private_key: str
+
+
+class BaseGatewayRepo(ABC):
+    @abstractmethod
+    async def get_service(self, project_name: str, run_name: str) -> Optional[Service]:
+        pass
+
+    @abstractmethod
+    async def add_service(self, project_name: str, service: Service) -> None:
+        pass
+
+    @abstractmethod
+    async def get_project(self, name: str) -> Optional[Project]:
+        pass
+
+    @abstractmethod
+    async def add_project(self, project: Project) -> None:
+        pass

src/dstack/_internal/gateway/repos/memory.py

@@ -0,0 +1,21 @@
+from typing import Dict, Optional
+
+from dstack._internal.gateway.repos.base import BaseGatewayRepo, Project, Service
+
+
+class InMemoryGatewayRepo(BaseGatewayRepo):
+    def __init__(self) -> None:
+        self.services: Dict[str, Dict[str, Service]] = {}
+        self.projects: Dict[str, Project] = {}
+
+    async def get_service(self, project_name: str, run_name: str) -> Optional[Service]:
+        return self.services.get(project_name, {}).get(run_name)
+
+    async def add_service(self, project_name: str, service: Service) -> None:
+        self.services.setdefault(project_name, {})[service.run_name] = service
+
+    async def get_project(self, name: str) -> Optional[Project]:
+        return self.projects.get(name)
+
+    async def add_project(self, project: Project) -> None:
+        self.projects[project.name] = project

src/dstack/_internal/gateway/routers/service_proxy.py

@@ -0,0 +1,35 @@
+from fastapi import APIRouter, Depends, Request, status
+from fastapi.datastructures import URL
+from fastapi.responses import RedirectResponse, Response
+from typing_extensions import Annotated
+
+from dstack._internal.gateway.deps import get_gateway_repo
+from dstack._internal.gateway.repos.base import BaseGatewayRepo
+from dstack._internal.gateway.services import service_proxy
+
+REDIRECTED_HTTP_METHODS = ["GET", "POST", "PUT", "DELETE", "PATCH", "HEAD"]
+PROXIED_HTTP_METHODS = REDIRECTED_HTTP_METHODS + ["OPTIONS"]
+
+
+router = APIRouter()
+
+
+@router.api_route("/{project_name}/{run_name}", methods=REDIRECTED_HTTP_METHODS)
+async def redirect_to_service_root(request: Request) -> Response:
+    url = URL(str(request.url))
+    url = url.replace(path=url.path + "/")
+    return RedirectResponse(url, status.HTTP_308_PERMANENT_REDIRECT)
+
+
+@router.api_route("/{project_name}/{run_name}/{path:path}", methods=PROXIED_HTTP_METHODS)
+async def service_reverse_proxy(
+    project_name: str,
+    run_name: str,
+    path: str,
+    request: Request,
+    repo: Annotated[BaseGatewayRepo, Depends(get_gateway_repo)],
+) -> Response:
+    return await service_proxy.proxy(project_name, run_name, path, request, repo)
+
+
+# TODO(#1595): support websockets

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

@@ -0,0 +1,105 @@
+import asyncio
+from pathlib import Path
+from tempfile import TemporaryDirectory
+from typing import Dict, Optional
+
+from httpx import AsyncClient, AsyncHTTPTransport
+
+from dstack._internal.core.services.ssh.tunnel import (
+    SSH_DEFAULT_OPTIONS,
+    IPSocket,
+    SocketPair,
+    SSHTunnel,
+    UnixSocket,
+)
+from dstack._internal.gateway.repos.base import Project, Replica, Service
+from dstack._internal.utils.logging import get_logger
+from dstack._internal.utils.path import FileContent
+
+logger = get_logger(__name__)
+OPEN_TUNNEL_TIMEOUT = 10
+
+
+class ServiceReplicaConnection:
+    def __init__(self, project: Project, service: Service, replica: Replica) -> None:
+        self._temp_dir = TemporaryDirectory()
+        app_socket_path = (Path(self._temp_dir.name) / "replica.sock").absolute()
+        self._tunnel = SSHTunnel(
+            destination=replica.ssh_destination,
+            port=replica.ssh_port,
+            ssh_proxy=replica.ssh_proxy,
+            identity=FileContent(project.ssh_private_key),
+            forwarded_sockets=[
+                SocketPair(
+                    remote=IPSocket("localhost", service.app_port),
+                    local=UnixSocket(app_socket_path),
+                ),
+            ],
+            options={
+                **SSH_DEFAULT_OPTIONS,
+                "ConnectTimeout": str(OPEN_TUNNEL_TIMEOUT),
+            },
+        )
+        self._client = AsyncClient(
+            transport=AsyncHTTPTransport(uds=str(app_socket_path)),
+            # The hostname in base_url is normally a placeholder, it will be overwritten
+            # by proxied requests' Host header unless they don't have it (HTTP/1.0)
+            base_url="http://service/",
+        )
+        self._is_open = asyncio.locks.Event()
+
+    async def open(self) -> None:
+        await self._tunnel.aopen()
+        self._is_open.set()
+
+    async def close(self) -> None:
+        self._is_open.clear()
+        await self._client.aclose()
+        await self._tunnel.aclose()
+
+    async def client(self) -> AsyncClient:
+        await asyncio.wait_for(self._is_open.wait(), timeout=OPEN_TUNNEL_TIMEOUT)
+        return self._client
+
+
+class ServiceReplicaConnectionPool:
+    def __init__(self) -> None:
+        # TODO(#1595): remove connections to stopped replicas
+        self.connections: Dict[str, ServiceReplicaConnection] = {}
+
+    async def get(self, replica_id: str) -> Optional[ServiceReplicaConnection]:
+        return self.connections.get(replica_id)
+
+    async def add(
+        self, project: Project, service: Service, replica: Replica
+    ) -> ServiceReplicaConnection:
+        connection = self.connections.get(replica.id)
+        if connection is not None:
+            return connection
+        connection = ServiceReplicaConnection(project, service, replica)
+        self.connections[replica.id] = connection
+        try:
+            await connection.open()
+        except BaseException:
+            self.connections.pop(replica.id, None)
+            raise
+        return connection
+
+    async def remove(self, replica_id: str) -> None:
+        connection = self.connections.pop(replica_id, None)
+        if connection is not None:
+            await connection.close()
+
+    async def remove_all(self) -> None:
+        replica_ids = list(self.connections)
+        results = await asyncio.gather(
+            *(self.remove(replica_id) for replica_id in replica_ids), return_exceptions=True
+        )
+        for i, exc in enumerate(results):
+            if isinstance(exc, Exception):
+                logger.error(
+                    "Error removing connection to service replica %s: %s", replica_ids[i], exc
+                )
+
+
+service_replica_connection_pool: ServiceReplicaConnectionPool = ServiceReplicaConnectionPool()