@tuned-tensor/cli 0.13.1 → 0.15.0

package/training/adapter/src/serve.py

@@ -1,280 +1,146 @@
+"""Artifact preparation only; vLLM owns HTTP, batching, streaming and tools."""
 from __future__ import annotations
 
-import hmac
 import json
-import math
+import hmac
 import os
+import signal
+import socket
 import sys
-import threading
-import time
-import uuid
-from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
 from pathlib import Path
 from tempfile import TemporaryDirectory
-from typing import Any
-
-from evaluate import (
-    configure_hugging_face_cache,
-    import_runtime_dependencies,
-    load_text_model,
-    resolve_adapter_path,
-    sampling_kwargs,
-)
-from model_contract import assert_certified_base_model_revision, chat_template_kwargs
-
-
-MODEL_ARTIFACT = os.environ.get("TT_MODEL_ARTIFACT")
-BASE_MODEL = os.environ["TT_BASE_MODEL"]
-MODEL_SOURCE = os.environ.get("TT_MODEL_SOURCE", BASE_MODEL)
-BASE_MODEL_REVISION = os.environ.get("TT_BASE_MODEL_REVISION")
-MODEL_NAME = os.environ.get("TT_MODEL_NAME", "tuned-tensor-local")
-MODEL_LOADER = os.environ.get("TT_MODEL_LOADER", "causal_lm")
-SYSTEM_PROMPT = os.environ.get("TT_SYSTEM_PROMPT", "").strip()
-HOST = os.environ.get("TT_HOST", "127.0.0.1")
-PORT = int(os.environ.get("TT_PORT", "8000"))
-DEVICE_REQUEST = os.environ.get("TT_DEVICE", "cuda")
-DEFAULT_MAX_TOKENS = int(os.environ.get("TT_MAX_TOKENS", "512"))
-DEFAULT_TEMPERATURE = float(os.environ.get("TT_TEMPERATURE", "0"))
-DEFAULT_TOP_P = float(os.environ.get("TT_TOP_P", "1"))
-MAX_REQUEST_BYTES = 2 * 1024 * 1024
-MAX_PROMPT_CHARS = 100_000
-MAX_PROMPT_TOKENS = 16_384
-MAX_MESSAGES = 128
-API_KEY = os.environ.get("TT_API_KEY", "")
-MAX_CONCURRENT_REQUESTS = int(os.environ.get("TT_MAX_CONCURRENT_REQUESTS", "1"))
-
-
-if MODEL_LOADER not in ("causal_lm", "image_text_to_text"):
-    raise ValueError(
-        f"The bundled model server does not support TT_MODEL_LOADER={MODEL_LOADER!r}; "
-        "expected causal_lm or image_text_to_text"
-    )
-assert_certified_base_model_revision(BASE_MODEL, BASE_MODEL_REVISION, "Serving base model revision")
-configure_hugging_face_cache(os.environ.get("HF_HOME"))
-import_runtime_dependencies()
-TEMP_DIR = TemporaryDirectory(prefix="tt-local-serve-")
-ADAPTER_PATH = (
-    resolve_adapter_path(MODEL_ARTIFACT, Path(TEMP_DIR.name))
-    if MODEL_ARTIFACT
-    else None
-)
-MODEL_PAYLOAD = {
-    "base_model": BASE_MODEL,
-    "model_source": MODEL_SOURCE,
-    "base_model_revision": BASE_MODEL_REVISION,
-    "device": DEVICE_REQUEST,
-    "model_loader": MODEL_LOADER,
-}
-MODEL, TOKENIZER, DEVICE = load_text_model(MODEL_PAYLOAD, ADAPTER_PATH, MODEL_LOADER)
-GENERATION_LOCK = threading.Lock()
-REQUEST_SLOTS = threading.BoundedSemaphore(MAX_CONCURRENT_REQUESTS)
-
 
