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

dstack/_internal/proxy/gateway/testing/common.py

@@ -0,0 +1,13 @@
+from dataclasses import dataclass
+from typing import Union
+from unittest.mock import AsyncMock, MagicMock
+
+AnyMock = Union[MagicMock, AsyncMock]
+
+
+@dataclass
+class Mocks:
+    reload_nginx: AnyMock
+    run_certbot: AnyMock
+    open_conn: AnyMock
+    close_conn: AnyMock

dstack/_internal/proxy/lib/auth.py

@@ -0,0 +1,7 @@
+from abc import ABC, abstractmethod
+
+
+class BaseProxyAuthProvider(ABC):
+    @abstractmethod
+    async def is_project_member(self, project_name: str, token: str) -> bool:
+        pass

dstack/_internal/proxy/lib/deps.py

@@ -0,0 +1,106 @@
+from abc import ABC, abstractmethod
+from typing import AsyncGenerator, Optional
+
+from fastapi import Depends, FastAPI, Request, Security, status
+from fastapi.security import HTTPAuthorizationCredentials, HTTPBearer
+from typing_extensions import Annotated
+
+from dstack._internal.proxy.lib.auth import BaseProxyAuthProvider
+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 ServiceConnectionPool
+
+
+class ProxyDependencyInjector(ABC):
+    """
+    An injector instance stored in FastAPI's app.state.proxy_dependency_injector
+    configures dstack-proxy to use a specific set of dependencies, e.g.
+    a specific repo implementation.
+    """
+
+    def __init__(self) -> None:
+        self._service_conn_pool = ServiceConnectionPool()
+
+    # Abstract AsyncGenerator does not need async def since
+    # type checkers infer a different type without yield in body.
+    # https://mypy.readthedocs.io/en/stable/more_types.html#asynchronous-iterators
+
+    @abstractmethod
+    def get_repo(self) -> AsyncGenerator[BaseProxyRepo, None]:
+        pass
+
+    @abstractmethod
+    def get_auth_provider(self) -> AsyncGenerator[BaseProxyAuthProvider, None]:
+        pass
+
+    async def get_service_connection_pool(self) -> ServiceConnectionPool:
+        return self._service_conn_pool
+
+
+def get_injector_from_app(app: FastAPI) -> ProxyDependencyInjector:
+    injector = app.state.proxy_dependency_injector
+    if not isinstance(injector, ProxyDependencyInjector):
+        raise UnexpectedProxyError(f"Unexpected proxy_dependency_injector type {type(injector)}")
+    return injector
+
+
+async def get_injector(request: Request) -> ProxyDependencyInjector:
+    return get_injector_from_app(request.app)
+
+
+async def get_proxy_repo(
+    injector: Annotated[ProxyDependencyInjector, Depends(get_injector)],
+) -> AsyncGenerator[BaseProxyRepo, None]:
+    async for repo in injector.get_repo():
+        yield repo
+
+
+async def get_proxy_auth_provider(
+    injector: Annotated[ProxyDependencyInjector, Depends(get_injector)],
+) -> AsyncGenerator[BaseProxyAuthProvider, None]:
+    async for provider in injector.get_auth_provider():
+        yield provider
+
+
+async def get_service_connection_pool(
+    injector: Annotated[ProxyDependencyInjector, Depends(get_injector)],
+) -> ServiceConnectionPool:
+    return await injector.get_service_connection_pool()
+
+
+class ProxyAuthContext:
+    def __init__(self, project_name: str, token: Optional[str], provider: BaseProxyAuthProvider):
+        self._project_name = project_name
+        self._token = token
+        self._provider = provider
+
+    async def enforce(self) -> None:
+        if self._token is None or not await self._provider.is_project_member(
+            self._project_name, self._token
+        ):
+            raise ProxyError(
+                f"Unauthenticated or unauthorized to access project {self._project_name}",
+                status.HTTP_403_FORBIDDEN,
+            )
+
+
+class ProxyAuth:
+    def __init__(self, auto_enforce: bool):
+        self._auto_enforce = auto_enforce
+
+    async def __call__(
+        self,
+        project_name: str,
+        token: Annotated[
+            Optional[HTTPAuthorizationCredentials], Security(HTTPBearer(auto_error=False))
+        ],
+        provider: Annotated[BaseProxyAuthProvider, Depends(get_proxy_auth_provider)],
+    ) -> ProxyAuthContext:
+        context = ProxyAuthContext(
+            project_name=project_name,
+            token=token.credentials if token is not None else None,
+            provider=provider,
+        )
+        if self._auto_enforce:
+            await context.enforce()
+        return context

