dstack 0.0.9__py3-none-any.whl → 0.20.7__py3-none-any.whl

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

@@ -0,0 +1,426 @@
+import asyncio
+from asyncio import Lock
+from datetime import datetime
+from pathlib import Path
+from typing import Iterable, Optional
+
+import dstack._internal.proxy.gateway.schemas.registry as schemas
+from dstack._internal.core.models.instances import SSHConnectionParams
+from dstack._internal.core.models.routers import AnyRouterConfig, RouterType
+from dstack._internal.proxy.gateway import models as gateway_models
+from dstack._internal.proxy.gateway.repo.repo import GatewayProxyRepo
+from dstack._internal.proxy.gateway.services.nginx import (
+    LimitReqConfig,
+    LimitReqZoneConfig,
+    LocationConfig,
+    ModelEntrypointConfig,
+    Nginx,
+    ReplicaConfig,
+    ServiceConfig,
+)
+from dstack._internal.proxy.lib import models
+from dstack._internal.proxy.lib.errors import ProxyError, UnexpectedProxyError
+from dstack._internal.proxy.lib.repo import BaseProxyRepo
+from dstack._internal.proxy.lib.services.service_connection import (
+    ServiceConnection,
+    ServiceConnectionPool,
+)
+from dstack._internal.utils.logging import get_logger
+
+ACCESS_LOG_PATH = Path("/var/log/nginx/dstack.access.log")
+logger = get_logger(__name__)
+lock = Lock()
+
+
+async def register_service(
+    project_name: str,
+    run_name: str,
+    domain: str,
+    https: bool,
+    rate_limits: tuple[models.RateLimit, ...],
+    auth: bool,
+    client_max_body_size: int,
+    model: Optional[schemas.AnyModel],
+    ssh_private_key: str,
+    repo: GatewayProxyRepo,
+    nginx: Nginx,
+    service_conn_pool: ServiceConnectionPool,
+    router: Optional[AnyRouterConfig] = None,
+) -> None:
+    service = models.Service(
+        project_name=project_name,
+        run_name=run_name,
+        domain=domain,
+        https=https,
+        rate_limits=rate_limits,
+        auth=auth,
+        client_max_body_size=client_max_body_size,
+        replicas=(),
+        router=router,
+    )
+
+    async with lock:
+        if await repo.get_service(project_name, run_name) is not None:
+            raise ProxyError(f"Service {service.fmt()} is already registered")
+
+        old_project = await repo.get_project(project_name)
+        new_project = models.Project(name=project_name, ssh_private_key=ssh_private_key)
+        if old_project is not None and old_project.ssh_private_key != new_project.ssh_private_key:
+            logger.warning(
+                "SSH key for service %s is different from the previous one", service.fmt()
+            )
+        await repo.set_project(new_project)
+
+        logger.debug("Registering service %s", service.fmt())
+
+        await apply_service(
+            service=service,
+            old_service=None,
+            repo=repo,
+            nginx=nginx,
+            service_conn_pool=service_conn_pool,
+        )
+        await repo.set_service(service)
+
+        if model is not None:
+            await repo.set_model(
+                models.ChatModel(
+                    project_name=project_name,
+                    name=model.name,
+                    created_at=datetime.now(),
+                    run_name=run_name,
+                    format_spec=model_schema_to_format_spec(model),
+                ),
+            )
+
+    logger.info("Service %s is registered now", service.fmt())
+
+
+async def unregister_service(
+    project_name: str,
+    run_name: str,
+    repo: GatewayProxyRepo,
+    nginx: Nginx,
+    service_conn_pool: ServiceConnectionPool,
+) -> 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} is not registered, cannot unregister"
+            )
+
+        logger.debug("Unregistering service %s", service.fmt())
+
+        await stop_replica_connections(
+            ids=(r.id for r in service.replicas),
+            service_conn_pool=service_conn_pool,
+        )
+        await nginx.unregister(service.domain_safe)
+        await repo.delete_models_by_run(project_name, run_name)
+        await repo.delete_service(project_name, run_name)
+
+    logger.info("Service %s is unregistered now", service.fmt())
+
+
+async def register_replica(
+    project_name: str,
+    run_name: str,
+    replica_id: str,
+    app_port: int,
+    ssh_destination: str,
+    ssh_port: int,
+    ssh_proxy: Optional[SSHConnectionParams],
+    ssh_head_proxy: Optional[SSHConnectionParams],
+    ssh_head_proxy_private_key: Optional[str],
+    repo: GatewayProxyRepo,
+    nginx: Nginx,
+    service_conn_pool: ServiceConnectionPool,
+) -> None:
+    replica = models.Replica(
+        id=replica_id,
+        app_port=app_port,
+        ssh_destination=ssh_destination,
+        ssh_port=ssh_port,
+        ssh_proxy=ssh_proxy,
+        ssh_head_proxy=ssh_head_proxy,
+        ssh_head_proxy_private_key=ssh_head_proxy_private_key,
+    )
+
+    async with lock:
+        old_service = await repo.get_service(project_name, run_name)
+        if old_service is None:
+            raise ProxyError(
+                f"Service {project_name}/{run_name} does not exist, cannot register replica"
+            )
+
+        if old_service.find_replica(replica_id) is not None:
+            # NOTE: as of 0.19.25, the dstack server relies on the exact text of this error.
+            # See dstack._internal.server.services.services.register_replica
+            raise ProxyError(f"Replica {replica_id} already exists in service {old_service.fmt()}")
+
+        service = old_service.with_replicas(old_service.replicas + (replica,))
+
+        logger.debug("Registering replica %s in service %s", replica.id, service.fmt())
+        failures = await apply_service(
+            service=service,
+            old_service=old_service,
+            repo=repo,
+            nginx=nginx,
+            service_conn_pool=service_conn_pool,
+        )
+        if replica in failures:
+            raise ProxyError(
+                f"Cannot register replica {replica.id}"
+                f" in service {service.fmt()}: {failures[replica]}"
+            )
+        await repo.set_service(service)
+
+    logger.info("Replica %s in service %s is registered now", replica.id, service.fmt())
+
+
+async def unregister_replica(
+    project_name: str,
+    run_name: str,
+    replica_id: str,
+    repo: GatewayProxyRepo,
+    nginx: Nginx,
+    service_conn_pool: ServiceConnectionPool,
+) -> None:
+    async with lock:
+        old_service = await repo.get_service(project_name, run_name)
+        if old_service is None:
+            raise ProxyError(
+                f"Service {project_name}/{run_name} does not exist, cannot unregister replica"
+            )
+
+        replica = old_service.find_replica(replica_id)
+        if replica is None:
+            raise ProxyError(
+                f"Replica {replica_id} does not exist in service {old_service.fmt()},"
+                " cannot unregister"
+            )
+
+        service = old_service.with_replicas(tuple(r for r in old_service.replicas if r != replica))
+
+        logger.debug("Unregistering replica %s in service %s", replica.id, service.fmt())
+
+        await apply_service(
+            service=service,
+            old_service=old_service,
+            repo=repo,
+            nginx=nginx,
+            service_conn_pool=service_conn_pool,
+        )
+        await repo.set_service(service)
+
+    logger.info("Replica %s in service %s is unregistered now", replica_id, service.fmt())
+
+
+async def register_model_entrypoint(
+    project_name: str,
+    domain: str,
+    https: bool,
+    repo: GatewayProxyRepo,
+    nginx: Nginx,
+) -> None:
+    entrypoint = gateway_models.ModelEntrypoint(
+        project_name=project_name,
+        domain=domain,
+        https=https,
+    )
+    logger.debug("Registering entrypoint %s in project %s", domain, project_name)
+    await apply_entrypoint(entrypoint, repo, nginx)
+    await repo.set_entrypoint(entrypoint)
+    logger.info("Entrypoint %s is now registered in project %s", domain, project_name)
+
+
+async def apply_service(
+    service: models.Service,
+    old_service: Optional[models.Service],
+    repo: GatewayProxyRepo,
+    nginx: Nginx,
+    service_conn_pool: ServiceConnectionPool,
+) -> dict[models.Replica, BaseException]:
+    if old_service is not None:
+        if service.domain != old_service.domain:
+            raise UnexpectedProxyError(
+                f"Did not expect service {service.fmt()}"
+                f" domain name to change ({old_service.domain} -> {service.domain})"
+            )
+        await stop_replica_connections(
+            ids=(
+                replica.id for replica in old_service.replicas if replica not in service.replicas
+            ),
+            service_conn_pool=service_conn_pool,
+        )
+    replica_conns, replica_failures = await get_or_add_replica_connections(
+        service, repo, service_conn_pool
+    )
+    replica_configs = [
+        ReplicaConfig(id=replica.id, socket=conn.app_socket_path)
+        for replica, conn in replica_conns.items()
+    ]
+    service_config = await get_nginx_service_config(service, replica_configs)
+    await nginx.register(service_config, (await repo.get_config()).acme_settings)
+    return replica_failures
+
+
+async def get_or_add_replica_connections(
+    service: models.Service, repo: BaseProxyRepo, service_conn_pool: ServiceConnectionPool
+) -> tuple[dict[models.Replica, ServiceConnection], dict[models.Replica, BaseException]]:
+    project = await repo.get_project(service.project_name)
+    if project is None:
+        raise UnexpectedProxyError(
+            f"Project {service.project_name} unexpectedly missing, even though service"
+            f" {service.fmt()} exists."
+        )
+    replica_conns, replica_failures = {}, {}
+    tasks = [
+        service_conn_pool.get_or_add(project, service, replica) for replica in service.replicas
+    ]
+    results = await asyncio.gather(*tasks, return_exceptions=True)
+    for replica, conn_or_err in zip(service.replicas, results):
+        if isinstance(conn_or_err, BaseException):
+            replica_failures[replica] = conn_or_err
+            logger.warning(
+                "Failed starting connection to replica %s in service %s: %s",
+                replica.id,
+                service.fmt(),
+                conn_or_err,
+            )
+        else:
+            replica_conns[replica] = conn_or_err
+    return replica_conns, replica_failures
+
+
+async def stop_replica_connections(
+    ids: Iterable[str], service_conn_pool: ServiceConnectionPool
+) -> None:
+    tasks = map(service_conn_pool.remove, ids)
+    results = await asyncio.gather(*tasks, return_exceptions=True)
+    for replica_id, exc in zip(ids, results):
+        if isinstance(exc, Exception):
+            logger.error("Error stopping connection to replica %s: %s", replica_id, exc)
+
+
+async def get_nginx_service_config(
+    service: models.Service, replicas: Iterable[ReplicaConfig]
+) -> ServiceConfig:
+    limit_req_zones: list[LimitReqZoneConfig] = []
+    locations: list[LocationConfig] = []
+    is_sglang = service.router and service.router.type == RouterType.SGLANG
+    sglang_whitelisted_paths = [
+        "/generate",
+        "/v1/",
+        "/chat/completions",
+    ]  # Prefix match for paths that end with a slash and exact match for paths that don't
+    sglang_limits: dict[str, LimitReqConfig] = {}
+    sglang_prefix_lengths: dict[str, int] = {}  # Track prefix lengths for most-specific selection
+
+    for i, rate_limit in enumerate(service.rate_limits):
+        zone_name = f"{i}.{service.domain_safe}"
+        if isinstance(rate_limit.key, models.IPAddressPartitioningKey):
+            key = "$binary_remote_addr"
+        elif isinstance(rate_limit.key, models.HeaderPartitioningKey):
+            key = f"$http_{rate_limit.key.header.lower().replace('-', '_')}"
+        else:
+            raise TypeError(f"Unexpected key type {type(rate_limit.key)}")
+        limit_req_zones.append(
+            LimitReqZoneConfig(name=zone_name, key=key, rpm=round(rate_limit.rps * 60))
+        )
+        if is_sglang:
+            for path in sglang_whitelisted_paths:
+                if rate_limit.prefix == path or path.startswith(rate_limit.prefix):
+                    # Use the longest prefix if multiple prefixes match the same path
+                    current_prefix_len = len(rate_limit.prefix)
+                    if path not in sglang_limits or current_prefix_len > sglang_prefix_lengths.get(
+                        path, 0
+                    ):
+                        sglang_limits[path] = LimitReqConfig(
+                            zone=zone_name, burst=rate_limit.burst
+                        )
+                        sglang_prefix_lengths[path] = current_prefix_len
+        else:
+            locations.append(
+                LocationConfig(
+                    prefix=rate_limit.prefix,
+                    limit_req=LimitReqConfig(zone=zone_name, burst=rate_limit.burst),
+                )
+            )
+
+    # Add SGLang whitelisted paths as locations
+    if is_sglang:
+        for path in sglang_whitelisted_paths:
+            # Use prefix match for paths that end with a slash and exact match for paths that don't
+            if path.endswith("/"):
+                locations.append(LocationConfig(prefix=path, limit_req=sglang_limits.get(path)))
+            else:
+                locations.append(
+                    LocationConfig(prefix=f"= {path}", limit_req=sglang_limits.get(path))
+                )
+
+    # Don't auto-add / location for SGLang routers (catch-all 403 handles it)
+    if not any(location.prefix == "/" for location in locations) and not is_sglang:
+        locations.append(LocationConfig(prefix="/", limit_req=None))
+    return ServiceConfig(
+        domain=service.domain_safe,
+        https=service.https_safe,
+        project_name=service.project_name,
+        auth=service.auth,
+        client_max_body_size=service.client_max_body_size,
+        access_log_path=ACCESS_LOG_PATH,
+        limit_req_zones=limit_req_zones,
+        locations=locations,
+        replicas=sorted(replicas, key=lambda r: r.id),  # sort for reproducible configs
+        router=service.router,
+    )
+
+
+async def apply_entrypoint(
+    entrypoint: gateway_models.ModelEntrypoint, repo: GatewayProxyRepo, nginx: Nginx
+) -> None:
+    config = ModelEntrypointConfig(
+        domain=entrypoint.domain,
+        https=entrypoint.https,
+        project_name=entrypoint.project_name,
+    )
+    acme = (await repo.get_config()).acme_settings
+    await nginx.register(config, acme)
+
+
+async def apply_all(
+    repo: GatewayProxyRepo, nginx: Nginx, service_conn_pool: ServiceConnectionPool
+) -> None:
+    service_tasks = [
+        apply_service(
+            service=service,
+            old_service=None,
+            repo=repo,
+            nginx=nginx,
+            service_conn_pool=service_conn_pool,
+        )
+        for service in await repo.list_services()
+    ]
+    entrypoint_tasks = [
+        apply_entrypoint(entrypoint, repo, nginx) for entrypoint in await repo.list_entrypoints()
+    ]
+    results = await asyncio.gather(*service_tasks, *entrypoint_tasks, return_exceptions=True)
+    for exc in results:
+        if isinstance(exc, Exception):
+            logger.error("Exception restoring gateway: %s", exc)
+
+
+def model_schema_to_format_spec(model: schemas.AnyModel) -> models.AnyModelFormat:
+    if model.type == "chat":
+        if model.format == "openai":
+            return models.OpenAIChatModelFormat(prefix=model.prefix)
+        elif model.format == "tgi":
+            return models.TGIChatModelFormat(
+                chat_template=model.chat_template,
+                eos_token=model.eos_token,
+            )
+        else:
+            raise UnexpectedProxyError(f"Unexpected model format {model.format}")
+    else:
+        raise UnexpectedProxyError(f"Unexpected model type {model.type}")

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