-def text_content(content: Any) -> str:
-    if isinstance(content, str):
-        return content
-    if content is None:
-        return ""
-    if not isinstance(content, list):
-        raise ValueError("Message content must be text or an array of text parts.")
-    parts: list[str] = []
-    for part in content:
-        if not isinstance(part, dict) or part.get("type") != "text":
-            raise ValueError("The bundled model server accepts text content only.")
-        text = part.get("text")
-        if not isinstance(text, str):
-            raise ValueError("Every text content part must contain a string text field.")
-        parts.append(text)
-    return "\n".join(parts)
-
-
-def normalize_messages(raw_messages: Any) -> list[dict[str, str]]:
-    if not isinstance(raw_messages, list) or not raw_messages:
-        raise ValueError("Request must include a non-empty messages array.")
-    if len(raw_messages) > MAX_MESSAGES:
-        raise ValueError(f"Request exceeds the {MAX_MESSAGES}-message limit.")
-
-    system_parts = [SYSTEM_PROMPT] if SYSTEM_PROMPT else []
-    conversation: list[dict[str, str]] = []
-    for raw in raw_messages:
-        if not isinstance(raw, dict) or raw.get("role") not in {"system", "user", "assistant", "tool"}:
-            raise ValueError("Each message must have a supported role and text content.")
-        role = str(raw["role"])
-        content = text_content(raw.get("content", ""))
-        if role == "system":
-            if content.strip():
-                system_parts.append(content.strip())
+from evaluate import configure_hugging_face_cache, resolve_adapter_path
+from model_contract import assert_certified_base_model_revision, assert_certified_model_config
+
+class BearerAuthMiddleware:
+    """Keep TT's auth boundary on ALL routes, including upstream admin APIs.
+
+    vLLM's default bearer guard only covers selected API prefixes. This is
+    transport admission only; requests/responses and generation stay upstream.
+    """
+    def __init__(self, app):
+        self.app = app
+        self.expected = os.environ["TT_API_KEY"].encode()
+
+    async def __call__(self, scope, receive, send):
+        if scope["type"] in {"http", "websocket"}:
+            authorization = dict(scope.get("headers", [])).get(b"authorization", b"")
+            scheme, _, token = authorization.partition(b" ")
+            if scheme.lower() != b"bearer" or not hmac.compare_digest(token, self.expected):
+                if scope["type"] == "websocket":
+                    await send({"type": "websocket.close", "code": 1008})
+                else:
+                    await send({"type": "http.response.start", "status": 401,
+                                "headers": [(b"content-type", b"application/json")]})
+                    await send({"type": "http.response.body", "body": b'{"error":"Unauthorized"}'})
+                return
+        await self.app(scope, receive, send)
+
+
+def build_vllm_args(model_source: str, adapter_path: str | None, temp_dir: Path) -> list[str]:
+    args = [
+        "serve", model_source,
+        "--served-model-name", os.environ["TT_MODEL_NAME"],
+        "--host", os.environ.get("TT_HOST", "127.0.0.1"),
+        "--port", os.environ.get("TT_PORT", "8000"),
+        "--max-model-len", os.environ.get("TT_CONTEXT_LENGTH", "16384"),
+        "--max-num-seqs", os.environ.get("TT_MAX_CONCURRENT_REQUESTS", "1"),
+        "--gpu-memory-utilization", os.environ.get("TT_GPU_MEMORY_UTILIZATION", "0.8"),
+        "--generation-config", "vllm",
+        "--override-generation-config", json.dumps({
+            "max_new_tokens": int(os.environ.get("TT_MAX_TOKENS", "512")),
+            "temperature": float(os.environ.get("TT_TEMPERATURE", "0")),
+            "top_p": float(os.environ.get("TT_TOP_P", "1")),
+        }),
+        "--language-model-only",
+        "--no-enable-log-requests", "--no-enable-log-outputs",
+    ]
+    if os.environ.get("TT_API_KEY"):
+        args += ["--middleware", "serve.BearerAuthMiddleware"]
+    base = os.environ["TT_BASE_MODEL"]
+    # Upstream model-specific parsers, not a TT parser or model-size heuristic.
+    parser = {
+        "Qwen/Qwen3.5-2B": "qwen3_xml",
+        "nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-BF16": "qwen3_coder",
+        "meta-models/Muse-Glimmer-30B": "muse_glimmer",
+    }[base]
+    args += ["--enable-auto-tool-choice", "--tool-call-parser", parser]
+    if base != "meta-models/Muse-Glimmer-30B":
+        args += ["--default-chat-template-kwargs", '{"enable_thinking":false}']
+    if adapter_path:
+        config = json.loads((Path(adapter_path) / "adapter_config.json").read_text())
+        adapter_base = config.get("base_model_name_or_path")
+        if adapter_base not in {os.environ["TT_BASE_MODEL"], model_source}:
+            raise ValueError("The adapter base does not match the verified serving model.")
+        base_name = "base:" + os.environ["TT_BASE_MODEL"]
+        args[args.index("--served-model-name") + 1] = base_name
+        args += ["--enable-lora", "--max-lora-rank", str(max(8, int(config["r"]))),
+                 "--lora-modules", json.dumps({"name": os.environ["TT_MODEL_NAME"],
+                 "path": adapter_path, "base_model_name": base_name})]
+    prompt = os.environ.get("TT_SYSTEM_PROMPT", "").strip()
+    if prompt:
+        source = Path(model_source)
+        template_file = source / "chat_template.jinja"
+        if template_file.is_file():
+            template = template_file.read_text(encoding="utf-8")
         else:
-            conversation.append({"role": role, "content": content})
-
-    if not conversation:
-        raise ValueError("Request must contain at least one non-system message.")
-    messages: list[dict[str, str]] = []
-    if system_parts:
-        # The certified chat templates accept one leading system message. Merge the invariant owner
-        # prompt and any client context instead of triggering an implicit
-        # template fallback with duplicate system turns.
-        messages.append({"role": "system", "content": "\n\n".join(system_parts)})
-    messages.extend(conversation)
-    if sum(len(message["content"]) for message in messages) > MAX_PROMPT_CHARS:
-        raise ValueError(f"Prompt exceeds the {MAX_PROMPT_CHARS}-character limit.")
-    return messages
-
-
-def generate_text(messages: list[dict[str, str]], generation: dict[str, Any]) -> tuple[str, int, int]:
-    try:
-        prompt = TOKENIZER.apply_chat_template(
-            messages,
-            tokenize=False,
-            add_generation_prompt=True,
-            **chat_template_kwargs(BASE_MODEL),
-        )
-    except Exception as exc:
-        raise ValueError(f"Model chat-template rendering failed: {exc}") from exc
-    inputs = TOKENIZER(prompt, return_tensors="pt")
-    if int(inputs["input_ids"].shape[-1]) > MAX_PROMPT_TOKENS:
-        raise ValueError(f"Prompt exceeds the {MAX_PROMPT_TOKENS}-token limit.")
-    target_device = next(MODEL.parameters()).device
-    inputs = {key: value.to(target_device) for key, value in inputs.items()}
-    with GENERATION_LOCK:
-        import torch
-
-        with torch.inference_mode():
-            generated = MODEL.generate(
-                **inputs,
-                max_new_tokens=int(generation["max_new_tokens"]),
-                pad_token_id=TOKENIZER.eos_token_id,
-                **sampling_kwargs(generation),
-            )
-    prompt_tokens = int(inputs["input_ids"].shape[-1])
-    completion_tokens = generated[0][prompt_tokens:]
-    content = TOKENIZER.decode(completion_tokens, skip_special_tokens=True).strip()
-    return content, prompt_tokens, int(completion_tokens.shape[-1])
-
-
-def bounded_number(value: Any, default: float, minimum: float, maximum: float) -> float:
-    number = float(default if value is None else value)
-    if not math.isfinite(number) or number < minimum or number > maximum:
-        raise ValueError(f"Generation value must be between {minimum} and {maximum}.")
-    return number
-
-
-def bounded_integer(value: Any, default: int, minimum: int, maximum: int) -> int:
-    number = bounded_number(value, default, minimum, maximum)
-    if not number.is_integer():
-        raise ValueError(f"Generation value must be an integer between {minimum} and {maximum}.")
-    return int(number)
-
-
-class Handler(BaseHTTPRequestHandler):
-    server_version = "TunedTensorLocalServer/0.1"
-
-    def setup(self) -> None:
-        super().setup()
-        self.connection.settimeout(30)
-
-    def authorized(self) -> bool:
-        if not API_KEY:
-            return True
-        supplied = self.headers.get("authorization", "")
-        return hmac.compare_digest(supplied, "Bearer " + API_KEY)
-
-    def send_json(self, status: int, payload: Any) -> None:
-        body = json.dumps(payload).encode("utf-8")
-        self.send_response(status)
-        self.send_header("Content-Type", "application/json")
-        self.send_header("Content-Length", str(len(body)))
-        self.end_headers()
-        self.wfile.write(body)
-
-    def do_GET(self) -> None:  # noqa: N802
-        if not self.authorized():
-            self.send_json(401, {"error": {"message": "Unauthorized"}})
-            return
-        if self.path in {"/", "/health"}:
-            self.send_json(200, {"status": "ok", "model": MODEL_NAME, "device": DEVICE})
-            return
-        if self.path == "/v1/models":
-            self.send_json(200, {
-                "object": "list",
-                "data": [{"id": MODEL_NAME, "object": "model", "owned_by": "tuned-tensor-local"}],
-            })
-            return
-        self.send_json(404, {"error": {"message": "Not found"}})
-
-    def do_POST(self) -> None:  # noqa: N802
-        if not self.authorized():
-            self.send_json(401, {"error": {"message": "Unauthorized"}})
-            return
-        if self.path not in {"/v1/chat/completions", "/chat/completions"}:
-            self.send_json(404, {"error": {"message": "Not found"}})
-            return
+            template = json.loads((source / "tokenizer_config.json").read_text())["chat_template"]
+        if not isinstance(template, str):
+            raise ValueError("Serving requires an unambiguous string chat template.")
+        # Literal data, never prompt text interpolated as Jinja source. Preserve
+        # tool calls/results verbatim while merging leading system context.
+        prefix = "{%- set tt = namespace(system=" + json.dumps(prompt) + ", history=[]) -%}"
+        prefix += "{%- for message in messages -%}{%- if message.role == 'system' -%}"
+        prefix += "{%- set tt.system = tt.system + '\\n\\n' + message.content -%}"
+        prefix += "{%- else -%}{%- set tt.history = tt.history + [message] -%}{%- endif -%}{%- endfor -%}"
+        prefix += "{%- set messages = [{'role': 'system', 'content': tt.system}] + tt.history -%}"
+        path = temp_dir / "chat-template.jinja"
+        path.write_text(prefix + template, encoding="utf-8")
+        args += ["--chat-template", str(path)]
+    return args
+
+
+def prepare_model_source() -> str:
+    base = os.environ["TT_BASE_MODEL"]
+    revision = os.environ.get("TT_BASE_MODEL_REVISION")
+    assert_certified_base_model_revision(base, revision, "Serving base model revision")
+    configure_hugging_face_cache(os.environ.get("HF_HOME"))
+    source = os.environ.get("TT_MODEL_SOURCE", base)
+    if not Path(source).is_dir():
+        if source != base:
+            raise ValueError("Explicit serving snapshot is not an existing directory.")
+        from huggingface_hub import snapshot_download
+        source = snapshot_download(repo_id=base, revision=revision, local_files_only=True)
+    config = json.loads((Path(source) / "config.json").read_text(encoding="utf-8"))
+    assert_certified_model_config(config, "Serving model config", base)
+    return str(Path(source).resolve())
+
+
+def run_server() -> None:
+    host = os.environ.get("TT_HOST", "127.0.0.1")
+    port = int(os.environ.get("TT_PORT", "8000"))
+    # Early diagnostic; upstream still owns the real bind and handles races.
+    family = socket.AF_INET6 if ":" in host else socket.AF_INET
+    with socket.socket(family, socket.SOCK_STREAM) as probe:
+        probe.bind((host, port))
+    model_source = prepare_model_source()
+    if os.environ.get("TT_API_KEY"):
+        os.environ["VLLM_API_KEY"] = os.environ["TT_API_KEY"]
+    with TemporaryDirectory(prefix="tt-serve-") as temp:
+        adapter = resolve_adapter_path(os.environ.get("TT_MODEL_ARTIFACT"), Path(temp))
+        args = build_vllm_args(model_source, adapter, Path(temp))
+        from vllm.entrypoints.cli.main import main
+        previous = sys.argv
         try:
-            length = int(self.headers.get("content-length", "0"))
-            if length <= 0 or length > MAX_REQUEST_BYTES:
-                raise ValueError("Request body is empty or too large.")
-            body = json.loads(self.rfile.read(length))
-            if not isinstance(body, dict):
-                raise ValueError("Request body must be a JSON object.")
-            if body.get("stream"):
-                self.send_json(400, {"error": {"message": "Streaming is not supported yet."}})
-                return
-            messages = normalize_messages(body.get("messages"))
-            generation = {
-                "max_new_tokens": bounded_integer(body.get("max_tokens"), DEFAULT_MAX_TOKENS, 1, 8192),
-                "temperature": bounded_number(body.get("temperature"), DEFAULT_TEMPERATURE, 0, 5),
-                "top_p": bounded_number(body.get("top_p"), DEFAULT_TOP_P, 0, 1),
-            }
-            if not REQUEST_SLOTS.acquire(blocking=False):
-                self.send_json(429, {"error": {"message": "The local model is busy; retry shortly."}})
-                return
-            started = time.perf_counter()
-            try:
-                content, prompt_tokens, completion_tokens = generate_text(messages, generation)
-            finally:
-                REQUEST_SLOTS.release()
-            latency_ms = round((time.perf_counter() - started) * 1000)
-            self.send_json(200, {
-                "id": "chatcmpl-" + uuid.uuid4().hex,
-                "object": "chat.completion",
-                "created": int(time.time()),
-                "model": MODEL_NAME,
-                "choices": [{
-                    "index": 0,
-                    "message": {"role": "assistant", "content": content},
-                    "finish_reason": "stop",
-                }],
-                "usage": {
-                    "prompt_tokens": prompt_tokens,
-                    "completion_tokens": completion_tokens,
-                    "total_tokens": prompt_tokens + completion_tokens,
-                },
-                "tt_local": {"latency_ms": latency_ms, "device": DEVICE},
-            })
-        except (ValueError, TypeError, json.JSONDecodeError) as exc:
-            self.send_json(400, {"error": {"message": str(exc)}})
-        except Exception as exc:
-            print(f"Model server error: {type(exc).__name__}: {exc}", file=sys.stderr, flush=True)
-            self.send_json(500, {"error": {"message": "Internal model server error."}})
-
-    def log_message(self, fmt: str, *args: Any) -> None:
-        print("%s - %s" % (self.address_string(), fmt % args), flush=True)
-
-
-def main() -> None:
-    print(f"Serving {MODEL_NAME} on http://{HOST}:{PORT}", flush=True)
-    print(f"OpenAI-compatible endpoint: http://{HOST}:{PORT}/v1/chat/completions", flush=True)
-    print(f"Device: {DEVICE}", flush=True)
-    server = ThreadingHTTPServer((HOST, PORT), Handler)
-    server.daemon_threads = True
-    server.request_queue_size = max(2, MAX_CONCURRENT_REQUESTS * 2)
-    try:
-        server.serve_forever()
-    except KeyboardInterrupt:
-        print("Model server stopped.", flush=True)
-    finally:
-        server.server_close()
+            sys.argv = ["vllm", *args]
+            main()
+        finally:
+            sys.argv = previous
 
 
 if __name__ == "__main__":
-    main()
+    # Ensure temporary artifacts unwind if stopped before upstream installs its
+    # own graceful shutdown handlers. Node forwards signals to the whole group.
+    signal.signal(signal.SIGTERM, lambda *_: sys.exit(0))
+    run_server()

package/training/adapter/src/train.py

@@ -92,8 +92,8 @@ def assert_certified_request() -> None:
 def require_cuda() -> dict[str, Any]:
     if not torch.cuda.is_available():
         raise RuntimeError(
-            "TT Local bundled training requires an NVIDIA CUDA GPU. "
-            "Run this workflow on DGX Spark and verify it first with tt-local doctor."
+            "TT local training requires an NVIDIA CUDA GPU. "
+            "Run this workflow on a supported CUDA host and verify it first with tt doctor."
         )
     device_index = torch.cuda.current_device()
     properties = torch.cuda.get_device_properties(device_index)

package/training/serving/pyproject.toml

@@ -0,0 +1,9 @@
+[project]
+name = "tuned-tensor-serving"
+version = "0.1.0"
+requires-python = ">=3.12,<3.13"
+dependencies = ["vllm==0.28.0"]
+
+[tool.uv]
+package = false
+environments = ["sys_platform == 'linux'"]