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

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

@@ -0,0 +1,208 @@
+import datetime
+import json
+import uuid
+from typing import AsyncIterator, Dict, List
+
+import httpx
+import jinja2
+import jinja2.sandbox
+from fastapi import status
+
+from dstack._internal.proxy.lib.errors import ProxyError
+from dstack._internal.proxy.lib.schemas.model_proxy import (
+    ChatCompletionsChoice,
+    ChatCompletionsChunk,
+    ChatCompletionsChunkChoice,
+    ChatCompletionsRequest,
+    ChatCompletionsResponse,
+    ChatCompletionsUsage,
+    ChatMessage,
+)
+from dstack._internal.proxy.lib.services.model_proxy.clients.base import ChatCompletionsClient
+
+
+class TGIChatCompletions(ChatCompletionsClient):
+    # https://huggingface.github.io/text-generation-inference/
+    def __init__(self, http_client: httpx.AsyncClient, chat_template: str, eos_token: str):
+        self.client = http_client
+        self.eos_token = eos_token
+
+        try:
+            jinja_env = jinja2.sandbox.ImmutableSandboxedEnvironment(
+                trim_blocks=True, lstrip_blocks=True
+            )
+            jinja_env.globals["raise_exception"] = raise_exception
+            self.chat_template = jinja_env.from_string(chat_template)
+        except jinja2.TemplateError as e:
+            raise ProxyError(f"Failed to compile chat template: {e}")
+
+    async def generate(self, request: ChatCompletionsRequest) -> ChatCompletionsResponse:
+        payload = self.get_payload(request)
+        try:
+            resp = await self.client.post("/generate", json=payload)
+            await self.propagate_error(resp)
+        except httpx.RequestError as e:
+            raise ProxyError(f"Error requesting model: {e!r}", status.HTTP_502_BAD_GATEWAY)
+
+        data = resp.json()
+
+        choices = [
+            ChatCompletionsChoice(
+                finish_reason=self.finish_reason(data["details"]["finish_reason"]),
+                index=0,
+                message=ChatMessage(
+                    role="assistant",
+                    content=self.trim_stop_tokens(
+                        data["generated_text"], payload["parameters"]["stop"]
+                    ),
+                ),
+            )
+        ]
+        completion_tokens = data["details"]["generated_tokens"]
+        prompt_tokens = len(data["details"]["prefill"])
+
+        for i, sequence in enumerate(data["details"].get("best_of_sequences", []), start=1):
+            choices.append(
+                ChatCompletionsChoice(
+                    finish_reason=self.finish_reason(sequence["finish_reason"]),
+                    index=i,
+                    message=ChatMessage(
+                        role="assistant",
+                        content=self.trim_stop_tokens(
+                            sequence["generated_text"], payload["parameters"]["stop"]
+                        ),
+                    ),
+                )
+            )
+            completion_tokens += sequence["generated_tokens"]
+
+        return ChatCompletionsResponse(
+            id=uuid.uuid4().hex,
+            choices=choices,
+            created=int(datetime.datetime.utcnow().timestamp()),
+            model=request.model,
+            system_fingerprint=f"fp_{data['details']['seed']}",
+            usage=ChatCompletionsUsage(
+                completion_tokens=completion_tokens,
+                prompt_tokens=prompt_tokens,  # TODO(egor-s): do we need to multiply by number of sequences?
+                total_tokens=completion_tokens + prompt_tokens,
+            ),
+        )
+
+    async def stream(self, request: ChatCompletionsRequest) -> AsyncIterator[ChatCompletionsChunk]:
+        completion_id = uuid.uuid4().hex
+        created = int(datetime.datetime.utcnow().timestamp())
+
+        payload = self.get_payload(request)
+        try:
+            async with self.client.stream("POST", "/generate_stream", json=payload) as resp:
+                await self.propagate_error(resp)
+                async for line in resp.aiter_lines():
+                    if line.startswith("data:"):
+                        yield self.parse_chunk(
+                            data=json.loads(line[len("data:") :].strip("\n")),
+                            model=request.model,
+                            completion_id=completion_id,
+                            created=created,
+                        )
+        except httpx.RequestError as e:
+            raise ProxyError(f"Error requesting model: {e!r}", status.HTTP_502_BAD_GATEWAY)
+
+    def parse_chunk(
+        self, data: dict, model: str, completion_id: str, created: int
+    ) -> ChatCompletionsChunk:
+        if "error" in data:
+            raise ProxyError(data["error"])
+        chunk = ChatCompletionsChunk(
+            id=completion_id,
+            choices=[],
+            created=created,
+            model=model,
+            system_fingerprint="",
+        )
+        if data["details"] is not None:
+            chunk.choices = [
+                ChatCompletionsChunkChoice(
+                    delta={},
+                    logprobs=None,
+                    finish_reason=self.finish_reason(data["details"]["finish_reason"]),
+                    index=0,
+                )
+            ]
+        else:
+            chunk.choices = [
+                ChatCompletionsChunkChoice(
+                    delta={"content": data["token"]["text"], "role": "assistant"},
+                    logprobs=None,
+                    finish_reason=None,
+                    index=0,
+                )
+            ]
+        return chunk
+
+    def get_payload(self, request: ChatCompletionsRequest) -> Dict:
+        try:
+            inputs = self.chat_template.render(
+                messages=request.messages,
+                add_generation_prompt=True,
+            )
+        except jinja2.TemplateError as e:
+            raise ProxyError(f"Failed to render chat template: {e}")
+
+        stop = ([request.stop] if isinstance(request.stop, str) else request.stop) or []
+        if self.eos_token not in stop:
+            stop.append(self.eos_token)
+
+        parameters = {
+            "do_sample": True,  # activate logits sampling
+            "max_new_tokens": request.max_tokens,
+            # TODO(egor-s): OpenAI parameters do not convert to `repetition_penalty`
+            # "repetition_penalty": None,
+            # "return_full_text": False,
+            "stop": stop,
+            "seed": request.seed,
+            "temperature": request.temperature,
+            # OpenAI doesn't specify `top_k` parameter
+            # "top_k": None,
+            # "truncate": None,
+            # "typical_p": None,
+            "best_of": request.n,
+            # "watermark": False,
+            "details": True,  # to get best_of_sequences
+            "decoder_input_details": not request.stream,
+        }
+        if request.top_p < 1.0:
+            parameters["top_p"] = request.top_p
+        return {
+            "inputs": inputs,
+            "parameters": parameters,
+        }
+
+    @staticmethod
+    def finish_reason(reason: str) -> str:
+        if reason == "stop_sequence" or reason == "eos_token":
+            return "stop"
+        if reason == "length":
+            return "length"
+        raise ProxyError(f"Unknown finish reason: {reason}")
+
+    @staticmethod
+    def trim_stop_tokens(text: str, stop_tokens: List[str]) -> str:
+        for stop_token in stop_tokens:
+            if text.endswith(stop_token):
+                return text[: -len(stop_token)]
+        return text
+
+    @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)
+
+
+def raise_exception(message: str):
+    raise jinja2.TemplateError(message)

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