@@ -0,0 +1,95 @@
+import datetime
+import logging
+import random
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import Container, Dict, Generator, List
+
+import httpx
+
+logger = logging.getLogger(__name__)
+BASE_URL = "http://dstack/"  # any hostname will work
+
+
+@dataclass
+class CachedClientInfo:
+    client: httpx.AsyncClient
+    socket: Path
+    connect_errors: List[datetime.datetime] = field(default_factory=lambda: [])
+
+    def seems_disconnected(self) -> bool:
+        if len(self.connect_errors) < 2:
+            return False
+        return self.connect_errors[-1] - self.connect_errors[0] >= datetime.timedelta(minutes=2)
+
+
+class HTTPMultiClient(httpx.AsyncClient):
+    """
+    An HTTP client that sends requests to randomly chosen Unix sockets from a specified
+    directory. This allows to balance the load between multiple HTTP server replicas.
+    Automatically deletes sockets that stop responding.
+    Used for requesting random dstack-server replicas from the gateway.
+    """
+
+    def __init__(self, sockets_dir: Path):
+        super().__init__(base_url=BASE_URL)
+        self._sockets_dir = sockets_dir.expanduser()
+        self._clients_cache: Dict[str, CachedClientInfo] = {}
+
+    async def send(self, request: httpx.Request, *args, **kwargs) -> httpx.Response:
+        errors: List[httpx.RequestError] = []
+        clients_count = 0
+
+        for clients_count, client in enumerate(self._iter_clients_rand(), start=1):
+            try:
+                resp = await client.client.send(request, *args, **kwargs)
+                client.connect_errors = []
+                return resp
+            except httpx.ConnectError:
+                client.connect_errors.append(datetime.datetime.now())
+                if client.seems_disconnected():
+                    logging.debug(
+                        "Removing socket %s after several failed connection attempts",
+                        client.socket,
+                    )
+                    client.socket.unlink()
+            except httpx.RequestError as e:
+                errors.append(e)
+                logger.warning("Request failed with socket %s: %r", client.socket, e)
+
+        msg = f"Cannot request {request.url.path}: "
+        if not clients_count:
+            msg += f"no sockets found in {self._sockets_dir}"
+        elif not errors:
+            msg += f"all {clients_count} socket(s) in {self._sockets_dir} are disconnected"
+        else:
+            msg += f"{len(errors)} socket(s) failed. Last error: {errors[-1]!r}"
+        raise httpx.RequestError(msg, request=request)
+
+    def _iter_clients_rand(self) -> Generator[CachedClientInfo, None, None]:
+        sockets = list(self._sockets_dir.glob("*.sock"))
+        self._evict_clients(stems_to_keep={s.stem for s in sockets})
+        random.shuffle(sockets)
+
+        for socket in sockets:
+            if socket.stem in self._clients_cache:
+                cached_client = self._clients_cache[socket.stem]
+            else:
+                cached_client = self._clients_cache[socket.stem] = self._make_client(socket)
+            yield cached_client
+
+    @staticmethod
+    def _make_client(socket: Path) -> CachedClientInfo:
+        client = httpx.AsyncClient(
+            transport=httpx.AsyncHTTPTransport(uds=str(socket.absolute())),
+            base_url=BASE_URL,
+        )
+        return CachedClientInfo(
+            client=client,
+            socket=socket,
+        )
+
+    def _evict_clients(self, stems_to_keep: Container[str]) -> None:
+        self._clients_cache = {
+            stem: client for stem, client in self._clients_cache.items() if stem in stems_to_keep
+        }

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

