From 4a5268b04827a92481cf0ce99627af4b0a8c7b46 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E9=87=8E=E7=94=9F=E3=81=AE=E7=94=B7?= Date: Mon, 21 Sep 2026 20:49:43 +0900 Subject: [PATCH] Add optional Windows ROCm backend for DiffusionGemma --- .gitignore | 4 + README.md | 13 +- docs/gfx1151.md | 181 +++++++++++++++++++ package.json | 4 + requirements-rocm.txt | 5 + scripts/http-smoke.ts | 174 ++++++++++++++++++ scripts/rocm_server.py | 376 +++++++++++++++++++++++++++++++++++++++ scripts/start-rocm.ps1 | 277 ++++++++++++++++++++++++++++ src/index.ts | 26 ++- test/http-smoke.test.ts | 121 +++++++++++++ test/index.test.ts | 40 +++++ test/test_rocm_server.py | 204 +++++++++++++++++++++ 12 files changed, 1416 insertions(+), 9 deletions(-) create mode 100644 docs/gfx1151.md create mode 100644 requirements-rocm.txt create mode 100644 scripts/http-smoke.ts create mode 100644 scripts/rocm_server.py create mode 100644 scripts/start-rocm.ps1 create mode 100644 test/http-smoke.test.ts create mode 100644 test/index.test.ts create mode 100644 test/test_rocm_server.py diff --git a/.gitignore b/.gitignore index 71d8b6f..1dba1e5 100644 --- a/.gitignore +++ b/.gitignore @@ -7,3 +7,7 @@ eval/runs/ .DS_Store coverage/ dist/ +.venv-rocm/ +.runtime/ +__pycache__/ +*.pyc diff --git a/README.md b/README.md index 39939e6..2829b09 100644 --- a/README.md +++ b/README.md @@ -4,7 +4,12 @@ A local, Jev-compatible `POST /v1/systemone` API written in TypeScript for [Bun](https://bun.sh/), backed by DiffusionGemma through an OpenAI-compatible Chat Completions endpoint. -The defaults target: +An optional **Windows / AMD ROCm backend**, tested on Radeon 8060S (gfx1151), +runs the original DiffusionGemma checkpoint in FP16 with PyTorch and +Transformers. See [the ROCm guide](docs/gfx1151.md) for setup, the PowerShell +launcher and an end-to-end HTTP smoke check. + +The oMLX defaults target: - inference server: `http://127.0.0.1:8000` - model: `diffusiongemma-26B-A4B-it-4bit` @@ -35,7 +40,7 @@ them for consequential decisions. ## Run with oMLX -Requires Bun 1.2+ and a running oMLX server. +Requires Bun 1.4.2+ (for the checked-in lockfile) and a running oMLX server. ```sh bun install @@ -158,6 +163,7 @@ bun install bun test bun run typecheck bun run smoke # live call to the configured inference server +bun run smoke:http # /ready and all three decision types through LocalJev HTTP ``` ## Evaluate different models @@ -196,7 +202,8 @@ and [`lmstudio-ai/lmstudio-bug-tracker#2037`](https://github.com/lmstudio-ai/lmstudio-bug-tracker/issues/2037). The reported MLX backend fails to load `diffusion_gemma`, while the normal llama.cpp backend reports an unknown architecture. oMLX already loads and serves your exact -checkpoint successfully, so it is the better runner for this Mac today. +checkpoint successfully in the Mac setup. For Windows gfx1151, see the +[optional ROCm backend](docs/gfx1151.md). Even after LM Studio adds ordinary generation support, changing runners alone will not make the result OpenJev-equivalent. The runner must expose seeded diffusion diff --git a/docs/gfx1151.md b/docs/gfx1151.md new file mode 100644 index 0000000..7db070b --- /dev/null +++ b/docs/gfx1151.md @@ -0,0 +1,181 @@ +# Windows ROCm / Radeon 8060S (gfx1151) + +The optional backend in this guide runs the original +[`google/diffusiongemma-26B-A4B-it`](https://huggingface.co/google/diffusiongemma-26B-A4B-it) +checkpoint with ROCm PyTorch and Transformers. LocalJev still generates and +validates probability JSON; it does not read option probabilities directly +from logits. The OpenAI-compatible backend implements the text-only subset +needed by LocalJev. + +LocalJev can also connect to an existing OpenAI-compatible ROCm inference +server through `LOCALJEV_UPSTREAM` and `LOCALJEV_UPSTREAM_MODEL`. Python and +the adapter below are only needed when using the bundled Transformers +backend. The existing oMLX defaults remain available. + +## Connect an existing ROCm server + +Configure the endpoint and the model ID advertised by that server: + +```powershell +$env:LOCALJEV_UPSTREAM = 'http://127.0.0.1:8000' +$env:LOCALJEV_UPSTREAM_MODEL = 'your-served-model-id' +bun --no-env-file run src/index.ts +``` + +Set `LOCALJEV_UPSTREAM_API_KEY` if the server requires authentication. +The runner must support the text Chat Completions and model-list endpoints +used by LocalJev. The instructions below provide a standalone ROCm runner +for the original DiffusionGemma checkpoint. + +## Requirements + +- Windows with a working AMD driver and **ROCm-enabled PyTorch** for gfx1151. + A CUDA or CPU PyTorch wheel will not work. The backend refuses CPU fallback. +- A Python environment that already passes a ROCm GPU check (Python 3.13 on + the tested host). See [AMD's Windows PyTorch documentation](https://rocm.docs.amd.com/projects/radeon-ryzen/en/latest/docs/install/installrad/windows/install-pytorch.html) + for driver and wheel installation. +- Bun **1.4.2 or newer**, matching the upstream lockfile format. +- About **52 GB of disk space for the model**, plus dependencies, and enough + GPU-accessible memory for FP16 weights and working allocations. A 128 GB + Strix Halo host is the target; this is not a 4-bit low-memory implementation. + +## Set up without changing an existing ROCm environment + +Run from the repository root in PowerShell. Set `$RocmPython` to the Python +executable containing your working ROCm PyTorch installation. Do not rely on +an unrelated `python` executable earlier on PATH. + +```powershell +$RocmPython = "$env:LOCALAPPDATA\Programs\Python\Python313\python.exe" +& $RocmPython -c 'import torch; assert torch.version.hip and torch.cuda.is_available(); print(torch.__version__); print(torch.cuda.get_device_properties(0))' +& $RocmPython -m venv --system-site-packages .venv-rocm +& .\.venv-rocm\Scripts\python.exe -m pip install -r requirements-rocm.txt +bun install --frozen-lockfile +``` + +The overlay environment reuses the installed ROCm PyTorch while keeping the +new Transformers version separate. `requirements-rocm.txt` deliberately does +not install PyTorch. No native compilation is required. + +## Start both servers + +Check the Python packages, HIP device and available ports without loading +model weights: + +```powershell +.\scripts\start-rocm.ps1 -CheckOnly +``` + +Then start the model and API: + +```powershell +.\scripts\start-rocm.ps1 +``` + +The first start downloads the official checkpoint through Hugging Face and +loads it onto the GPU. The launcher waits for the backend before starting +LocalJev. It listens on loopback, uses one in-flight inference request, and +stops its own backend when LocalJev exits. Existing services are never stopped +to free ports. Use another terminal for requests. + +For an explicitly downloaded checkpoint, a local model directory also works: + +```powershell +.\scripts\start-rocm.ps1 -Model '.runtime/models/diffusiongemma-26B-A4B-it' -LocalFilesOnly +``` + +The served model identifier is the value passed as `-Model`; the launcher +configures LocalJev to match it. Use `-Python` and `-Bun` for explicit +executables. The launcher also detects the repository-local Bun executable +at `.runtime/bun/bun-windows-x64/bun.exe` when present. + +The launcher configures its environment directly and does not load `.env`. +Set optional client authentication in the launching shell: + +```powershell +$env:LOCALJEV_API_KEY = 'your-local-client-key' +.\scripts\start-rocm.ps1 +``` + +## Verify the complete HTTP path + +```powershell +Invoke-RestMethod http://127.0.0.1:8080/ready +bun run smoke:http +``` + +`smoke:http` checks readiness, sends a real decision request containing +`choice`, `score`, and `noul`, and validates the returned probabilities and +token usage. It prints the upstream model and elapsed time. If client +authentication is enabled, set `LOCALJEV_API_KEY` in this terminal too. +For a different API URL, set `LOCALJEV_URL`. + +For development checks without loading model weights: + +```powershell +bun test +bun run typecheck +& .\.venv-rocm\Scripts\python.exe -m unittest discover -s test -p 'test_rocm_server.py' +``` + +## Backend behavior + +- FP16, explicit HIP device placement, SDPA attention and eager PyTorch MoE + operations are used by default. + The startup report includes the actual GPU architecture and package versions. +- DiffusionGemma uses its own denoising schedule. Ordinary autoregressive + `temperature` is not mapped onto that schedule. A dynamic cache avoids the + Transformers compilation path on Windows. +- JSON schemas are included in the prompt. There is no grammar-constrained + diffusion decoder; LocalJev retains its validation and corrective retries. +- Generation is serialized and queue admission is bounded. This is a local + development server, not a public multi-user inference service. +- A different model must be selected explicitly. Generic causal models may + be used for diagnostics, but their output does not validate DiffusionGemma. + +The original oMLX configuration and evaluation results remain available in the +[main README](../README.md). Those measurements were made on a Mac and do not +describe the ROCm backend. + +## Verified on September 20, 2026 + +The official model at revision +`f7f5b7f5fa82ffc52addd066915886d497f5517b` completed a real +`/ready` → `/v1/systemone` request on Windows with the following configuration: + +| Component | Validated value | +|---|---| +| GPU | AMD Radeon 8060S, `gfx1151` | +| Python | 3.13.9 | +| PyTorch | `2.15.0a0+rocm10.1.0a20260909` | +| HIP reported by PyTorch | `7.16.26362` | +| Transformers | `5.11.0` | +| Bun | `1.4.2` | +| Model placement | Entire model on `cuda:0` (PyTorch's HIP device API) | +| Inference | FP16, SDPA, eager MoE, dynamic cache | +| Request output ceiling | 512 tokens | +| Input / output tokens | 522 / 72 | +| Full HTTP decision latency | 30.80 seconds; 41.47 seconds after a fresh restart | + +The smoke request returned `technical` with probability `0.9`, a frustration +score of `0.9`, and an urgent probability of `0.8`. All three response types +passed validation without a corrective retry in both runs. Each was the first +request after loading. This is an integration check, +not a calibration study or a representative throughput benchmark. + +The weights occupy 51,647,701,024 bytes. Loading can temporarily use much more +system memory than the final GPU allocation. These results validate the +listed nightly stack; other ROCm versions require their own runtime check. + +Development validation passed 27 Bun tests, 12 Python tests, and the TypeScript +type check. Launcher checks covered Windows PowerShell 5.1 and PowerShell 7, +including failure, timeout, cancellation and child-process cleanup. + +## Separate quantized-runner validation + +ROCmFPX was used as an external runner for a separate Q8 evaluation. It is +not an installation, build or runtime dependency of this ROCm adapter; +the setup above uses ROCm PyTorch, Transformers and the original checkpoint. +The [recorded Q8 evaluation](https://github.com/Yasei-no-otoko/localjev/blob/be015c7e41ba8823b5c596db72418ea8a82cdb01/eval/reports/2026-09-20-gfx1151-diffusiongemma-q8-rocmfpx-fa-off/report.md) +belongs to that different runner and quantization. Its 240/240 valid outputs +are not a full-suite result for the FP16 adapter described here. diff --git a/package.json b/package.json index f1c8e40..5475d0e 100644 --- a/package.json +++ b/package.json @@ -4,12 +4,16 @@ "description": "A Jev-compatible System One API backed by a local OpenAI-compatible model", "type": "module", "private": true, + "engines": { + "bun": ">=1.4.2" + }, "scripts": { "start": "bun run src/index.ts", "dev": "bun --watch run src/index.ts", "test": "bun test", "typecheck": "tsc --noEmit", "smoke": "bun run scripts/live-smoke.ts", + "smoke:http": "bun run scripts/http-smoke.ts", "eval": "bun run scripts/eval/run.ts", "eval:report": "bun run scripts/eval/report.ts" }, diff --git a/requirements-rocm.txt b/requirements-rocm.txt new file mode 100644 index 0000000..a9e499e --- /dev/null +++ b/requirements-rocm.txt @@ -0,0 +1,5 @@ +# Install into an environment that already exposes a working ROCm PyTorch. +# Deliberately omit torch: PyPI's default wheel must not replace the HIP build. +transformers==5.11.0 +accelerate>=1.10,<2 +Pillow>=11,<13 diff --git a/scripts/http-smoke.ts b/scripts/http-smoke.ts new file mode 100644 index 0000000..34117ae --- /dev/null +++ b/scripts/http-smoke.ts @@ -0,0 +1,174 @@ +import type { SystemOneRequest } from "../src/types"; + +// The API derives its seed from state and questions, so keep this request fixed. +const request: SystemOneRequest = { + model: "localjev-latest", + state: "Hi, I have been trying to connect Stripe but keep getting a 403 error.", + questions: { + department: { + type: "choice", + instructions: "Which team should handle this?", + criteria: { + billing: "Payment or subscription issues", + technical: "Bugs or integration problems", + sales: "Pricing or account questions", + }, + }, + frustration: { + type: "score", + instructions: "How frustrated does the customer appear?", + criteria: ["Calm", "Frustrated but civil", "Very angry"], + }, + urgent: { + type: "noul", + instructions: "Does this require an immediate response?", + criteria: null, + }, + }, +}; + +function object(value: unknown, path: string): Record { + if (typeof value !== "object" || value === null || Array.isArray(value)) { + throw new Error(`${path} must be an object`); + } + return value as Record; +} + +function keys(value: Record, expected: string[], path: string): void { + if ( + Object.keys(value).length !== expected.length || + expected.some((key) => !Object.hasOwn(value, key)) + ) { + throw new Error(`${path} must contain exactly: ${expected.join(", ")}`); + } +} + +function boundedNumber(value: unknown, maximum: number, path: string): number { + if (typeof value !== "number" || !Number.isFinite(value) || value < 0 || value > maximum) { + throw new Error(`${path} must be a finite number between 0 and ${maximum}`); + } + return value; +} + +function distribution(value: unknown, labels: string[], path: string): void { + const probabilities = object(value, path); + keys(probabilities, labels, path); + const total = labels.reduce( + (sum, label) => sum + boundedNumber(probabilities[label], 1, `${path}.${label}`), + 0, + ); + if (Math.abs(total - 1) > 1e-6) { + throw new Error(`${path} must sum to 1 (received ${total})`); + } +} + +function validateResponse(raw: unknown): Record { + const body = object(raw, "response"); + if (typeof body.model !== "string" || body.model.length === 0) { + throw new Error("response.model must be a nonempty string"); + } + const answers = object(body.answers, "answers"); + keys(answers, Object.keys(request.questions), "answers"); + for (const [name, question] of Object.entries(request.questions)) { + const path = `answers.${name}`; + const answer = object(answers[name], path); + if (answer.type !== question.type) throw new Error(`${path}.type must be ${question.type}`); + if (question.type === "noul") { + boundedNumber(answer.noul, 1, `${path}.noul`); + continue; + } + boundedNumber(answer.confidence, 1, `${path}.confidence`); + if (question.type === "choice") { + const labels = Object.keys(question.criteria); + distribution(answer.probabilities, labels, `${path}.probabilities`); + if (typeof answer.choice !== "string" || !labels.includes(answer.choice)) { + throw new Error(`${path}.choice must be one of the requested labels`); + } + } else { + const labels = question.criteria.map((_, index) => String(index)); + distribution(answer.probabilities, labels, `${path}.probabilities`); + boundedNumber(answer.score, question.criteria.length - 1, `${path}.score`); + const legend = object(answer.legend, `${path}.legend`); + keys(legend, labels, `${path}.legend`); + for (const [index, criterion] of question.criteria.entries()) { + if (legend[String(index)] !== criterion) { + throw new Error(`${path}.legend.${index} must match the requested criterion`); + } + } + } + } + const usage = object(body.usage, "usage"); + for (const name of ["input_tokens", "output_tokens"]) { + const count = boundedNumber(usage[name], Number.MAX_SAFE_INTEGER, `usage.${name}`); + if (!Number.isInteger(count)) throw new Error(`usage.${name} must be an integer`); + } + return body; +} + +export async function runHttpSmoke(options: { + url: string; + apiKey?: string; + timeoutMs?: number; +}): Promise> { + const base = new URL(options.url.replace(/\/+$/, "") + "/"); + if (!["http:", "https:"].includes(base.protocol) || base.username || base.password || base.search || base.hash) { + throw new Error("LocalJev URL must be HTTP(S) without credentials, query parameters, or a fragment"); + } + const timeoutMs = options.timeoutMs ?? 300_000; + if (!Number.isSafeInteger(timeoutMs) || timeoutMs <= 0) { + throw new Error("Smoke timeout must be a positive integer number of milliseconds"); + } + const headers: HeadersInit = { + "content-type": "application/json", + ...(options.apiKey ? { authorization: `Bearer ${options.apiKey}` } : {}), + }; + async function fetchJson(path: string, init?: RequestInit): Promise { + const response = await fetch(new URL(path, base), { + ...init, + headers, + signal: AbortSignal.timeout(timeoutMs), + }); + if (response.status !== 200) { + throw new Error(`${path} returned HTTP ${response.status}: ${(await response.text()).slice(0, 1_000)}`); + } + try { + return await response.json(); + } catch { + throw new Error(`${path} did not return valid JSON`); + } + } + const ready = object(await fetchJson("ready"), "ready"); + if (ready.status !== "ready" || typeof ready.upstream_model !== "string" || !ready.upstream_model) { + throw new Error("/ready must report status=ready and a nonempty upstream_model"); + } + const started = performance.now(); + const body = validateResponse(await fetchJson("v1/systemone", { + method: "POST", + body: JSON.stringify(request), + })); + return { + status: "ok", + url: base.href, + upstream_model: ready.upstream_model, + latency_ms: Math.round((performance.now() - started) * 100) / 100, + model: body.model, + answers: body.answers, + usage: body.usage, + }; +} + +if (import.meta.main) { + try { + const host = process.env.LOCALJEV_HOST ?? "127.0.0.1"; + const clientHost = host === "0.0.0.0" ? "127.0.0.1" : host === "::" ? "[::1]" : host; + const evidence = await runHttpSmoke({ + url: process.env.LOCALJEV_URL ?? `http://${clientHost.includes(":") && !clientHost.startsWith("[") ? `[${clientHost}]` : clientHost}:${process.env.LOCALJEV_PORT ?? "8080"}`, + apiKey: process.env.LOCALJEV_API_KEY ?? "", + timeoutMs: Number(process.env.LOCALJEV_SMOKE_TIMEOUT ?? "300") * 1_000, + }); + console.log(JSON.stringify(evidence)); + } catch (error) { + console.error(`HTTP smoke failed: ${error instanceof Error ? error.message : String(error)}`); + process.exitCode = 1; + } +} diff --git a/scripts/rocm_server.py b/scripts/rocm_server.py new file mode 100644 index 0000000..9051ffa --- /dev/null +++ b/scripts/rocm_server.py @@ -0,0 +1,376 @@ +"""Local text-only Chat Completions adapter for ROCm; JSON schemas are prompted. + +Install requirements-rocm.txt into an environment with a working HIP PyTorch. +The default model is the original DiffusionGemma checkpoint, not an MLX quant. +""" + +from __future__ import annotations + +import argparse +from dataclasses import dataclass +import hmac +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer +import json +import logging +import math +import os +import re +import threading +import time +import uuid +from typing import Any + + +DEFAULT_MODEL = "google/diffusiongemma-26B-A4B-it" +LOG = logging.getLogger("localjev.rocm") + + +class RequestError(Exception): + def __init__(self, message: str, status: int = 400): + super().__init__(message) + self.status = status + + +@dataclass(frozen=True) +class Limits: + max_body_bytes: int = 1_048_576 + max_messages: int = 64 + max_input_tokens: int = 8192 + max_output_tokens: int = 2048 + max_queue: int = 2 + queue_timeout: float = 180.0 + + +@dataclass(frozen=True) +class ChatRequest: + messages: list[dict[str, str]] + max_tokens: int + temperature: float + seed: int | None + + +@dataclass(frozen=True) +class Generation: + content: str + prompt_tokens: int + completion_tokens: int + finish_reason: str + + +def integer(value: Any, name: str, minimum: int, maximum: int) -> int: + if type(value) is not int or not minimum <= value <= maximum: + raise RequestError(f"{name} must be an integer from {minimum} to {maximum}") + return value + + +def parse_request(body: Any, model_id: str, limits: Limits) -> ChatRequest: + if not isinstance(body, dict): + raise RequestError("The request must be a JSON object") + if body.get("model") != model_id: + raise RequestError(f"Only model {model_id!r} is loaded", 404) + if body.get("stream", False) is not False: + raise RequestError("Only non-streaming requests are supported") + integer(body.get("n", 1), "n", 1, 1) + for field in ("tools", "tool_choice", "stop", "logprobs", "top_logprobs"): + if body.get(field) is not None: + raise RequestError(f"{field} is not supported by this text adapter") + messages = body.get("messages") + if not isinstance(messages, list) or not 1 <= len(messages) <= limits.max_messages: + raise RequestError(f"messages must contain 1 to {limits.max_messages} messages") + clean = [] + for message in messages: + if not isinstance(message, dict) or message.get("role") not in ("system", "user", "assistant"): + raise RequestError("Each message needs a system, user, or assistant role") + content = message.get("content") + if not isinstance(content, str): + raise RequestError("Only string message content is supported; no image/audio inputs") + clean.append({"role": message["role"], "content": content}) + max_tokens = integer(body.get("max_tokens", min(512, limits.max_output_tokens)), "max_tokens", 1, limits.max_output_tokens) + if "max_completion_tokens" in body: + raise RequestError("Use max_tokens instead of max_completion_tokens") + temperature = body.get("temperature", 0.0) + if type(temperature) not in (int, float) or not math.isfinite(temperature) or not 0 <= temperature <= 2: + raise RequestError("temperature must be a finite number from 0 to 2") + seed = body.get("seed") + if seed is not None: + seed = integer(seed, "seed", 0, 2**32 - 1) + template = body.get("chat_template_kwargs", {}) + if not isinstance(template, dict) or set(template) - {"enable_thinking"}: + raise RequestError("Only chat_template_kwargs.enable_thinking=false is supported") + if template.get("enable_thinking", False) is not False: + raise RequestError("Thinking must be disabled for this decision endpoint") + response_format = body.get("response_format") + if response_format is not None: + if not isinstance(response_format, dict): + raise RequestError("response_format must be an object") + format_type = response_format.get("type") + instruction = None + if format_type == "json_schema": + definition = response_format.get("json_schema") + if not isinstance(definition, dict) or not isinstance(definition.get("schema"), dict): + raise RequestError("response_format.json_schema.schema must be an object") + instruction = "Return only JSON matching this schema, with no explanation or markdown:\n" + json.dumps(definition["schema"], ensure_ascii=False) + elif format_type == "json_object": + instruction = "Return only a JSON object, with no explanation or markdown." + elif format_type != "text": + raise RequestError("response_format.type must be text, json_object, or json_schema") + # This is a prompt instruction, not constrained decoding or schema enforcement. + if instruction: + if clean[0]["role"] == "system": + clean[0]["content"] += "\n\n" + instruction + else: + clean.insert(0, {"role": "system", "content": instruction}) + return ChatRequest(clean, max_tokens, float(temperature), seed) + + +def final_text(text: str) -> str: + """Remove complete reasoning blocks while preserving the final answer.""" + text = re.sub(r".*?", "", text, flags=re.DOTALL) + text = re.sub(r"<\|channel>thought\s*.*?", "", text, flags=re.DOTALL) + return text.strip() + + +def trim_generated_tokens(tokens: list[int], eos_ids: list[int], maximum: int) -> tuple[list[int], str]: + # Diffusion generates a whole canvas, which can extend past max_tokens. + for index, token in enumerate(tokens[:maximum]): + if token in eos_ids: + return tokens[:index + 1], "stop" + return tokens[:maximum], "length" if len(tokens) >= maximum else "stop" + + +class TransformersRuntime: + def __init__(self, args: argparse.Namespace, limits: Limits): + # Import lazily so request validation tests never load torch or touch a GPU. + import torch + import transformers + from transformers import AutoConfig, AutoModelForCausalLM, AutoProcessor, AutoTokenizer, DiffusionGemmaForBlockDiffusion + + if not torch.version.hip or not torch.cuda.is_available(): + raise RuntimeError("A working ROCm/HIP PyTorch GPU is required; CPU/CUDA fallback is disabled") + if not 0 <= args.device < torch.cuda.device_count(): + raise RuntimeError(f"ROCm device {args.device} is not available") + properties = torch.cuda.get_device_properties(args.device) + architecture = getattr(properties, "gcnArchName", "unknown") + if architecture.split(":", 1)[0] != "gfx1151" and not args.allow_other_gpu: + raise RuntimeError(f"Expected gfx1151, found {architecture}; use --allow-other-gpu for another HIP GPU") + self.torch = torch + self.device = torch.device(f"cuda:{args.device}") + torch.cuda.set_device(self.device) + self.model_id = args.model + self.limits = limits + self.info = { + "device": str(self.device), "gpu": properties.name, "architecture": architecture, + "hip": torch.version.hip, "torch": torch.__version__, "transformers": transformers.__version__, + "dtype": args.dtype, "attention": "sdpa", "experts": "eager", + } + LOG.info("ROCm runtime: %s", json.dumps(self.info)) + load_options = {"local_files_only": args.local_files_only, "trust_remote_code": False} + if args.revision: + load_options["revision"] = args.revision + config = AutoConfig.from_pretrained(args.model, **load_options) + self.is_diffusion = config.model_type == "diffusion_gemma" + if self.is_diffusion: + model_class = DiffusionGemmaForBlockDiffusion + self.processor = AutoProcessor.from_pretrained(args.model, **load_options) + self.tokenizer = self.processor.tokenizer + else: + model_class = AutoModelForCausalLM + self.processor = AutoTokenizer.from_pretrained(args.model, **load_options) + self.tokenizer = self.processor + self.info["model_type"] = config.model_type + self.model = model_class.from_pretrained( + args.model, config=config, dtype=getattr(torch, args.dtype), + device_map={"": str(self.device)}, attn_implementation="sdpa", + # The 5.11 grouped_mm auto-selection checks CUDA-shaped capability + # without distinguishing HIP. Keep MoE on ordinary PyTorch linears. + experts_implementation="eager", **load_options, + ).eval() + misplaced = [name for name, value in self.model.named_parameters() if value.device != self.device] + if misplaced: + raise RuntimeError(f"Model was not fully placed on {self.device}: {misplaced[:3]}") + torch.cuda.synchronize(self.device) + LOG.info("Loaded %s (%s) entirely on %s", args.model, config.model_type, architecture) + if self.is_diffusion: + LOG.info("Diffusion sampling uses checkpoint t_min/t_max; request temperature does not replace its schedule. Dynamic cache disables compilation.") + + def generate(self, request: ChatRequest) -> Generation: + inputs = self.processor.apply_chat_template( + request.messages, tokenize=True, add_generation_prompt=True, + return_dict=True, return_tensors="pt", enable_thinking=False, + ) + prompt_tokens = inputs["input_ids"].shape[-1] + if prompt_tokens > self.limits.max_input_tokens: + raise RequestError(f"Prompt has {prompt_tokens} tokens; maximum is {self.limits.max_input_tokens}") + inputs = inputs.to(self.device) + # Transformers 5.11 DiffusionGemma always returns .sequences. Its + # compile path is entered only for a static (compileable) cache. + options = {"max_new_tokens": request.max_tokens, "cache_implementation": "dynamic"} + if not self.is_diffusion: + options["do_sample"] = request.temperature > 0 + options["disable_compile"] = True + if request.temperature > 0: + options["temperature"] = request.temperature + if self.tokenizer.pad_token_id is None: + options["pad_token_id"] = self.tokenizer.eos_token_id + if request.seed is not None: + self.torch.manual_seed(request.seed) + with self.torch.inference_mode(): + output = self.model.generate(**inputs, **options) + sequences = output.sequences if hasattr(output, "sequences") else output + generated = sequences[0, prompt_tokens:].tolist() + eos_ids = self.model.generation_config.eos_token_id + if eos_ids is None: + eos_ids = self.tokenizer.eos_token_id + if not isinstance(eos_ids, list): + eos_ids = [] if eos_ids is None else [eos_ids] + generated, finish_reason = trim_generated_tokens(generated, eos_ids, request.max_tokens) + # Skip control tokens individually after removing thought-channel text. + text = final_text(self.tokenizer.decode(generated, skip_special_tokens=False)) + for token in self.tokenizer.all_special_tokens: + text = text.replace(token, "") + return Generation(text.strip(), prompt_tokens, len(generated), finish_reason) + + +class InferenceService: + def __init__(self, runtime: Any, limits: Limits): + self.runtime = runtime + self.limits = limits + self.admission = threading.BoundedSemaphore(1 + limits.max_queue) + self.inference = threading.Lock() + + def complete(self, body: Any) -> dict[str, Any]: + request = parse_request(body, self.runtime.model_id, self.limits) + if not self.admission.acquire(blocking=False): + raise RequestError("Inference queue is full; retry shortly", 429) + try: + if not self.inference.acquire(timeout=self.limits.queue_timeout): + raise RequestError("Timed out waiting for inference; retry shortly", 429) + try: + started = time.monotonic() + result = self.runtime.generate(request) + LOG.info("Generated %d tokens in %.2f seconds (%s)", result.completion_tokens, time.monotonic() - started, result.finish_reason) + finally: + self.inference.release() + finally: + self.admission.release() + return { + "id": "chatcmpl-" + uuid.uuid4().hex, "object": "chat.completion", + "created": int(time.time()), "model": self.runtime.model_id, + "choices": [{"index": 0, "message": {"role": "assistant", "content": result.content}, "finish_reason": result.finish_reason}], + "usage": {"prompt_tokens": result.prompt_tokens, "completion_tokens": result.completion_tokens, "total_tokens": result.prompt_tokens + result.completion_tokens}, + } + + +def make_server(service: InferenceService, host: str, port: int, api_key: str = "") -> ThreadingHTTPServer: + class Handler(BaseHTTPRequestHandler): + def setup(self): + super().setup() + self.connection.settimeout(30) + + def log_message(self, format: str, *args): + LOG.info(format, *args) + + def send_json(self, status: int, body: dict[str, Any]): + encoded = json.dumps(body, ensure_ascii=False, allow_nan=False).encode("utf-8") + self.send_response(status) + self.send_header("Content-Type", "application/json; charset=utf-8") + self.send_header("Content-Length", str(len(encoded))) + if status == 429: + self.send_header("Retry-After", "1") + self.end_headers() + self.wfile.write(encoded) + + def authorize(self): + supplied = self.headers.get("Authorization", "") + if api_key and not hmac.compare_digest(supplied.encode(), ("Bearer " + api_key).encode()): + raise RequestError("Invalid upstream API key", 401) + + def run_request(self, action): + try: + self.authorize() + action() + except RequestError as error: + self.send_json(error.status, {"error": {"message": str(error), "type": "invalid_request_error" if error.status < 500 else "server_error"}}) + except (BrokenPipeError, ConnectionResetError, TimeoutError): + LOG.warning("Client disconnected or request body timed out") + except Exception: + LOG.exception("Inference request failed") + self.send_json(500, {"error": {"message": "Inference failed; inspect the backend log", "type": "server_error"}}) + + def do_GET(self): + def respond(): + if self.path == "/health": + self.send_json(200, {"status": "ok", "model": service.runtime.model_id, "runtime": service.runtime.info, "structured_output": "prompt_only"}) + elif self.path == "/v1/models": + self.send_json(200, {"object": "list", "data": [{"id": service.runtime.model_id, "object": "model", "created": 0, "owned_by": "local"}]}) + else: + raise RequestError("Unknown endpoint", 404) + self.run_request(respond) + + def do_POST(self): + def respond(): + if self.path != "/v1/chat/completions": + raise RequestError("Unknown endpoint", 404) + if self.headers.get("Transfer-Encoding"): + raise RequestError("Chunked request bodies are not supported") + try: + length = int(self.headers.get("Content-Length", "")) + except ValueError: + raise RequestError("Content-Length is required", 411) from None + if not 0 < length <= service.limits.max_body_bytes: + raise RequestError("Request body exceeds the configured size limit", 413) + payload = self.rfile.read(length) + if len(payload) != length: + raise RequestError("Incomplete request body") + # Consume a bounded body before rejecting its media type. Closing + # a Windows socket with unread bytes can reset the connection + # before the client receives the HTTP error response. + if self.headers.get_content_type() != "application/json": + raise RequestError("Content-Type must be application/json", 415) + def reject_constant(value): + raise ValueError(value) + try: + body = json.loads(payload, parse_constant=reject_constant) + except (ValueError, UnicodeDecodeError, RecursionError): + raise RequestError("Request body must contain valid JSON") from None + self.send_json(200, service.complete(body)) + self.run_request(respond) + + return ThreadingHTTPServer((host, port), Handler) + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--model", default=DEFAULT_MODEL) + parser.add_argument("--host", default="127.0.0.1") + parser.add_argument("--port", type=int, default=8000) + parser.add_argument("--device", type=int, default=0) + parser.add_argument("--dtype", choices=("float16", "bfloat16", "float32"), default="float16") + parser.add_argument("--allow-other-gpu", action="store_true") + parser.add_argument("--local-files-only", action="store_true") + parser.add_argument("--revision") + parser.add_argument("--max-input-tokens", type=int, default=8192) + parser.add_argument("--max-output-tokens", type=int, default=2048) + parser.add_argument("--max-queue", type=int, default=2) + parser.add_argument("--queue-timeout", type=float, default=180.0) + args = parser.parse_args() + if args.max_input_tokens < 1 or args.max_output_tokens < 1 or args.max_queue < 0 or not math.isfinite(args.queue_timeout) or args.queue_timeout <= 0: + parser.error("Token limits and queue timeout must be positive; max-queue must be nonnegative") + if not 1 <= args.port <= 65535: + parser.error("port must be from 1 to 65535") + logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s") + limits = Limits(max_input_tokens=args.max_input_tokens, max_output_tokens=args.max_output_tokens, max_queue=args.max_queue, queue_timeout=args.queue_timeout) + runtime = TransformersRuntime(args, limits) + service = InferenceService(runtime, limits) + server = make_server(service, args.host, args.port, os.environ.get("LOCALJEV_UPSTREAM_API_KEY", "")) + LOG.info("Ready at http://%s:%d; schema output is prompted, not grammar constrained", args.host, args.port) + try: + server.serve_forever() + except KeyboardInterrupt: + LOG.info("Stopping ROCm backend") + finally: + server.server_close() + + +if __name__ == "__main__": + main() diff --git a/scripts/start-rocm.ps1 b/scripts/start-rocm.ps1 new file mode 100644 index 0000000..bcc00aa --- /dev/null +++ b/scripts/start-rocm.ps1 @@ -0,0 +1,277 @@ +#Requires -Version 5.1 +<# +.SYNOPSIS +Start the native ROCm model server and LocalJev on Windows. +.DESCRIPTION +Uses an already prepared Python environment and model cache. The backend is +started in a hidden process, while LocalJev runs in the current terminal. +Press Ctrl+C to stop both. Existing listeners are never stopped. +Automatic Bun .env loading is disabled; parameters override inherited LocalJev +connection settings. LOCALJEV_API_KEY and LOCALJEV_UPSTREAM_API_KEY are inherited +unless their corresponding parameters are supplied. +.EXAMPLE +pwsh -File scripts/start-rocm.ps1 -LocalFilesOnly +.EXAMPLE +pwsh -File scripts/start-rocm.ps1 -PythonPath C:\Python313\python.exe -CheckOnly +#> +[CmdletBinding()] +param( + [Alias('Python')] + [ValidateNotNullOrEmpty()] + [string] $PythonPath = '.venv-rocm\Scripts\python.exe', + [string] $Bun, + [ValidateNotNullOrEmpty()] + [string] $Model = 'google/diffusiongemma-26B-A4B-it', + [string] $Revision, + [ValidateRange(1, 65535)] + [int] $BackendPort = 8000, + [ValidateRange(1, 65535)] + [int] $Port = 8080, + [System.Net.IPAddress] $ListenAddress = [System.Net.IPAddress]::Loopback, + [ValidateRange(0, 2147483647)] + [int] $Device = 0, + [ValidateSet('float16', 'bfloat16', 'float32')] + [string] $Dtype = 'float16', + [ValidateRange(1, 2147483647)] + [int] $MaxInputTokens = 8192, + [ValidateRange(1, 2147483647)] + [int] $MaxOutputTokens = 2048, + [ValidateRange(1, 86400)] + [int] $StartupTimeoutSeconds = 1200, + [ValidateRange(1, 86400)] + [int] $TimeoutSeconds = 300, + [string] $ApiKey = $env:LOCALJEV_API_KEY, + [string] $UpstreamApiKey = $env:LOCALJEV_UPSTREAM_API_KEY, + [switch] $LocalFilesOnly, + [switch] $CheckOnly +) + +Set-StrictMode -Version Latest +$ErrorActionPreference = 'Stop' +$repoRoot = Split-Path -Parent $PSScriptRoot + +function Resolve-Executable([string] $Name) { + if ([System.IO.Path]::IsPathRooted($Name) -or $Name.Contains('\') -or $Name.Contains('/')) { + $candidate = if ([System.IO.Path]::IsPathRooted($Name)) { $Name } else { Join-Path $repoRoot $Name } + if (-not (Test-Path -LiteralPath $candidate -PathType Leaf)) { + throw "Executable not found: $candidate. Prepare the ROCm environment first or specify -PythonPath." + } + return (Resolve-Path -LiteralPath $candidate).Path + } + $command = Get-Command -Name $Name -CommandType Application -ErrorAction SilentlyContinue + if (-not $command) { + throw "Executable '$Name' was not found on PATH. Install the required dependency first." + } + return @($command)[0].Source +} + +function Assert-PortAvailable([System.Net.IPAddress] $Address, [int] $Number) { + $listener = [System.Net.Sockets.TcpListener]::new($Address, $Number) + try { + $listener.Server.ExclusiveAddressUse = $true + $listener.Start() + } + catch { + throw "Cannot bind ${Address}:${Number}. Choose another port or stop the existing listener yourself. $($_.Exception.Message)" + } + finally { + $listener.Stop() + } +} + +function ConvertTo-WindowsArgument([string] $Argument) { + # Start-Process joins ArgumentList before passing it to CreateProcess. + # Quote using the Windows C-runtime rules, including trailing backslashes. + $escaped = [regex]::Replace($Argument, '(\\*)"', '$1$1\"') + $escaped = [regex]::Replace($escaped, '(\\+)$', '$1$1') + return '"' + $escaped + '"' +} + +if ([System.Environment]::OSVersion.Platform -ne [System.PlatformID]::Win32NT) { + throw 'This launcher requires Windows. Run scripts/rocm_server.py directly on other platforms.' +} +if ($BackendPort -eq $Port) { + throw '-BackendPort and -Port must be different.' +} +if ([string]::IsNullOrWhiteSpace($Model)) { + throw '-Model must not be blank.' +} + +$python = Resolve-Executable $PythonPath +$bundledBun = Join-Path $repoRoot '.runtime\bun\bun-windows-x64\bun.exe' +if (-not $Bun) { + $Bun = if (Test-Path -LiteralPath $bundledBun -PathType Leaf) { $bundledBun } else { 'bun' } +} +$bunExecutable = Resolve-Executable $Bun +$backendScript = Join-Path $PSScriptRoot 'rocm_server.py' +$bridgeScript = Join-Path $repoRoot 'src\index.ts' +foreach ($requiredFile in @($backendScript, $bridgeScript)) { + if (-not (Test-Path -LiteralPath $requiredFile -PathType Leaf)) { + throw "Required source file not found: $requiredFile" + } +} + +Assert-PortAvailable ([System.Net.IPAddress]::Loopback) $BackendPort +Assert-PortAvailable $ListenAddress $Port + +$backendProcess = $null +$savedEnvironment = @{} +$locationPushed = $false +try { + Push-Location -LiteralPath $repoRoot + $locationPushed = $true + $bunHelp = & $bunExecutable --help 2>&1 | Out-String + if ($LASTEXITCODE -ne 0 -or $bunHelp -notmatch '--no-env-file') { + throw 'This launcher requires a Bun version with --no-env-file support. Update Bun first.' + } + + # Fail before launching either server when the selected Python is incomplete + # or silently resolves to CPU/CUDA PyTorch instead of ROCm. + $preflight = @' +import importlib.util +from importlib.metadata import version +from pathlib import Path +import sys +required = ("torch", "transformers", "accelerate", "PIL") +missing = [name for name in required if importlib.util.find_spec(name) is None] +if missing: + raise SystemExit("Missing Python dependencies: " + ", ".join(missing)) +from packaging.version import Version +transformers_version = version("transformers") +transformers_root = Path(importlib.util.find_spec("transformers").origin).parent +diffusion_module = transformers_root / "models" / "diffusion_gemma" / "modeling_diffusion_gemma.py" +if Version(transformers_version) < Version("5.11.0") or not diffusion_module.is_file(): + raise SystemExit("Transformers 5.11.0 or newer with DiffusionGemma support is required.") +if "class DiffusionGemmaForBlockDiffusion(" not in diffusion_module.read_text(encoding="utf-8"): + raise SystemExit("The selected Transformers build has no DiffusionGemmaForBlockDiffusion model.") +import torch +if not torch.version.hip: + raise SystemExit("The selected Python does not have a ROCm PyTorch build.") +device = int(sys.argv[1]) +if not torch.cuda.is_available() or device >= torch.cuda.device_count(): + raise SystemExit("The selected ROCm GPU device is not available: " + str(device)) +props = torch.cuda.get_device_properties(device) +architecture = str(getattr(props, "gcnArchName", "unknown architecture")) +if architecture.split(":")[0] != "gfx1151": + raise SystemExit("This launcher targets gfx1151; selected GPU architecture: " + architecture) +print("Python: " + sys.executable) +print("ROCm: " + str(torch.version.hip)) +print("Transformers: " + transformers_version) +print("GPU: " + props.name + " (" + architecture + ")") +'@ + # Stdin avoids PowerShell 5.1's legacy native argument handling stripping + # embedded quotes from a multiline Python -c argument. + $preflight | & $python - $Device + if ($LASTEXITCODE -ne 0) { + throw 'ROCm preflight failed. See the dependency/GPU error above.' + } + if ($CheckOnly) { + Write-Host 'Launcher checks passed. No model was loaded and no server was started.' + return + } + + $bridgeEnvironment = @{ + LOCALJEV_UPSTREAM = "http://127.0.0.1:$BackendPort" + LOCALJEV_UPSTREAM_MODEL = $Model + LOCALJEV_HOST = $ListenAddress.ToString() + LOCALJEV_PORT = "$Port" + LOCALJEV_MAX_INFLIGHT = '1' + LOCALJEV_TIMEOUT = "$TimeoutSeconds" + LOCALJEV_MAX_OUTPUT_TOKENS = "$MaxOutputTokens" + LOCALJEV_API_KEY = $ApiKey + LOCALJEV_UPSTREAM_API_KEY = $UpstreamApiKey + } + foreach ($entry in $bridgeEnvironment.GetEnumerator()) { + $savedEnvironment[$entry.Key] = [System.Environment]::GetEnvironmentVariable($entry.Key, 'Process') + [System.Environment]::SetEnvironmentVariable($entry.Key, $entry.Value, 'Process') + } + $healthHeaders = @{} + if ($UpstreamApiKey) { $healthHeaders.Authorization = "Bearer $UpstreamApiKey" } + + $runtimeDirectory = Join-Path $repoRoot '.runtime' + $null = New-Item -ItemType Directory -Path $runtimeDirectory -Force + $runId = '{0}-{1}-{2}' -f (Get-Date -Format 'yyyyMMdd-HHmmss'), $PID, ([guid]::NewGuid().ToString('N').Substring(0, 8)) + $stdoutLog = Join-Path $runtimeDirectory "rocm-$runId.stdout.log" + $stderrLog = Join-Path $runtimeDirectory "rocm-$runId.stderr.log" + $backendArguments = @( + '-u', $backendScript, + '--model', $Model, + '--host', '127.0.0.1', + '--port', "$BackendPort", + '--device', "$Device", + '--dtype', $Dtype, + '--max-input-tokens', "$MaxInputTokens", + '--max-output-tokens', "$MaxOutputTokens", + '--max-queue', '2' + ) + if ($LocalFilesOnly) { $backendArguments += '--local-files-only' } + if ($Revision) { $backendArguments += @('--revision', $Revision) } + $argumentLine = ($backendArguments | ForEach-Object { ConvertTo-WindowsArgument $_ }) -join ' ' + $backendProcess = Start-Process -FilePath $python -ArgumentList $argumentLine -WorkingDirectory $repoRoot ` + -WindowStyle Hidden -RedirectStandardOutput $stdoutLog -RedirectStandardError $stderrLog -PassThru + $null = $backendProcess.Handle + Write-Host "ROCm backend PID $($backendProcess.Id); logs: $stdoutLog and $stderrLog" + Write-Host "Waiting up to $StartupTimeoutSeconds seconds for model '$Model'..." + + $timer = [System.Diagnostics.Stopwatch]::StartNew() + $healthy = $false + while ($timer.Elapsed.TotalSeconds -lt $StartupTimeoutSeconds) { + $backendProcess.Refresh() + if ($backendProcess.HasExited) { + throw "ROCm backend exited with code $($backendProcess.ExitCode). Inspect $stderrLog and $stdoutLog." + } + try { + $remainingSeconds = [Math]::Max(1, [Math]::Ceiling($StartupTimeoutSeconds - $timer.Elapsed.TotalSeconds)) + $health = Invoke-RestMethod -Uri "http://127.0.0.1:$BackendPort/health" -Headers $healthHeaders -TimeoutSec ([Math]::Min(5, $remainingSeconds)) + if ($health.status -eq 'ok' -and $health.model -eq $Model) { + $healthy = $true + break + } + } + catch { + # The backend can be unreachable until the model finishes loading. + } + Start-Sleep -Milliseconds 500 + } + if (-not $healthy) { + throw "ROCm model startup timed out after $StartupTimeoutSeconds seconds. Inspect $stderrLog and $stdoutLog." + } + $backendProcess.Refresh() + if ($backendProcess.HasExited) { + throw "ROCm backend exited after its health check. Inspect $stderrLog and $stdoutLog." + } + + Write-Host "Starting LocalJev at http://${ListenAddress}:$Port; press Ctrl+C to stop both servers." + & $bunExecutable --no-env-file run $bridgeScript + if ($LASTEXITCODE -ne 0) { + throw "LocalJev exited with code $LASTEXITCODE." + } +} +finally { + foreach ($entry in $savedEnvironment.GetEnumerator()) { + [System.Environment]::SetEnvironmentVariable($entry.Key, $entry.Value, 'Process') + } + if ($null -ne $backendProcess) { + $backendProcess.Refresh() + if (-not $backendProcess.HasExited) { + # Windows venv python.exe can launch the actual interpreter as a + # child. Stop the owned tree so that worker cannot outlive LocalJev. + try { + if ($PSVersionTable.PSVersion.Major -ge 7) { + $backendProcess.Kill($true) + } + else { + & "$env:SystemRoot\System32\taskkill.exe" /PID $backendProcess.Id /T /F | Out-Null + } + } + catch [System.InvalidOperationException] { + # The process can exit between HasExited and Kill. + } + if (-not $backendProcess.WaitForExit(10000)) { + Write-Warning "Unable to stop the owned ROCm backend tree (PID $($backendProcess.Id))." + } + } + $backendProcess.Dispose() + } + if ($locationPushed) { Pop-Location } +} diff --git a/src/index.ts b/src/index.ts index d2030fc..362c510 100644 --- a/src/index.ts +++ b/src/index.ts @@ -7,15 +7,29 @@ export { Engine } from "./engine"; export { LocalJevApp } from "./server"; export * from "./types"; +export function serveLocalJev(app: LocalJevApp, idleTimeout = 255) { + return Bun.serve({ + hostname: app.settings.host, + port: app.settings.port, + idleTimeout, + fetch(request, server) { + if ( + request.method === "POST" && + new URL(request.url).pathname === "/v1/systemone" + ) { + // Inference and validation retries can exceed Bun's idle timeout. + // The engine enforces the configured timeout on each upstream call. + server.timeout(request, 0); + } + return app.fetch(request); + }, + }); +} + if (import.meta.main) { const settings = loadSettings(); const app = new LocalJevApp(settings, new Engine(settings)); - const server = Bun.serve({ - hostname: settings.host, - port: settings.port, - idleTimeout: 255, - fetch: (request) => app.fetch(request), - }); + const server = serveLocalJev(app); console.log(`LocalJev listening on ${server.url}`); diff --git a/test/http-smoke.test.ts b/test/http-smoke.test.ts new file mode 100644 index 0000000..fad61ab --- /dev/null +++ b/test/http-smoke.test.ts @@ -0,0 +1,121 @@ +import { describe, expect, test } from "bun:test"; +import { fileURLToPath } from "node:url"; + +import { runHttpSmoke } from "../scripts/http-smoke"; +import { loadSettings } from "../src/config"; +import type { DecisionEngine } from "../src/engine"; +import { LocalJevApp } from "../src/server"; + +const validResponse = { + model: "localjev-0.2", + answers: { + department: { + type: "choice", + choice: "technical", + probabilities: { billing: 0.1, technical: 0.8, sales: 0.1 }, + confidence: 0.4, + }, + frustration: { + type: "score", + score: 0.9, + legend: { "0": "Calm", "1": "Frustrated but civil", "2": "Very angry" }, + probabilities: { "0": 0.2, "1": 0.7, "2": 0.1 }, + confidence: 0.2, + }, + urgent: { type: "noul", noul: 0.3 }, + }, + usage: { input_tokens: 120, output_tokens: 40 }, +}; + +describe("HTTP smoke", () => { + test("calls a listening authenticated LocalJev server with all question types and a stable seed", async () => { + const seeds: number[] = []; + const engine: DecisionEngine = { + async ready() { return true; }, + async decide(questions, _state, seed) { + expect(Object.values(questions).map((question) => question.type)).toEqual(["choice", "score", "noul"]); + seeds.push(seed); + return { + answers: validResponse.answers as Awaited>["answers"], + inputTokens: 120, + outputTokens: 40, + }; + }, + }; + const app = new LocalJevApp(loadSettings({ apiKey: "smoke-test-key", upstreamModel: "test-model" }), engine); + const server = Bun.serve({ hostname: "127.0.0.1", port: 0, fetch: (request) => app.fetch(request) }); + try { + const options = { url: server.url.href, apiKey: "smoke-test-key" }; + const evidence = await runHttpSmoke(options); + expect(evidence).toMatchObject({ + status: "ok", upstream_model: "test-model", ...validResponse, + }); + expect(evidence.latency_ms).toBeGreaterThanOrEqual(0); + await runHttpSmoke(options); + expect(seeds).toHaveLength(2); + expect(seeds[0]).toBe(seeds[1]); + } finally { + await server.stop(true); + await app.close(); + } + }); + + test("fails before inference when readiness is unavailable", async () => { + const paths: string[] = []; + const server = Bun.serve({ + hostname: "127.0.0.1", port: 0, + fetch(request) { + paths.push(new URL(request.url).pathname); + return Response.json({ status: "unavailable" }, { status: 503 }); + }, + }); + try { + await expect(runHttpSmoke({ url: server.url.href })).rejects.toThrow("ready returned HTTP 503"); + expect(paths).toEqual(["/ready"]); + } finally { + await server.stop(true); + } + }); + + test.each([ + ["incorrect question type", (body: typeof validResponse) => { body.answers.urgent.type = "score"; }, "answers.urgent.type"], + ["probability sum", (body: typeof validResponse) => { body.answers.department.probabilities.billing = 0.9; }, "must sum to 1"], + ["noul bounds", (body: typeof validResponse) => { body.answers.urgent.noul = 1.1; }, "answers.urgent.noul"], + ["token counts", (body: typeof validResponse) => { body.usage.input_tokens = -1; }, "usage.input_tokens"], + ] as const)("rejects invalid %s over HTTP", async (_name, mutate, expected) => { + const body = structuredClone(validResponse); + mutate(body); + const server = Bun.serve({ + hostname: "127.0.0.1", port: 0, + fetch(request) { + return Response.json(new URL(request.url).pathname === "/ready" + ? { status: "ready", upstream_model: "test-model" } + : body); + }, + }); + try { + await expect(runHttpSmoke({ url: server.url.href })).rejects.toThrow(expected); + } finally { + await server.stop(true); + } + }); + + test("CLI exits nonzero when readiness fails", async () => { + const server = Bun.serve({ + hostname: "127.0.0.1", port: 0, + fetch: () => Response.json({ status: "unavailable" }, { status: 503 }), + }); + try { + const child = Bun.spawn([process.execPath, "run", "scripts/http-smoke.ts"], { + cwd: fileURLToPath(new URL("..", import.meta.url)), + env: { ...process.env, LOCALJEV_URL: server.url.href }, + stdout: "pipe", stderr: "pipe", + }); + expect(await child.exited).toBe(1); + expect(await new Response(child.stderr).text()).toContain("HTTP smoke failed: ready returned HTTP 503"); + expect(await new Response(child.stdout).text()).toBe(""); + } finally { + await server.stop(true); + } + }); +}); diff --git a/test/index.test.ts b/test/index.test.ts new file mode 100644 index 0000000..ba9095e --- /dev/null +++ b/test/index.test.ts @@ -0,0 +1,40 @@ +import { expect, test } from "bun:test"; + +import { loadSettings, LocalJevApp, serveLocalJev } from "../src/index"; +import type { DecisionEngine } from "../src/engine"; + +test("inference can outlast the HTTP idle timeout", async () => { + const engine: DecisionEngine = { + async decide() { + // Bun checks a one-second idle limit on a roughly four-second timer. + await Bun.sleep(5_500); + return { + answers: { urgent: { type: "noul", noul: 0.9 } }, + inputTokens: 10, + outputTokens: 5, + }; + }, + }; + const app = new LocalJevApp( + loadSettings({ host: "127.0.0.1", port: 0, apiKey: "" }), + engine, + ); + const server = serveLocalJev(app, 1); + try { + const response = await fetch(new URL("v1/systemone", server.url), { + method: "POST", + headers: { "content-type": "application/json" }, + body: JSON.stringify({ + model: "localjev-latest", + state: "A production service is down.", + questions: { urgent: { type: "noul", instructions: "Is this urgent?" } }, + }), + signal: AbortSignal.timeout(10_000), + }); + expect(response.status).toBe(200); + expect((await response.json()).answers.urgent).toEqual({ type: "noul", noul: 0.9 }); + } finally { + await server.stop(true); + await app.close(); + } +}, 12_000); diff --git a/test/test_rocm_server.py b/test/test_rocm_server.py new file mode 100644 index 0000000..e8650aa --- /dev/null +++ b/test/test_rocm_server.py @@ -0,0 +1,204 @@ +"""CPU-only HTTP/queue tests; never import torch or download model weights.""" + +import copy +import http.client +import json +from pathlib import Path +import select +import sys +import threading +import unittest + +sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "scripts")) +from rocm_server import ( # noqa: E402 + Generation, InferenceService, Limits, RequestError, + final_text, make_server, parse_request, trim_generated_tokens, +) + + +class FakeRuntime: + model_id = "test-model" + info = {"architecture": "gfx1151", "device": "cuda:0"} + + def __init__(self): + self.requests = [] + self.failure = False + + def generate(self, request): + self.requests.append(request) + if self.failure: + raise RuntimeError("private backend error") + return Generation('{"billing":0.9}', 21, 8, "stop") + + +def request_body(**overrides): + body = {"model": "test-model", "messages": [{"role": "user", "content": "Charged twice"}], "max_tokens": 32, "seed": 1234} + body.update(overrides) + return body + + +class ServerTests(unittest.TestCase): + def setUp(self): + self.runtime = FakeRuntime() + self.service = InferenceService(self.runtime, Limits(max_body_bytes=4096)) + self.server = make_server(self.service, "127.0.0.1", 0, "test-key") + self.thread = threading.Thread(target=self.server.serve_forever, daemon=True) + self.thread.start() + + def tearDown(self): + self.server.shutdown() + self.server.server_close() + self.thread.join(timeout=2) + + def call(self, method, path, payload=None, key="test-key", content_type="application/json"): + connection = http.client.HTTPConnection(*self.server.server_address, timeout=3) + headers = {"Authorization": f"Bearer {key}", "Content-Type": content_type} + data = payload if isinstance(payload, (str, bytes)) else json.dumps(payload) if payload is not None else None + connection.request(method, path, body=data, headers=headers) + response = connection.getresponse() + status = response.status + result = json.loads(response.read()) + connection.close() + return status, result + + def test_health_and_model_inventory(self): + status, health = self.call("GET", "/health") + self.assertEqual(status, 200) + self.assertEqual(health["runtime"]["architecture"], "gfx1151") + self.assertEqual(health["structured_output"], "prompt_only") + status, models = self.call("GET", "/v1/models") + self.assertEqual(status, 200) + self.assertEqual(models["data"][0]["id"], "test-model") + + def test_localjev_request_and_openai_response(self): + schema = {"type": "object", "properties": {"billing": {"type": "number"}}, "required": ["billing"]} + body = request_body(temperature=0.0, chat_template_kwargs={"enable_thinking": False}, response_format={"type": "json_schema", "json_schema": {"name": "decision", "strict": True, "schema": schema}}) + original = copy.deepcopy(body) + status, response = self.call("POST", "/v1/chat/completions", body) + self.assertEqual(status, 200) + self.assertEqual(response["model"], "test-model") + self.assertEqual(response["choices"][0]["message"]["content"], '{"billing":0.9}') + self.assertEqual(response["choices"][0]["finish_reason"], "stop") + self.assertEqual(response["usage"], {"prompt_tokens": 21, "completion_tokens": 8, "total_tokens": 29}) + self.assertEqual(self.runtime.requests[0].seed, 1234) + self.assertIn(json.dumps(schema), self.runtime.requests[0].messages[0]["content"]) + self.assertEqual(body, original) + + def test_authentication_and_model_mismatch_do_not_generate(self): + self.assertEqual(self.call("GET", "/health", key="wrong")[0], 401) + self.assertEqual(self.call("POST", "/v1/chat/completions", request_body(model="other"))[0], 404) + self.assertEqual(self.runtime.requests, []) + + def test_invalid_transport_and_json_do_not_generate(self): + cases = [("not-json", "application/json", 400), ('{"x":NaN}', "application/json", 400), ("x" * 4097, "application/json", 413), ("{}", "text/plain", 415)] + for payload, content_type, expected in cases: + with self.subTest(expected=expected, payload=payload[:20]): + self.assertEqual(self.call("POST", "/v1/chat/completions", payload, content_type=content_type)[0], expected) + self.assertEqual(self.runtime.requests, []) + + def test_media_type_rejection_consumes_split_body_without_connection_reset(self): + for attempt in range(20): + with self.subTest(attempt=attempt): + connection = http.client.HTTPConnection(*self.server.server_address, timeout=3) + try: + connection.putrequest("POST", "/v1/chat/completions") + connection.putheader("Authorization", "Bearer test-key") + connection.putheader("Content-Type", "text/plain") + connection.putheader("Content-Length", "2") + connection.endheaders(b"{") + # The old handler closed here without consuming the body; + # sending the trailing byte could reset its error response. + readable, _, _ = select.select([connection.sock], [], [], 0.01) + self.assertEqual(readable, [], "Server rejected before reading the complete bounded body") + connection.send(b"}") + response = connection.getresponse() + self.assertEqual(response.status, 415) + self.assertIn("Content-Type", json.loads(response.read())["error"]["message"]) + finally: + connection.close() + self.assertEqual(self.runtime.requests, []) + + def test_backend_errors_are_redacted_and_capacity_is_released(self): + self.runtime.failure = True + with self.assertLogs("localjev.rocm", level="ERROR"): + status, response = self.call("POST", "/v1/chat/completions", request_body()) + self.assertEqual(status, 500) + self.assertNotIn("private backend error", json.dumps(response)) + self.runtime.failure = False + self.assertEqual(self.call("POST", "/v1/chat/completions", request_body())[0], 200) + + +class ValidationTests(unittest.TestCase): + def test_rejects_unsupported_or_unbounded_requests(self): + cases = [ + {"stream": True}, {"n": 2}, {"max_tokens": 0}, {"max_tokens": True}, + {"max_tokens": 4096}, {"seed": -1}, {"temperature": float("nan")}, + {"temperature": True}, {"messages": []}, {"messages": [{"role": "user", "content": [{"type": "image_url"}]}]}, + {"messages": [{"role": "tool", "content": "data"}]}, {"tools": []}, + {"chat_template_kwargs": {"enable_thinking": True}}, + {"response_format": {"type": "json_schema", "json_schema": {"schema": []}}}, + ] + for overrides in cases: + with self.subTest(overrides=overrides), self.assertRaises(RequestError): + parse_request(request_body(**overrides), "test-model", Limits()) + + def test_preserves_existing_system_prompt_without_mutating_input(self): + body = request_body(messages=[{"role": "system", "content": "Classify."}, {"role": "user", "content": "Hi"}], response_format={"type": "json_object"}) + result = parse_request(body, "test-model", Limits()) + self.assertTrue(result.messages[0]["content"].startswith("Classify.")) + self.assertEqual(body["messages"][0]["content"], "Classify.") + self.assertEqual(len(result.messages), 2) + + def test_diffusion_canvas_respects_output_budget_and_eos(self): + self.assertEqual(trim_generated_tokens([4, 5, 6, 1, 0, 0], [1], 3), ([4, 5, 6], "length")) + self.assertEqual(trim_generated_tokens([4, 5, 1, 0, 0], [1], 3), ([4, 5, 1], "stop")) + self.assertEqual(trim_generated_tokens([4, 1, 0, 0], [1], 10), ([4, 1], "stop")) + + def test_reasoning_markers_do_not_contaminate_probability_json(self): + self.assertEqual(final_text('analysis\n{"p": 0.9}'), '{"p": 0.9}') + self.assertEqual(final_text('<|channel>thought\n{"p": 0.9}'), '{"p": 0.9}') + self.assertEqual(final_text('{"p": 0.9}'), '{"p": 0.9}') + + def test_full_admission_queue_rejects_without_concurrent_generation(self): + runtime = FakeRuntime() + started = threading.Event() + release = threading.Event() + original_generate = runtime.generate + + def blocked_generate(request): + started.set() + if not release.wait(timeout=3): + raise RuntimeError("test timed out") + return original_generate(request) + + runtime.generate = blocked_generate + service = InferenceService(runtime, Limits(max_queue=0)) + worker = threading.Thread(target=lambda: service.complete(request_body())) + worker.start() + try: + self.assertTrue(started.wait(timeout=2)) + with self.assertRaises(RequestError) as captured: + service.complete(request_body()) + self.assertEqual(captured.exception.status, 429) + finally: + release.set() + worker.join(timeout=3) + self.assertFalse(worker.is_alive()) + self.assertEqual(len(runtime.requests), 1) + service.complete(request_body()) + self.assertEqual(len(runtime.requests), 2) + + def test_queue_timeout_releases_admission(self): + service = InferenceService(FakeRuntime(), Limits(max_queue=0, queue_timeout=0.01)) + service.inference.acquire() + try: + with self.assertRaises(RequestError) as captured: + service.complete(request_body()) + self.assertEqual(captured.exception.status, 429) + finally: + service.inference.release() + self.assertEqual(service.complete(request_body())["object"], "chat.completion") + + +if __name__ == "__main__": + unittest.main()