@@ -0,0 +1,23 @@
+import httpx
+
+from dstack._internal.proxy.lib.errors import UnexpectedProxyError
+from dstack._internal.proxy.lib.models import ChatModel
+from dstack._internal.proxy.lib.services.model_proxy.clients.base import ChatCompletionsClient
+from dstack._internal.proxy.lib.services.model_proxy.clients.openai import OpenAIChatCompletions
+from dstack._internal.proxy.lib.services.model_proxy.clients.tgi import TGIChatCompletions
+
+
+def get_chat_client(model: ChatModel, http_client: httpx.AsyncClient) -> ChatCompletionsClient:
+    if model.format_spec.format == "tgi":
+        return TGIChatCompletions(
+            http_client=http_client,
+            chat_template=model.format_spec.chat_template,
+            eos_token=model.format_spec.eos_token,
+        )
+    elif model.format_spec.format == "openai":
+        return OpenAIChatCompletions(
+            http_client=http_client,
+            prefix=model.format_spec.prefix,
+        )
+    else:
+        raise UnexpectedProxyError(f"Unsupported model format {model.format_spec.format}")

dstack/_internal/proxy/lib/services/service_connection.py

@@ -0,0 +1,160 @@
+import asyncio
+import os
+import random
+from pathlib import Path
+from tempfile import TemporaryDirectory
+from typing import Dict, Optional
+
+import httpx
+from httpx import AsyncHTTPTransport
+
+from dstack._internal.core.services.ssh.tunnel import (
+    SSH_DEFAULT_OPTIONS,
+    IPSocket,
+    SocketPair,
+    SSHTunnel,
+    UnixSocket,
+)
+from dstack._internal.proxy.lib.errors import UnexpectedProxyError
+from dstack._internal.proxy.lib.models import Project, Replica, Service
+from dstack._internal.proxy.lib.repo import BaseProxyRepo
+from dstack._internal.utils.common import get_or_error
+from dstack._internal.utils.env import environ
+from dstack._internal.utils.logging import get_logger
+from dstack._internal.utils.path import FileContent
+
+logger = get_logger(__name__)
+OPEN_TUNNEL_TIMEOUT = 10
+HTTP_TIMEOUT = environ.get_int("DSTACK_SERVICE_CLIENT_TIMEOUT", default=60)
+# Same as default Nginx proxy timeout; override via DSTACK_SERVICE_CLIENT_TIMEOUT
+
+
+class ServiceClient(httpx.AsyncClient):
+    def build_request(self, *args, **kwargs) -> httpx.Request:
+        self.cookies.clear()  # the client is shared by all users, don't leak cookies
+        return super().build_request(*args, **kwargs)
+
+
+class ServiceConnection:
+    def __init__(self, project: Project, service: Service, replica: Replica) -> None:
+        self._temp_dir = TemporaryDirectory()
+        options = {
+            **SSH_DEFAULT_OPTIONS,
+            "ConnectTimeout": str(OPEN_TUNNEL_TIMEOUT),
+            "ServerAliveInterval": "60",
+        }
+        if service.domain is not None:
+            # expose socket for Nginx
+            os.chmod(self._temp_dir.name, 0o755)
+            options["StreamLocalBindMask"] = "0111"
+        self._app_socket_path = (Path(self._temp_dir.name) / "replica.sock").absolute()
+        ssh_proxies = []
+        if replica.ssh_head_proxy is not None:
+            ssh_head_proxy_private_key = get_or_error(replica.ssh_head_proxy_private_key)
+            ssh_proxies.append((replica.ssh_head_proxy, FileContent(ssh_head_proxy_private_key)))
+        if replica.ssh_proxy is not None:
+            ssh_proxies.append((replica.ssh_proxy, None))
+        self._tunnel = SSHTunnel(
+            destination=replica.ssh_destination,
+            port=replica.ssh_port,
+            ssh_proxies=ssh_proxies,
+            identity=FileContent(project.ssh_private_key),
+            forwarded_sockets=[
+                SocketPair(
+                    remote=IPSocket("localhost", replica.app_port),
+                    local=UnixSocket(self._app_socket_path),
+                ),
+            ],
+            options=options,
+        )
+        self._client = ServiceClient(
+            transport=AsyncHTTPTransport(uds=str(self._app_socket_path)),
+            # The hostname in base_url is there for troubleshooting, as it may appear in
+            # logs and in the Host header. The actual destination is the Unix socket.
+            base_url=f"http://{replica.id}-{service.run_name}/",
+            timeout=HTTP_TIMEOUT,
+        )
+        self._is_open = asyncio.locks.Event()
+
+    @property
+    def app_socket_path(self) -> Path:
+        return self._app_socket_path
+
+    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) -> ServiceClient:
+        await asyncio.wait_for(self._is_open.wait(), timeout=OPEN_TUNNEL_TIMEOUT)
+        return self._client
+
+
+class ServiceConnectionPool:
+    def __init__(self) -> None:
+        # TODO(#2238): remove connections to stopped replicas in-server
+        self.connections: Dict[str, ServiceConnection] = {}
+
+    async def get(self, replica_id: str) -> Optional[ServiceConnection]:
+        return self.connections.get(replica_id)
+
+    async def get_or_add(
+        self, project: Project, service: Service, replica: Replica
+    ) -> ServiceConnection:
+        connection = self.connections.get(replica.id)
+        if connection is not None:
+            return connection
+        connection = ServiceConnection(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
+                )
+
+
+async def get_service_replica_client(
+    service: Service, repo: BaseProxyRepo, service_conn_pool: ServiceConnectionPool
+) -> httpx.AsyncClient:
+    """
+    `service` must have at least one replica
+    """
+    if service.domain is not None:
+        # Forward to Nginx so that requests are visible to StatsCollector in the access log
+        return httpx.AsyncClient(
+            base_url="http://127.0.0.1",
+            headers={"Host": service.domain},
+            timeout=HTTP_TIMEOUT,
+        )
+    # Nginx not available, forward directly to the tunnel
+    replica = random.choice(service.replicas)
+    connection = await service_conn_pool.get(replica.id)
+    if connection is None:
+        project = await repo.get_project(service.project_name)
+        if project is None:
+            raise UnexpectedProxyError(
+                f"Expected to find project {service.project_name} but could not"
+            )
+        connection = await service_conn_pool.get_or_add(project, service, replica)
+    return await connection.client()