dstack/_internal/proxy/lib/errors.py

@@ -0,0 +1,14 @@
+from fastapi import HTTPException, status
+
+
+class ProxyError(HTTPException):
+    """Errors in dstack-proxy that are caused by and should be reported to the user"""
+
+    def __init__(self, detail: str, code: int = status.HTTP_400_BAD_REQUEST) -> None:
+        super().__init__(detail=detail, status_code=code)
+
+
+class UnexpectedProxyError(RuntimeError):
+    """Internal errors in dstack-proxy that should have never happened"""
+
+    pass

dstack/_internal/proxy/lib/models.py

@@ -0,0 +1,112 @@
+"""Things stored in BaseProxyRepo implementations."""
+
+from datetime import datetime
+from typing import Iterable, Literal, Optional, Union
+
+from pydantic import BaseModel, Field
+from typing_extensions import Annotated
+
+from dstack._internal.core.models.instances import SSHConnectionParams
+from dstack._internal.core.models.routers import AnyRouterConfig
+from dstack._internal.proxy.lib.errors import UnexpectedProxyError
+
+
+# Models should be immutable so that they can be stored in memory and safely shared by
+# coroutines without copying on every read operation.
+class ImmutableModel(BaseModel):
+    class Config:
+        frozen = True
+
+
+class Replica(ImmutableModel):
+    id: str
+    app_port: int
+    ssh_destination: str
+    ssh_port: int
+    ssh_proxy: Optional[SSHConnectionParams]
+    # Optional outer proxy, a head node/bastion
+    ssh_head_proxy: Optional[SSHConnectionParams] = None
+    ssh_head_proxy_private_key: Optional[str] = None
+
+
+class IPAddressPartitioningKey(ImmutableModel):
+    type: Literal["ip_address"] = "ip_address"
+
+
+class HeaderPartitioningKey(ImmutableModel):
+    type: Literal["header"] = "header"
+    header: Annotated[str, Field(regex=r"^[a-zA-Z0-9-_]+$")]  # prevent Nginx config injection
+
+
+class RateLimit(ImmutableModel):
+    prefix: Annotated[str, Field(regex=r"^/[^\s\\{}]*$")]  # prevent Nginx config injection
+    key: Annotated[
+        Union[IPAddressPartitioningKey, HeaderPartitioningKey],
+        Field(discriminator="type"),
+    ]
+    rps: float
+    burst: int
+
+
+class Service(ImmutableModel):
+    project_name: str
+    run_name: str
+    domain: Optional[str]  # only used on gateways
+    https: Optional[bool]  # only used on gateways
+    rate_limits: tuple[RateLimit, ...] = ()  # only used on gateways
+    auth: bool
+    client_max_body_size: int  # only enforced on gateways
+    strip_prefix: bool = True  # only used in-server
+    replicas: tuple[Replica, ...]
+    router: Optional[AnyRouterConfig] = None
+
+    @property
+    def domain_safe(self) -> str:
+        if self.domain is None:
+            raise UnexpectedProxyError(f"domain unexpectedly missing for service {self.fmt()}")
+        return self.domain
+
+    @property
+    def https_safe(self) -> bool:
+        if self.https is None:
+            raise UnexpectedProxyError(f"https unexpectedly missing for service {self.fmt()}")
+        return self.https
+
+    def with_replicas(self, new_replicas: Iterable[Replica]) -> "Service":
+        return Service(**{**self.dict(), "replicas": tuple(new_replicas)})
+
+    def find_replica(self, replica_id: str) -> Optional[Replica]:
+        for replica in self.replicas:
+            if replica.id == replica_id:
+                return replica
+        return None
+
+    def fmt(self) -> str:
+        return f"{self.project_name}/{self.run_name}"
+
+
+class Project(ImmutableModel):
+    name: str
+    ssh_private_key: str
+
+
+class TGIChatModelFormat(ImmutableModel):
+    format: Literal["tgi"] = "tgi"
+    chat_template: str
+    eos_token: str
+
+
+class OpenAIChatModelFormat(ImmutableModel):
+    format: Literal["openai"] = "openai"
+    prefix: str
+
+
+AnyModelFormat = Union[TGIChatModelFormat, OpenAIChatModelFormat]
+
+
+class ChatModel(ImmutableModel):
+    project_name: str
+    name: str
+    created_at: datetime
+    run_name: str
+    format_spec: Annotated[AnyModelFormat, Field(discriminator="format")]

