pytensorforge 0.1.0__py3-none-any.whl

src/serving/protocol.py

@@ -0,0 +1,377 @@
+import time
+import uuid
+from dataclasses import dataclass, field
+from typing import List, Optional, Union
+
+from src.serving.errors import bad_request
+
+CHAT_ROLES = ("system", "user", "assistant", "developer")
+
+NEUTRAL_PARAMS = {
+    "n": 1,
+    "presence_penalty": 0,
+    "frequency_penalty": 0,
+    "logprobs": False,
+    "logit_bias": None,
+    "echo": False,
+    "best_of": 1,
+    "suffix": None,
+    "parallel_tool_calls": None,
+    "top_logprobs": None,
+}
+
+IGNORED_PARAMS = ("user", "metadata", "store", "service_tier", "response_format")
+
+UNSUPPORTED_PARAMS = ("tools", "tool_choice", "functions", "function_call", "audio", "modalities", "prediction")
+
+
+@dataclass
+class SamplingParams:
+    max_tokens: Optional[int] = None
+    temperature: Optional[float] = None
+    top_p: Optional[float] = None
+    top_k: Optional[int] = None
+    repetition_penalty: Optional[float] = None
+    seed: Optional[int] = None
+    stop: List[str] = field(default_factory=list)
+
+
+@dataclass
+class ChatRequest:
+    model: str
+    messages: List[dict]
+    sampling: SamplingParams
+    stream: bool = False
+    include_usage: bool = False
+
+
+@dataclass
+class CompletionRequest:
+    model: str
+    prompt: Union[str, List[int]]
+    sampling: SamplingParams
+    stream: bool = False
+    include_usage: bool = False
+
+
+def _is_int(v):
+    return isinstance(v, int) and not isinstance(v, bool)
+
+
+def _is_num(v):
+    return (isinstance(v, (int, float))) and not isinstance(v, bool)
+
+
+def _number(body, name, lo=None, hi=None, lo_open=False, integer=False):
+    if name not in body or body[name] is None:
+        return None
+
+    v = body[name]
+
+    if integer and not _is_int(v):
+        raise bad_request(f"'{name}' must be an integer", param=name)
+
+    if not integer and not _is_num(v):
+        raise bad_request(f"'{name}' must be a number", param=name)
+
+    if lo is not None and (v <= lo if lo_open else v < lo):
+        raise bad_request(f"'{name}' must be {'>' if lo_open else '>='} {lo}", param=name)
+
+    if hi is not None and v > hi:
+        raise bad_request(f"'{name}' must be <= {hi}", param=name)
+
+    return v
+
+
+def _check_params(body, allowed):
+    if not isinstance(body, dict):
+        raise bad_request("request body must be a JSON object")
+
+    for name in UNSUPPORTED_PARAMS:
+        if body.get(name) not in (None, [], {}):
+            raise bad_request(f"'{name}' is not supported by this server", param=name, code="unsupported_parameter")
+
+    for name, neutral in NEUTRAL_PARAMS.items():
+        if name in body and body[name] is not None and body[name] != neutral:
+            raise bad_request(
+                f"'{name}'={body[name]!r} is not supported; only {neutral!r} is accepted",
+                param=name, code="unsupported_parameter",
+            )
+
+    known = set(allowed) | set(NEUTRAL_PARAMS) | set(IGNORED_PARAMS) | set(UNSUPPORTED_PARAMS)
+    unknown = sorted(k for k in body if k not in known)
+
+    if unknown:
+        raise bad_request(f"unrecognized request parameter(s): {unknown}", param=unknown[0])
+
+
+def _model(body):
+    model = body.get("model")
+
+    if not isinstance(model, str) or not model:
+        raise bad_request("'model' is required and must be a string", param="model")
+
+    return model
+
+
+def _stop(body, limits):
+    stop = body.get("stop")
+
+    if stop is None:
+        return []
+
+    if isinstance(stop, str):
+        stop = [stop]
+
+    if not isinstance(stop, list) or not all(isinstance(s, str) for s in stop):
+        raise bad_request("'stop' must be a string or a list of strings", param="stop")
+
+    stop = [s for s in stop if s]
+
+    if len(stop) > limits.max_stop_sequences:
+        raise bad_request(f"at most {limits.max_stop_sequences} stop sequences are allowed", param="stop")
+
+    if any(len(s) > limits.max_stop_sequence_chars for s in stop):
+        raise bad_request(f"stop sequences must be at most {limits.max_stop_sequence_chars} characters", param="stop")
+
+    return stop
+
+
+def _stream(body):
+    stream = body.get("stream", False)
+
+    if stream is None:
+        stream = False
+
+    if not isinstance(stream, bool):
+        raise bad_request("'stream' must be a boolean", param="stream")
+
+    opts = body.get("stream_options")
+    include_usage = False
+
+    if opts is not None:
+        if not isinstance(opts, dict) or set(opts) - {"include_usage"}:
+            raise bad_request("'stream_options' supports only 'include_usage'", param="stream_options")
+
+        if not stream:
+            raise bad_request("'stream_options' is only allowed when 'stream' is true", param="stream_options")
+
+        include_usage = bool(opts.get("include_usage", False))
+
+    return stream, include_usage
+
+
+def _sampling(body, limits, token_field_names):
+    max_tokens = None
+    source = None
+
+    for name in token_field_names:
+        v = _number(body, name, lo=1, integer=True)
+
+        if v is not None:
+            if max_tokens is not None and v != max_tokens:
+                raise bad_request(f"conflicting values for {' and '.join(token_field_names)}", param=name)
+            max_tokens = v
+            source = source or name
+
+    if max_tokens is not None and max_tokens > limits.max_generation_tokens:
+        raise bad_request(
+            f"{source}={max_tokens} exceeds this server's limit of {limits.max_generation_tokens}",
+            param=source, code="max_tokens_exceeded",
+        )
+
+    return SamplingParams(
+        max_tokens=max_tokens,
+        temperature=_number(body, "temperature", lo=0, hi=2),
+        top_p=_number(body, "top_p", lo=0, hi=1, lo_open=True),
+        top_k=_number(body, "top_k", lo=0, integer=True),
+        repetition_penalty=_number(body, "repetition_penalty", lo=0, hi=10, lo_open=True),
+        seed=_number(body, "seed", integer=True),
+        stop=_stop(body, limits),
+    )
+
+
+def _content_text(content, idx):
+    if isinstance(content, str):
+        return content
+
+    if isinstance(content, list):
+        parts = []
+
+        for part in content:
+            if not isinstance(part, dict) or part.get("type") != "text" or not isinstance(part.get("text"), str):
+                raise bad_request(
+                    f"messages[{idx}].content: only text content parts are supported", param=f"messages[{idx}].content"
+                )
+            parts.append(part["text"])
+
+        return "".join(parts)
+
+    if content is None:
+        return ""
+
+    raise bad_request(f"messages[{idx}].content must be a string or a list of text parts",
+                      param=f"messages[{idx}].content")
+
+
+CHAT_FIELDS = ("model", "messages", "stream", "stream_options", "max_tokens", "max_completion_tokens",
+               "temperature", "top_p", "top_k", "repetition_penalty", "seed", "stop")
+
+COMPLETION_FIELDS = ("model", "prompt", "stream", "stream_options", "max_tokens", "temperature", "top_p",
+                     "top_k", "repetition_penalty", "seed", "stop")
+
+
+def parse_chat_request(body, limits):
+    _check_params(body, CHAT_FIELDS)
+
+    messages = body.get("messages")
+
+    if not isinstance(messages, list) or not messages:
+        raise bad_request("'messages' must be a non-empty list", param="messages")
+
+    if len(messages) > limits.max_messages:
+        raise bad_request(f"at most {limits.max_messages} messages are allowed", param="messages")
+
+    out = []
+    total_chars = 0
+
+    for i, m in enumerate(messages):
+        if not isinstance(m, dict):
+            raise bad_request(f"messages[{i}] must be an object", param=f"messages[{i}]")
+
+        role = m.get("role")
+
+        if role not in CHAT_ROLES:
+            raise bad_request(f"messages[{i}].role must be one of {list(CHAT_ROLES)}", param=f"messages[{i}].role")
+
+        if m.get("tool_calls") or role == "tool":
+            raise bad_request("tool messages are not supported", param=f"messages[{i}]")
+
+        text = _content_text(m.get("content"), i)
+        total_chars += len(text)
+        out.append({"role": "system" if role == "developer" else role, "content": text})
+
+    if total_chars > limits.max_prompt_chars:
+        raise bad_request(f"messages exceed the {limits.max_prompt_chars}-character limit",
+                          param="messages", code="prompt_too_large")
+
+    if out[-1]["role"] == "assistant":
+        raise bad_request("the last message must not be from the assistant", param="messages")
+
+    stream, include_usage = _stream(body)
+
+    return ChatRequest(
+        model=_model(body),
+        messages=out,
+        sampling=_sampling(body, limits, ("max_tokens", "max_completion_tokens")),
+        stream=stream,
+        include_usage=include_usage,
+    )
+
+
+def parse_completion_request(body, limits):
+    _check_params(body, COMPLETION_FIELDS)
+
+    prompt = body.get("prompt")
+
+    if isinstance(prompt, list) and len(prompt) == 1 and isinstance(prompt[0], str):
+        prompt = prompt[0]
+    elif isinstance(prompt, list) and len(prompt) == 1 and isinstance(prompt[0], list):
+        prompt = prompt[0]
+
+    if isinstance(prompt, str):
+        if len(prompt) > limits.max_prompt_chars:
+            raise bad_request(f"prompt exceeds the {limits.max_prompt_chars}-character limit",
+                              param="prompt", code="prompt_too_large")
+    elif isinstance(prompt, list) and prompt and all(_is_int(t) for t in prompt):
+        if len(prompt) > limits.max_prompt_chars:
+            raise bad_request("prompt is too long", param="prompt", code="prompt_too_large")
+    elif isinstance(prompt, list):
+        raise bad_request("batched prompts are not supported; send one prompt per request", param="prompt")
+    else:
+        raise bad_request("'prompt' must be a string or a list of token ids", param="prompt")
+
+    stream, include_usage = _stream(body)
+
+    return CompletionRequest(
+        model=_model(body),
+        prompt=prompt,
+        sampling=_sampling(body, limits, ("max_tokens",)),
+        stream=stream,
+        include_usage=include_usage,
+    )
+
+
+def new_id(prefix):
+    return f"{prefix}-{uuid.uuid4().hex[:24]}"
+
+
+def usage_dict(prompt_tokens, completion_tokens):
+    return {
+        "prompt_tokens": prompt_tokens,
+        "completion_tokens": completion_tokens,
+        "total_tokens": prompt_tokens + completion_tokens,
+    }
+
+
+def chat_response(rid, model, fingerprint, text, finish_reason, usage):
+    return {
+        "id": rid,
+        "object": "chat.completion",
+        "created": int(time.time()),
+        "model": model,
+        "system_fingerprint": fingerprint,
+        "choices": [{
+            "index": 0,
+            "message": {"role": "assistant", "content": text},
+            "logprobs": None,
+            "finish_reason": finish_reason,
+        }],
+        "usage": usage,
+    }
+
+
+def chat_chunk(rid, created, model, fingerprint, delta, finish_reason=None, usage=None, include_choice=True):
+    chunk = {
+        "id": rid,
+        "object": "chat.completion.chunk",
+        "created": created,
+        "model": model,
+        "system_fingerprint": fingerprint,
+        "choices": [{"index": 0, "delta": delta, "logprobs": None, "finish_reason": finish_reason}]
+        if include_choice else [],
+    }
+
+    if usage is not None:
+        chunk["usage"] = usage
+
+    return chunk
+
+
+def completion_response(rid, model, fingerprint, text, finish_reason, usage):
+    return {
+        "id": rid,
+        "object": "text_completion",
+        "created": int(time.time()),
+        "model": model,
+        "system_fingerprint": fingerprint,
+        "choices": [{"index": 0, "text": text, "logprobs": None, "finish_reason": finish_reason}],
+        "usage": usage,
+    }
+
+
+def completion_chunk(rid, created, model, fingerprint, text, finish_reason=None, usage=None, include_choice=True):
+    chunk = {
+        "id": rid,
+        "object": "text_completion",
+        "created": created,
+        "model": model,
+        "system_fingerprint": fingerprint,
+        "choices": [{"index": 0, "text": text, "logprobs": None, "finish_reason": finish_reason}]
+        if include_choice else [],
+    }
+
+    if usage is not None:
+        chunk["usage"] = usage
+
+    return chunk