dstack/_internal/proxy/lib/testing/auth.py

@@ -0,0 +1,11 @@
+from typing import Container, Optional
+
+from dstack._internal.proxy.lib.auth import BaseProxyAuthProvider
+
+
+class ProxyTestAuthProvider(BaseProxyAuthProvider):
+    def __init__(self, project_to_tokens: Optional[dict[str, Container[str]]] = None) -> None:
+        self._project_to_tokens = project_to_tokens or {}
+
+    async def is_project_member(self, project_name: str, token: str) -> bool:
+        return token in self._project_to_tokens.get(project_name, set())

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

@@ -0,0 +1,51 @@
+from typing import AsyncGenerator, Optional
+
+from dstack._internal.proxy.lib.auth import BaseProxyAuthProvider
+from dstack._internal.proxy.lib.deps import ProxyDependencyInjector
+from dstack._internal.proxy.lib.models import Project, Replica, Service
+from dstack._internal.proxy.lib.repo import BaseProxyRepo
+
+
+class ProxyTestDependencyInjector(ProxyDependencyInjector):
+    def __init__(self, repo: BaseProxyRepo, auth: BaseProxyAuthProvider) -> None:
+        super().__init__()
+        self._repo = repo
+        self._auth = auth
+
+    async def get_repo(self) -> AsyncGenerator[BaseProxyRepo, None]:
+        yield self._repo
+
+    async def get_auth_provider(self) -> AsyncGenerator[BaseProxyAuthProvider, None]:
+        yield self._auth
+
+
+def make_project(name: str) -> Project:
+    return Project(name=name, ssh_private_key="secret")
+
+
+def make_service(
+    project_name: str,
+    run_name: str,
+    domain: Optional[str] = None,
+    https: Optional[bool] = None,
+    auth: bool = False,
+    strip_prefix: bool = True,
+) -> Service:
+    return Service(
+        project_name=project_name,
+        run_name=run_name,
+        domain=domain,
+        https=https,
+        auth=auth,
+        client_max_body_size=2**20,
+        strip_prefix=strip_prefix,
+        replicas=(
+            Replica(
+                id="xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx",
+                app_port=80,
+                ssh_destination="ubuntu@server",
+                ssh_port=22,
+                ssh_proxy=None,
+            ),
+        ),
+    )