dstack/_internal/proxy/lib/repo.py

@@ -0,0 +1,27 @@
+from abc import ABC, abstractmethod
+from typing import List, Optional
+
+from dstack._internal.proxy.lib.models import ChatModel, Project, Service
+
+
+class BaseProxyRepo(ABC):
+    """
+    Data access methods relevant for both in-server and gateway environments.
+    Implementations can have additional environment-specific methods.
+    """
+
+    @abstractmethod
+    async def get_service(self, project_name: str, run_name: str) -> Optional[Service]:
+        pass
+
+    @abstractmethod
+    async def list_models(self, project_name: str) -> List[ChatModel]:
+        pass
+
+    @abstractmethod
+    async def get_model(self, project_name: str, name: str) -> Optional[ChatModel]:
+        pass
+
+    @abstractmethod
+    async def get_project(self, name: str) -> Optional[Project]:
+        pass

dstack/_internal/proxy/lib/routers/model_proxy.py

@@ -0,0 +1,102 @@
+from typing import AsyncIterator, Optional
+
+from fastapi import APIRouter, Depends, status
+from fastapi.responses import StreamingResponse
+from typing_extensions import Annotated
+
+from dstack._internal.proxy.lib.deps import ProxyAuth, get_proxy_repo, get_service_connection_pool
+from dstack._internal.proxy.lib.errors import ProxyError, UnexpectedProxyError
+from dstack._internal.proxy.lib.repo import BaseProxyRepo
+from dstack._internal.proxy.lib.schemas.model_proxy import (
+    ChatCompletionsChunk,
+    ChatCompletionsRequest,
+    ChatCompletionsResponse,
+    Model,
+    ModelsResponse,
+)
+from dstack._internal.proxy.lib.services.model_proxy.model_proxy import get_chat_client
+from dstack._internal.proxy.lib.services.service_connection import (
+    ServiceConnectionPool,
+    get_service_replica_client,
+)
+
+router = APIRouter(dependencies=[Depends(ProxyAuth(auto_enforce=True))])
+
+
+@router.get("/{project_name}/models")
+async def get_models(
+    project_name: str, repo: Annotated[BaseProxyRepo, Depends(get_proxy_repo)]
+) -> ModelsResponse:
+    models = await repo.list_models(project_name)
+    data = [
+        Model(id=m.name, created=int(m.created_at.timestamp()), owned_by=project_name)
+        for m in models
+    ]
+    return ModelsResponse(data=data)
+
+
+@router.post("/{project_name}/chat/completions", response_model=ChatCompletionsResponse)
+async def post_chat_completions(
+    project_name: str,
+    body: ChatCompletionsRequest,
+    repo: Annotated[BaseProxyRepo, Depends(get_proxy_repo)],
+    service_conn_pool: Annotated[ServiceConnectionPool, Depends(get_service_connection_pool)],
+):
+    model = await repo.get_model(project_name, body.model)
+    if model is None:
+        raise ProxyError(
+            f"Model {body.model} not found in project {project_name}", status.HTTP_404_NOT_FOUND
+        )
+    service = await repo.get_service(project_name, model.run_name)
+    if service is None or not service.replicas:
+        raise UnexpectedProxyError(
+            f"Model {model.name} in project {project_name} references run {model.run_name}"
+            " that does not exist or has no replicas"
+        )
+    http_client = await get_service_replica_client(service, repo, service_conn_pool)
+    client = get_chat_client(model, http_client)
+    if not body.stream:
+        return await client.generate(body)
+    else:
+        return StreamingResponse(
+            await StreamingAdaptor(client.stream(body)).get_stream(),
+            media_type="text/event-stream",
+            headers={"X-Accel-Buffering": "no"},
+        )
+
+
+class StreamingAdaptor:
+    """
+    Converts a stream of ChatCompletionsChunk to an SSE stream.
+    Also pre-fetches the first chunk **before** starting streaming to downstream,
+    so that upstream request errors can propagate to the downstream client.
+    """
+
+    def __init__(self, stream: AsyncIterator[ChatCompletionsChunk]) -> None:
+        self._stream = stream
+
+    async def get_stream(self) -> AsyncIterator[bytes]:
+        try:
+            first_chunk = await self._stream.__anext__()
+        except StopAsyncIteration:
+            first_chunk = None
+        return self._adaptor(first_chunk)
+
+    async def _adaptor(self, first_chunk: Optional[ChatCompletionsChunk]) -> AsyncIterator[bytes]:
+        if first_chunk is not None:
+            yield self._encode_chunk(first_chunk)
+
+            try:
+                async for chunk in self._stream:
+                    yield self._encode_chunk(chunk)
+            except ProxyError as e:
+                # No standard way to report errors while streaming,
+                # but we'll at least send them as comments
+                yield f": {e.detail!r}\n\n".encode()  # !r to avoid line breaks
+                return
+
+        yield "data: [DONE]\n\n".encode()
+
+    @staticmethod
+    def _encode_chunk(chunk: ChatCompletionsChunk) -> bytes:
+        return f"data:{chunk.json()}\n\n".encode()