src/serving/security.py

@@ -0,0 +1,200 @@
+import hashlib
+import hmac
+import ipaddress
+import os
+import time
+from collections import OrderedDict
+from dataclasses import dataclass
+
+from src.serving import errors
+
+
+def hash_key(key):
+    return hashlib.sha256(key.encode("utf-8")).hexdigest()
+
+
+def _parse_entry(entry):
+    entry = entry.strip()
+
+    if not entry or entry.startswith("#"):
+        return None
+
+    if entry.startswith("sha256:"):
+        digest = entry[len("sha256:"):].strip().lower()
+
+        if len(digest) != 64 or any(c not in "0123456789abcdef" for c in digest):
+            raise ValueError("malformed sha256 key entry")
+
+        return digest
+
+    if len(entry) < 16:
+        raise ValueError("API keys must be at least 16 characters")
+
+    return hash_key(entry)
+
+
+def _read_key_file(path):
+    with open(path) as f:
+        return [line for line in f.read().splitlines()]
+
+
+def _collect(inline, path, env_name):
+    raw = list(inline or [])
+
+    if path:
+        raw.extend(_read_key_file(path))
+
+    if env_name and os.environ.get(env_name):
+        raw.extend(os.environ[env_name].split(","))
+
+    return [d for d in (_parse_entry(e) for e in raw) if d]
+
+
+@dataclass(frozen=True)
+class Principal:
+    id: str
+    authenticated: bool
+    admin: bool = False
+
+
+class KeyStore:
+
+    def __init__(self, user_digests, admin_digests, allow_unauthenticated=False):
+        self._user = [bytes.fromhex(d) for d in dict.fromkeys(user_digests)]
+        self._admin = [bytes.fromhex(d) for d in dict.fromkeys(admin_digests)]
+        self.allow_unauthenticated = allow_unauthenticated
+
+    @classmethod
+    def from_config(cls, sec):
+        users = _collect(sec.api_keys, sec.api_keys_file, sec.api_keys_env)
+        admins = _collect(sec.admin_keys, sec.admin_keys_file, None)
+        return cls(users, admins, sec.allow_unauthenticated)
+
+    @property
+    def enabled(self):
+        return bool(self._user or self._admin)
+
+    @staticmethod
+    def extract(headers):
+        auth = headers.get("authorization", "")
+
+        if auth[:7].lower() == "bearer ":
+            return auth[7:].strip() or None
+
+        key = headers.get("x-api-key", "").strip()
+
+        return key or None
+
+    @staticmethod
+    def _match(digest, pool):
+        found = False
+
+        for candidate in pool:
+            found |= hmac.compare_digest(digest, candidate)
+
+        return found
+
+    def authenticate(self, headers, client_ip):
+        key = self.extract(headers)
+
+        if key is None:
+            if self.enabled and not self.allow_unauthenticated:
+                raise errors.unauthorized("missing API key; send 'Authorization: Bearer <key>'")
+
+            return Principal(f"ip:{client_ip}", authenticated=False)
+
+        digest = hashlib.sha256(key.encode("utf-8")).digest()
+        is_admin = self._match(digest, self._admin)
+        is_user = self._match(digest, self._user)
+
+        if not (is_admin or is_user):
+            raise errors.unauthorized()
+
+        return Principal(f"key:{digest.hex()[:12]}", authenticated=True, admin=is_admin)
+
+
+def is_loopback(host):
+    if host in ("localhost", ""):
+        return True
+
+    try:
+        return ipaddress.ip_address(host).is_loopback
+    except ValueError:
+        return False
+
+
+class _Bucket:
+    __slots__ = ("tokens", "updated", "in_flight")
+
+    def __init__(self, capacity, now):
+        self.tokens = float(capacity)
+        self.updated = now
+        self.in_flight = 0
+
+
+class RateLimiter:
+
+    def __init__(self, requests_per_minute, burst, max_concurrent, max_tracked=10_000, clock=time.monotonic):
+        self.rate = requests_per_minute / 60.0
+        self.capacity = burst
+        self.max_concurrent = max_concurrent
+        self.max_tracked = max_tracked
+        self.clock = clock
+        self._buckets = OrderedDict()
+
+    @classmethod
+    def from_config(cls, rl):
+        return cls(rl.requests_per_minute, rl.burst, rl.max_concurrent_per_principal, rl.max_tracked_principals)
+
+    def _bucket(self, key, now):
+        b = self._buckets.get(key)
+
+        if b is None:
+            b = _Bucket(self.capacity, now)
+            self._buckets[key] = b
+
+            while len(self._buckets) > self.max_tracked:
+                oldest, victim = next(iter(self._buckets.items()))
+
+                if victim.in_flight:
+                    self._buckets.move_to_end(oldest)
+                    break
+
+                self._buckets.popitem(last=False)
+        else:
+            self._buckets.move_to_end(key)
+
+        b.tokens = min(self.capacity, b.tokens + (now - b.updated) * self.rate)
+        b.updated = now
+
+        return b
+
+    def acquire(self, key):
+        now = self.clock()
+        b = self._bucket(key, now)
+
+        if b.in_flight >= self.max_concurrent:
+            raise errors.rate_limited(
+                f"too many concurrent requests for this API key (limit {self.max_concurrent})", 1
+            )
+
+        if b.tokens < 1.0:
+            raise errors.rate_limited("request rate limit exceeded", (1.0 - b.tokens) / self.rate)
+
+        b.tokens -= 1.0
+        b.in_flight += 1
+
+    def release(self, key):
+        b = self._buckets.get(key)
+
+        if b is not None and b.in_flight > 0:
+            b.in_flight -= 1
+
+    def charge(self, key):
+        now = self.clock()
+        b = self._bucket(key, now)
+
+        if b.tokens < 1.0:
+            raise errors.rate_limited("request rate limit exceeded", (1.0 - b.tokens) / self.rate)
+
+        b.tokens -= 1.0