dstack/_internal/server/alembic.ini

@@ -0,0 +1,100 @@
+# A generic, single database configuration.
+
+[alembic]
+# path to migration scripts
+script_location = migrations
+
+# template used to generate migration file names; The default value is %%(rev)s_%%(slug)s
+# Uncomment the line below if you want the files to be prepended with date and time
+# see https://alembic.sqlalchemy.org/en/latest/tutorial.html#editing-the-ini-file
+# for all available tokens
+# file_template = %%(year)d_%%(month).2d_%%(day).2d_%%(hour).2d%%(minute).2d-%%(rev)s_%%(slug)s
+
+# sys.path path, will be prepended to sys.path if present.
+# defaults to the current working directory.
+prepend_sys_path = .
+
+# timezone to use when rendering the date within the migration file
+# as well as the filename.
+# If specified, requires the python-dateutil library that can be
+# installed by adding `alembic[tz]` to the pip requirements
+# string value is passed to dateutil.tz.gettz()
+# leave blank for localtime
+# timezone =
+
+# max length of characters to apply to the
+# "slug" field
+# truncate_slug_length = 40
+
+# set to 'true' to run the environment during
+# the 'revision' command, regardless of autogenerate
+# revision_environment = false
+
+# set to 'true' to allow .pyc and .pyo files without
+# a source .py file to be detected as revisions in the
+# versions/ directory
+# sourceless = false
+
+# version location specification; This defaults
+# to alembic/versions.  When using multiple version
+# directories, initial revisions must be specified with --version-path.
+# The path separator used here should be the separator specified by "path_separator" below.
+# version_locations = %(here)s/bar:%(here)s/bat:alembic/versions
+
+# path_separator; This indicates what character is used to split lists of file
+# paths, including version_locations and prepend_sys_path.
+#
+# Use os.pathsep. Default configuration used for new projects.
+path_separator = os
+
+# set to 'true' to search source files recursively
+# in each "version_locations" directory
+# new in Alembic version 1.10
+# recursive_version_locations = false
+
+# the output encoding used when revision files
+# are written from script.py.mako
+# output_encoding = utf-8
+
+; sqlalchemy.url = driver://user:pass@localhost/dbname
+
+
+[post_write_hooks]
+# post_write_hooks defines scripts or Python functions that are run
+# on newly generated revision scripts.  See the documentation for further
+# detail and examples
+
+# Logging configuration
+[loggers]
+keys = root,sqlalchemy,alembic
+
+[handlers]
+keys = console
+
+[formatters]
+keys = generic
+
+[logger_root]
+level = WARN
+handlers = console
+qualname =
+
+[logger_sqlalchemy]
+level = WARN
+handlers =
+qualname = sqlalchemy.engine
+
+[logger_alembic]
+level = INFO
+handlers =
+qualname = alembic
+
+[handler_console]
+class = StreamHandler
+args = (sys.stderr,)
+level = NOTSET
+formatter = generic
+
+[formatter_generic]
+format = %(levelname)-5.5s [%(name)s] %(message)s
+datefmt = %H:%M:%S