@@ -0,0 +1,170 @@
+import asyncio
+import datetime
+import logging
+import os
+from collections import deque
+from collections.abc import Reversible
+from pathlib import Path
+from typing import Iterable, Optional, TextIO
+
+from pydantic import BaseModel
+
+from dstack._internal.proxy.gateway.repo.repo import GatewayProxyRepo
+from dstack._internal.proxy.gateway.schemas.stats import PerWindowStats, ServiceStats, Stat
+from dstack._internal.proxy.lib.errors import UnexpectedProxyError
+from dstack._internal.utils.common import run_async
+
+logger = logging.getLogger(__name__)
+WINDOWS = (30, 60, 300)
+TTL = WINDOWS[-1]
+EMPTY_STATS = {window: Stat(requests=0, request_time=0.0) for window in WINDOWS}
+
+
+class StatFrame(BaseModel):
+    """Service metrics aggregated over a 1s frame"""
+
+    timestamp: int
+    requests: int
+    requests_time_total: float
+
+
+class LogEntry(BaseModel):
+    """A line from the access log"""
+
+    timestamp: datetime.datetime
+    host: str
+    status: int
+    request_time: float
+    is_replica_hit: bool
+
+
+class StatsCollector:
+    """
+    StatCollector parses nginx access log and calculates average request time and requests count.
+    """
+
+    def __init__(self, access_log: Path) -> None:
+        self._path = access_log
+        self._file: Optional[TextIO] = None
+        self._stats: dict[str, deque[StatFrame]] = {}
+        self._lock = asyncio.Lock()
+
+    async def collect(self) -> dict[str, PerWindowStats]:
+        """
+        :return: stats per host aggregated by 30s, 1m, 5m
+        """
+        result = {}
+        async with self._lock:
+            await run_async(self._collect)
+            now = datetime.datetime.now(tz=datetime.timezone.utc)
+            for host, frames in self._stats.items():
+                result[host] = self._aggregate(frames, now)
+        return result
+
+    @staticmethod
+    def _aggregate(frames: Reversible[StatFrame], now: datetime.datetime) -> PerWindowStats:
+        """
+        Aggregate 1s `frames` into windows 30s, 1m, 5m before `now`
+        """
+        result = {}
+        for window in WINDOWS:
+            req_count = 0
+            req_time_total = 0.0
+            for frame in reversed(frames):
+                if now.timestamp() - frame.timestamp > window:
+                    break
+                req_time_total += frame.requests_time_total
+                req_count += frame.requests
+            if req_count > 0:
+                result[window] = Stat(
+                    requests=req_count,
+                    request_time=round(req_time_total / req_count, 3),
+                )
+            else:
+                result[window] = Stat(requests=0, request_time=0.0)
+        return result
+
+    def _collect(self) -> None:
+        now = datetime.datetime.now(tz=datetime.timezone.utc)
+
+        for entry in self._read_access_log(now - datetime.timedelta(seconds=TTL)):
+            # only include requests that hit or should hit a service replica
+            if not entry.is_replica_hit:
+                continue
+
+            frame_timestamp = int(entry.timestamp.timestamp())
+            frames = self._stats.setdefault(entry.host, deque(maxlen=TTL))
+
+            # presume that log entries are sorted by timestamp
+            if not frames or frames[-1].timestamp != frame_timestamp:
+                latest_frame = StatFrame(
+                    timestamp=frame_timestamp, requests=1, requests_time_total=entry.request_time
+                )
+                frames.append(latest_frame)
+            else:
+                latest_frame = frames[-1]
+                latest_frame.requests += 1
+                latest_frame.requests_time_total += entry.request_time
+
+        for host in list(self._stats.keys()):
+            if self._stats[host][-1].timestamp < now.timestamp() - TTL:
+                del self._stats[host]
+
+    def _read_access_log(self, after: datetime.datetime) -> Iterable[LogEntry]:
+        try:
+            st_ino = os.stat(self._path).st_ino
+        except FileNotFoundError:
+            st_ino = None
+
+        if self._file is not None:
+            while True:
+                line = self._file.readline()
+                if not line:
+                    break
+                cells = line.split()
+                if len(cells) == 4:  # compatibility with pre-0.19.11 logs
+                    cells.append("0" if cells[2] in ["403", "404"] else "1")
+                timestamp_str, host, status, request_time, dstack_replica_hit = cells
+                timestamp = datetime.datetime.fromisoformat(timestamp_str)
+                if timestamp < after:
+                    continue
+                yield LogEntry(
+                    timestamp=timestamp,
+                    host=host,
+                    status=int(status),
+                    request_time=float(request_time),
+                    is_replica_hit=_parse_nginx_bool(dstack_replica_hit),
+                )
+            if os.fstat(self._file.fileno()).st_ino != st_ino:
+                # file was rotated
+                self._file.close()
+                self._file = None
+
+        if self._file is None and st_ino is not None:
+            logger.info("Opening access log file: %s", self._path)
+            self._file = open(self._path, "r")
+            # normally, recursion will not exceed depth of 2
+            yield from self._read_access_log(after)
+
+
+async def get_service_stats(
+    repo: GatewayProxyRepo, collector: StatsCollector
+) -> list[ServiceStats]:
+    stats_per_host = await collector.collect()
+    services = await repo.list_services()
+    return [
+        ServiceStats(
+            project_name=service.project_name,
+            run_name=service.run_name,
+            stats=stats_per_host.get(service.domain_safe, EMPTY_STATS),
+        )
+        for service in services
+    ]
+
+
+def _parse_nginx_bool(v: str) -> bool:
+    if v == "0":
+        return False
+    if v == "1":
+        return True
+    raise UnexpectedProxyError(f"Cannot parse boolean value: expected '0' or '1', got {v!r}")