src/serving/server.py

@@ -0,0 +1,121 @@
+import asyncio
+import logging
+import os
+import signal
+import threading
+
+from src.serving.app import APIApp
+from src.serving.http import HTTPServer
+from src.serving.model_server import ModelServer
+from src.serving.security import is_loopback
+
+log = logging.getLogger("ptf.server")
+
+DEFAULT_UI_DIR = os.path.join(os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))),
+                              "web", "dist")
+
+
+class InsecureConfiguration(ValueError):
+    pass
+
+
+def check_security(config, keys):
+    if keys.enabled:
+        return
+
+    if config.security.allow_unauthenticated:
+        if not is_loopback(config.host):
+            log.warning("serving WITHOUT authentication on non-loopback address %s", config.host)
+        return
+
+    if not is_loopback(config.host):
+        raise InsecureConfiguration(
+            f"refusing to listen on {config.host} without API keys; configure security.api_keys / "
+            f"api_keys_file / ${config.security.api_keys_env}, or set security.allow_unauthenticated"
+        )
+
+    config.security.allow_unauthenticated = True
+    keys.allow_unauthenticated = True
+    log.warning("no API keys configured; accepting unauthenticated requests on loopback %s", config.host)
+
+
+class APIServer:
+
+    def __init__(self, config):
+        self.config = config.validate()
+        self.models = ModelServer(self.config)
+        ui_dir = self.config.ui.directory or DEFAULT_UI_DIR
+        self.app = APIApp(self.models, self.config, ui_directory=ui_dir)
+        check_security(self.config, self.app.keys)
+
+        cors = {"Server": "pytensorforge"}
+        self.http = HTTPServer(self.app, self.config.host, self.config.port, self.config.limits,
+                               on_request_done=self.app.on_request_done, extra_headers=cors)
+        self.http.on_connection_change = lambda n: self.app.metrics.open_connections.set(n)
+        self._loop = None
+        self._thread = None
+        self._stopped = None
+        self._stop_requested = None
+
+    @property
+    def port(self):
+        return self.http.bound_port
+
+    async def _run(self, ready=None, install_signals=True):
+        self._loop = asyncio.get_running_loop()
+        self._stop_requested = asyncio.Event()
+
+        await self._loop.run_in_executor(None, self.models.preload)
+        await self.http.start()
+        log.info("listening on http://%s:%d", self.config.host, self.port)
+
+        if install_signals:
+            for sig in (signal.SIGINT, signal.SIGTERM):
+                try:
+                    self._loop.add_signal_handler(sig, self._stop_requested.set)
+                except (NotImplementedError, RuntimeError):
+                    pass
+
+        if ready is not None:
+            ready.set()
+
+        await self._stop_requested.wait()
+        log.info("shutting down: draining in-flight requests (up to %.0fs)", self.config.limits.shutdown_drain_s)
+
+        await self.http.shutdown(self.config.limits.shutdown_drain_s)
+        await self._loop.run_in_executor(None, self.models.shutdown, 5.0)
+        self.app.close()
+        log.info("shutdown complete")
+
+    def serve_forever(self):
+        asyncio.run(self._run())
+
+    def start_background(self, timeout=60.0):
+        ready = threading.Event()
+        failure = []
+
+        def target():
+            try:
+                asyncio.run(self._run(ready=ready, install_signals=False))
+            except BaseException as exc:
+                failure.append(exc)
+                ready.set()
+
+        self._thread = threading.Thread(target=target, name="ptf-api", daemon=True)
+        self._thread.start()
+
+        if not ready.wait(timeout):
+            raise TimeoutError("server did not start in time")
+
+        if failure:
+            raise failure[0]
+
+        return self
+
+    def stop(self, timeout=30.0):
+        if self._loop is not None and self._stop_requested is not None:
+            self._loop.call_soon_threadsafe(self._stop_requested.set)
+
+        if self._thread is not None:
+            self._thread.join(timeout)
+            self._thread = None