dstack/_internal/proxy/lib/schemas/model_proxy.py

@@ -0,0 +1,77 @@
+from typing import Any, Dict, List, Literal, Optional, Union
+
+from dstack._internal.core.models.common import CoreModel
+
+
+class ChatMessage(CoreModel):
+    role: str  # TODO(egor-s) types
+    content: str
+
+
+class ChatCompletionsRequest(CoreModel):
+    messages: List[ChatMessage]
+    model: str
+    frequency_penalty: Optional[float] = 0.0
+    logit_bias: Dict[str, float] = {}
+    max_tokens: Optional[int] = None
+    n: int = 1
+    presence_penalty: float = 0.0
+    response_format: Optional[Dict] = None
+    seed: Optional[int] = None
+    stop: Optional[Union[str, List[str]]] = None
+    stream: bool = False
+    temperature: Optional[float] = 1.0
+    top_p: Optional[float] = 1.0
+    tools: List[Any] = []
+    tool_choice: Union[Literal["none", "auto"], Dict] = {}
+    user: Optional[str] = None
+
+
+class ChatCompletionsChoice(CoreModel):
+    finish_reason: str
+    index: int
+    message: ChatMessage
+
+
+class ChatCompletionsChunkChoice(CoreModel):
+    delta: object
+    logprobs: object = {}
+    finish_reason: Optional[str]
+    index: int
+
+
+class ChatCompletionsUsage(CoreModel):
+    completion_tokens: int
+    prompt_tokens: int
+    total_tokens: int
+
+
+class ChatCompletionsResponse(CoreModel):
+    id: str
+    choices: List[ChatCompletionsChoice]
+    created: int
+    model: str
+    system_fingerprint: str = ""
+    object: Literal["chat.completion"] = "chat.completion"
+    usage: ChatCompletionsUsage
+
+
+class ChatCompletionsChunk(CoreModel):
+    id: Optional[str] = None
+    choices: List[ChatCompletionsChunkChoice]
+    created: Optional[int] = None
+    model: str
+    system_fingerprint: Optional[str] = ""
+    object: Literal["chat.completion.chunk"] = "chat.completion.chunk"
+
+
+class Model(CoreModel):
+    object: Literal["model"] = "model"
+    id: str
+    created: int
+    owned_by: str
+
+
+class ModelsResponse(CoreModel):
+    object: Literal["list"] = "list"
+    data: List[Model]

dstack/_internal/proxy/lib/services/model_proxy/clients/base.py

@@ -0,0 +1,18 @@
+from abc import ABC, abstractmethod
+from typing import AsyncIterator
+
+from dstack._internal.proxy.lib.schemas.model_proxy import (
+    ChatCompletionsChunk,
+    ChatCompletionsRequest,
+    ChatCompletionsResponse,
+)
+
+
+class ChatCompletionsClient(ABC):
+    @abstractmethod
+    async def generate(self, request: ChatCompletionsRequest) -> ChatCompletionsResponse:
+        pass
+
+    @abstractmethod
+    async def stream(self, request: ChatCompletionsRequest) -> AsyncIterator[ChatCompletionsChunk]:
+        yield

dstack/_internal/proxy/lib/services/model_proxy/clients/openai.py

@@ -0,0 +1,67 @@
+from typing import AsyncIterator
+
+import httpx
+from fastapi import status
+from pydantic import ValidationError
+
+from dstack._internal.proxy.lib.errors import ProxyError
+from dstack._internal.proxy.lib.schemas.model_proxy import (
+    ChatCompletionsChunk,
+    ChatCompletionsRequest,
+    ChatCompletionsResponse,
+)
+from dstack._internal.proxy.lib.services.model_proxy.clients.base import ChatCompletionsClient
+
+
+class OpenAIChatCompletions(ChatCompletionsClient):
+    def __init__(self, http_client: httpx.AsyncClient, prefix: str):
+        self._http = http_client
+        self._prefix = prefix
+
+    async def generate(self, request: ChatCompletionsRequest) -> ChatCompletionsResponse:
+        try:
+            resp = await self._http.post(
+                f"{self._prefix}/chat/completions", json=request.dict(exclude_unset=True)
+            )
+            await self._propagate_error(resp)
+        except httpx.RequestError as e:
+            raise ProxyError(f"Error requesting model: {e!r}", status.HTTP_502_BAD_GATEWAY)
+
+        try:
+            return ChatCompletionsResponse.__response__.parse_raw(resp.content)
+        except ValidationError as e:
+            raise ProxyError(f"Invalid response from model: {e}", status.HTTP_502_BAD_GATEWAY)
+
+    async def stream(self, request: ChatCompletionsRequest) -> AsyncIterator[ChatCompletionsChunk]:
+        try:
+            async with self._http.stream(
+                "POST", f"{self._prefix}/chat/completions", json=request.dict(exclude_unset=True)
+            ) as resp:
+                await self._propagate_error(resp)
+
+                async for line in resp.aiter_lines():
+                    if not line.startswith("data:"):
+                        continue
+                    data = line[len("data:") :].strip()
+                    if data == "[DONE]":
+                        break
+                    yield self._parse_chunk_data(data)
+        except httpx.RequestError as e:
+            raise ProxyError(f"Error requesting model: {e!r}", status.HTTP_502_BAD_GATEWAY)
+
+    @staticmethod
+    def _parse_chunk_data(data: str) -> ChatCompletionsChunk:
+        try:
+            return ChatCompletionsChunk.__response__.parse_raw(data)
+        except ValidationError as e:
+            raise ProxyError(f"Invalid chunk in model stream: {e}", status.HTTP_502_BAD_GATEWAY)
+
+    @staticmethod
+    async def _propagate_error(resp: httpx.Response) -> None:
+        """
+        Propagates HTTP error by raising ProxyError if status is not 200.
+        May also raise httpx.RequestError if there are issues reading the response.
+        """
+        if resp.status_code != 200:
+            resp_body = await resp.aread()
+            raise ProxyError(resp_body.decode(errors="replace"), code=resp.status_code)