| """Prefill-only scoring through stock SGLang HTTP APIs. |
| |
| Chat tokenization uses the inference server by default. An optional native CPU |
| tokenizer preserves control-token IDs for engines with a lossy string chat API. |
| No SGLang, CUDA, or model weights are imported into this process. |
| """ |
|
|
| import asyncio |
| import math |
| import uuid |
| from typing import Any |
|
|
| import httpx |
|
|
| from .backend import AdapterError, ScoringResult |
| from .labels import ( |
| DEFAULT_LABEL_SCHEME, |
| MAX_LABELS, |
| candidate_tiers, |
| validate_label_scheme, |
| ) |
|
|
|
|
| _IMAGE_MODEL_TYPES = { |
| "mage_vl", |
| "qwen2_vl", |
| "qwen2_5_vl", |
| "qwen3_vl", |
| "qwen3_vl_moe", |
| "qwen3_5", |
| "qwen3_5_moe", |
| "mistral3", |
| "pixtral", |
| } |
| _IMAGE_ARCHITECTURES = { |
| "MageVLForConditionalGeneration", |
| "Qwen2VLForConditionalGeneration", |
| "Qwen2_5_VLForConditionalGeneration", |
| "Qwen3VLForConditionalGeneration", |
| "Qwen3VLMoeForConditionalGeneration", |
| "Qwen3_5ForConditionalGeneration", |
| "Qwen3_5MoeForConditionalGeneration", |
| "Mistral3ForConditionalGeneration", |
| "PixtralForConditionalGeneration", |
| } |
|
|
|
|
| def _bad_response(message: str) -> AdapterError: |
| return AdapterError("invalid_engine_response", message, status=502) |
|
|
|
|
| def _token_list(value: Any) -> bool: |
| return isinstance(value, list) and all( |
| type(token) is int and token >= 0 for token in value |
| ) |
|
|
|
|
| class SGLangBackend: |
| def __init__( |
| self, |
| base_url: str, |
| model: str, |
| api_key: str | None = None, |
| timeout: float = 120, |
| max_connections: int = 32, |
| client: httpx.AsyncClient | None = None, |
| native_tokenizer: Any | None = None, |
| prompt_wording: str = "served", |
| native_system_prompt: str | None = None, |
| label_scheme: str = DEFAULT_LABEL_SCHEME, |
| ): |
| if not math.isfinite(timeout) or timeout <= 0 or max_connections < 1: |
| raise ValueError("timeout and max_connections must be positive") |
| if prompt_wording not in ("served", "native"): |
| raise ValueError('prompt_wording must be "served" or "native"') |
| |
| self.label_scheme = validate_label_scheme(label_scheme) |
| url = httpx.URL(base_url) |
| if ( |
| url.scheme not in ("http", "https") |
| or not url.host |
| or url.query |
| or url.fragment |
| ): |
| raise ValueError( |
| "base_url must be an HTTP(S) origin or path without query or fragment" |
| ) |
| self.base_url = base_url.rstrip("/") |
| self.model = model |
| self.timeout = timeout |
| self._headers = {"Authorization": f"Bearer {api_key}"} if api_key else {} |
| self._owns_client = client is None |
| self._client = client or httpx.AsyncClient( |
| limits=httpx.Limits(max_connections=max_connections), |
| timeout=timeout, |
| ) |
| self._startup_lock = asyncio.Lock() |
| self._label_strings: tuple[str, ...] = () |
| self._label_ids: tuple[int, ...] = () |
| self._image_capable = False |
| self._image_format_supported = False |
| self._ready = False |
| self._native_tokenizer = native_tokenizer |
| self._prompt_wording = prompt_wording |
| |
| |
| |
| |
| |
| self._native_system_prompt = native_system_prompt if prompt_wording == "native" else None |
|
|
| async def _request( |
| self, method: str, path: str, payload: dict | None = None |
| ) -> dict: |
| try: |
| response = await self._client.request( |
| method, |
| self.base_url + path, |
| json=payload, |
| headers=self._headers, |
| timeout=self.timeout, |
| ) |
| except httpx.TimeoutException as error: |
| raise AdapterError( |
| "engine_timeout", "The inference engine timed out.", status=504 |
| ) from error |
| except httpx.HTTPError as error: |
| raise AdapterError( |
| "engine_unavailable", |
| "Could not reach the inference engine.", |
| status=502, |
| ) from error |
| if response.status_code == 404: |
| raise AdapterError( |
| "engine_endpoint_missing", |
| f"The inference engine does not provide {path}.", |
| status=502, |
| ) |
| if response.status_code in (400, 422): |
| raise AdapterError( |
| "engine_rejected_request", |
| "The inference engine rejected the scoring input.", |
| status=422, |
| ) |
| if not response.is_success: |
| raise AdapterError( |
| "engine_http_error", |
| f"The inference engine returned HTTP {response.status_code}.", |
| status=502, |
| ) |
| try: |
| result = response.json() |
| except ValueError as error: |
| raise _bad_response( |
| "The inference engine returned invalid JSON." |
| ) from error |
| if not isinstance(result, dict) or "error" in result: |
| raise _bad_response("The inference engine returned an unexpected response.") |
| return result |
|
|
| async def _information(self, path: str, legacy: str) -> dict: |
| try: |
| return await self._request("GET", path) |
| except AdapterError as error: |
| if error.code != "engine_endpoint_missing": |
| raise |
| return await self._request("GET", legacy) |
|
|
| async def start(self) -> None: |
| async with self._startup_lock: |
| if self._ready: |
| return |
| info = await self._information("/model_info", "/get_model_info") |
| if info.get("is_generation") is not True: |
| raise AdapterError( |
| "unsupported_model", "A generation model is required." |
| ) |
| served = info.get("served_model_name") |
| if not isinstance(served, str) or not served: |
| raise _bad_response( |
| "The inference engine did not identify its served model." |
| ) |
| if self.model not in (served, info.get("model_path")): |
| raise AdapterError( |
| "model_not_found", |
| "Configured model is not served by the inference engine.", |
| "model", |
| 404, |
| ) |
| config = await self._information("/server_info", "/get_server_info") |
| if config.get("skip_tokenizer_init"): |
| raise AdapterError( |
| "tokenizer_unavailable", |
| "The inference engine must initialize its tokenizer.", |
| ) |
| if config.get("speculative_algorithm") is not None: |
| raise AdapterError( |
| "speculation_not_supported", |
| "Disable speculative decoding and MTP for prefill-only scoring.", |
| ) |
| self.model = served |
| self._image_capable = info.get("has_image_understanding") is True |
| architectures = info.get("architectures") or [] |
| self._image_format_supported = info.get( |
| "model_type" |
| ) in _IMAGE_MODEL_TYPES or ( |
| isinstance(architectures, list) |
| and any( |
| arch in _IMAGE_ARCHITECTURES |
| for arch in architectures |
| if isinstance(arch, str) |
| ) |
| ) |
| valid: list[tuple[str, int]] = [] |
| seen: set[int] = set() |
| |
| for candidates in candidate_tiers(self.label_scheme): |
| if len(valid) == MAX_LABELS: |
| break |
| await self._select_tier(list(candidates), valid, seen) |
| if len(valid) < 2: |
| raise AdapterError( |
| "labels_unavailable", |
| "At least two distinct, non-special single-token labels are required.", |
| ) |
| self._label_strings, self._label_ids = tuple(zip(*valid)) |
| self._ready = True |
|
|
| async def _select_tier( |
| self, candidates: list[str], valid: list[tuple[str, int]], seen: set[int] |
| ) -> None: |
| """Append this group's reversible single-token labels, in order.""" |
| tokenized = await self._request( |
| "POST", |
| "/v1/tokenize", |
| { |
| "model": self.model, |
| "prompt": candidates, |
| "add_special_tokens": False, |
| }, |
| ) |
| batches = tokenized.get("tokens") |
| if ( |
| not isinstance(batches, list) |
| or len(batches) != len(candidates) |
| or not all(_token_list(ids) for ids in batches) |
| ): |
| raise _bad_response( |
| "The inference engine returned malformed label tokenization." |
| ) |
| singles = [ |
| (label, ids[0]) for label, ids in zip(candidates, batches) if len(ids) == 1 |
| ] |
| if not singles: |
| return |
| decoded = await self._request( |
| "POST", |
| "/v1/detokenize", |
| { |
| "model": self.model, |
| "tokens": [[token] for _, token in singles], |
| "skip_special_tokens": True, |
| }, |
| ) |
| texts = decoded.get("text") |
| if ( |
| not isinstance(texts, list) |
| or len(texts) != len(singles) |
| or not all(isinstance(item, str) for item in texts) |
| ): |
| raise _bad_response( |
| "The inference engine returned malformed label detokenization." |
| ) |
| for (label, token), decoded_label in zip(singles, texts): |
| if label == decoded_label and token not in seen: |
| valid.append((label, token)) |
| seen.add(token) |
| if len(valid) == MAX_LABELS: |
| return |
|
|
| async def close(self) -> None: |
| if self._owns_client: |
| await self._client.aclose() |
|
|
| def labels(self, count: int) -> tuple[tuple[str, ...], tuple[int, ...]]: |
| if not self._ready: |
| raise AdapterError( |
| "engine_not_ready", "The scoring backend has not started.", status=503 |
| ) |
| if type(count) is not int or count < 2 or count > len(self._label_strings): |
| raise AdapterError( |
| "labels_unavailable", |
| f"This tokenizer supports 2 to {len(self._label_strings)} answer labels.", |
| "questions", |
| ) |
| return self._label_strings[:count], self._label_ids[:count] |
|
|
| async def evaluate( |
| self, |
| prompt: str, |
| images: list[str], |
| labels: tuple[str, ...], |
| token_ids: tuple[int, ...], |
| assistant_prefix: str | None, |
| ) -> ScoringResult: |
| if not self._ready: |
| raise AdapterError( |
| "engine_not_ready", "The scoring backend has not started.", status=503 |
| ) |
| if ( |
| len(labels) < 2 |
| or len(labels) != len(token_ids) |
| or len(set(token_ids)) != len(token_ids) |
| ): |
| raise AdapterError( |
| "invalid_labels", "Distinct label/token pairs are required." |
| ) |
| known = dict(zip(self._label_strings, self._label_ids)) |
| if any(known.get(label) != token for label, token in zip(labels, token_ids)): |
| raise AdapterError( |
| "invalid_labels", "Labels must come from this backend's tokenizer." |
| ) |
| if images and self._native_tokenizer is not None: |
| raise AdapterError( |
| "native_images_not_supported", |
| "Native tokenizer scoring currently accepts text only.", |
| "images", |
| ) |
| if images and not self._image_capable: |
| raise AdapterError( |
| "images_not_supported", |
| "The loaded model does not support images.", |
| "images", |
| ) |
| if images and not self._image_format_supported: |
| raise AdapterError( |
| "image_model_not_supported", |
| "Image scoring currently supports Mage-VL, Qwen VL/Qwen3.5, and Mistral3/Pixtral chat formats.", |
| "images", |
| ) |
| if self._native_tokenizer is not None: |
| final_ids = await asyncio.to_thread( |
| self._native_tokenizer.prepare, |
| prompt, |
| labels, |
| token_ids, |
| assistant_prefix, |
| self._native_system_prompt, |
| ) |
| if not _token_list(final_ids) or not final_ids: |
| raise _bad_response("The native tokenizer returned invalid prompt IDs.") |
| final_text = None |
| else: |
| final_ids, final_text = await self._prepare_http( |
| prompt, |
| images, |
| labels, |
| token_ids, |
| assistant_prefix, |
| self._native_system_prompt, |
| ) |
| payload = { |
| "sampling_params": { |
| "max_new_tokens": 0, |
| "temperature": 1, |
| "top_p": 1, |
| "top_k": -1, |
| }, |
| "return_logprob": True, |
| "logprob_start_len": -1, |
| "token_ids_logprob": list(token_ids), |
| "top_logprobs_num": 0, |
| "return_text_in_logprobs": False, |
| "stream": False, |
| "rid": "jev-" + uuid.uuid4().hex, |
| } |
| if images: |
| payload.update(text=final_text, image_data=images) |
| else: |
| payload["input_ids"] = final_ids |
| result = await self._request("POST", "/generate", payload) |
| meta = result.get("meta_info") |
| if not isinstance(meta, dict): |
| raise _bad_response("The inference engine omitted scoring metadata.") |
| if ( |
| type(meta.get("completion_tokens")) is not int |
| or meta["completion_tokens"] != 0 |
| ): |
| raise _bad_response( |
| "The inference engine did not honor zero-output scoring." |
| ) |
| input_tokens = meta.get("prompt_tokens") |
| if type(input_tokens) is not int or input_tokens < 1: |
| raise _bad_response( |
| "The inference engine returned invalid input-token usage." |
| ) |
| positions = meta.get("output_token_ids_logprobs") |
| if ( |
| not isinstance(positions, list) |
| or len(positions) != 1 |
| or not isinstance(positions[0], list) |
| ): |
| raise _bad_response( |
| "The inference engine omitted the requested next-token logprobs." |
| ) |
| values: dict[int, float] = {} |
| for entry in positions[0]: |
| if not isinstance(entry, (list, tuple)) or len(entry) < 2: |
| raise _bad_response("The inference engine returned malformed logprobs.") |
| value, token = entry[:2] |
| if ( |
| type(token) is not int |
| or type(value) not in (int, float) |
| or not math.isfinite(value) |
| or value > 0 |
| or token in values |
| ): |
| raise _bad_response("The inference engine returned invalid logprobs.") |
| values[token] = float(value) |
| if set(values) != set(token_ids): |
| raise _bad_response( |
| "The inference engine did not return every requested label logprob." |
| ) |
| return ScoringResult(tuple(values[token] for token in token_ids), input_tokens) |
|
|
| async def _prepare_http( |
| self, |
| prompt: str, |
| images: list[str], |
| labels: tuple[str, ...], |
| token_ids: tuple[int, ...], |
| assistant_prefix: str | None, |
| system_prompt: str | None = None, |
| ) -> tuple[list[int], str]: |
| content: Any = prompt |
| if images: |
| content = [ |
| {"type": "image_url", "image_url": {"url": image}} for image in images |
| ] |
| content.append({"type": "text", "text": prompt}) |
| messages = ( |
| [{"role": "system", "content": system_prompt}] if system_prompt else [] |
| ) + [{"role": "user", "content": content}] |
| rendered = await self._request( |
| "POST", |
| "/v1/tokenize", |
| { |
| "model": self.model, |
| "messages": messages, |
| "chat_template_kwargs": {"enable_thinking": False}, |
| }, |
| ) |
| prompt_ids = rendered.get("tokens") |
| if not _token_list(prompt_ids) or not prompt_ids: |
| raise _bad_response( |
| "The inference engine returned malformed chat tokenization." |
| ) |
| decoded = await self._request( |
| "POST", |
| "/v1/detokenize", |
| { |
| "model": self.model, |
| "tokens": prompt_ids, |
| "skip_special_tokens": False, |
| }, |
| ) |
| rendered_text = decoded.get("text") |
| if not isinstance(rendered_text, str) or not rendered_text: |
| raise _bad_response( |
| "The inference engine returned malformed chat detokenization." |
| ) |
| final_text = rendered_text + (assistant_prefix or "") |
| |
| |
| prompts = [rendered_text, final_text] + [final_text + label for label in labels] |
| checked = await self._request( |
| "POST", |
| "/v1/tokenize", |
| { |
| "model": self.model, |
| "prompt": prompts, |
| "add_special_tokens": False, |
| }, |
| ) |
| variants = checked.get("tokens") |
| if ( |
| not isinstance(variants, list) |
| or len(variants) != len(prompts) |
| or not all(_token_list(ids) for ids in variants) |
| ): |
| raise _bad_response( |
| "The inference engine returned malformed boundary tokenization." |
| ) |
| if variants[0] != prompt_ids: |
| raise AdapterError( |
| "unsupported_tokenizer_roundtrip", |
| "The rendered chat prompt cannot be losslessly tokenized through this HTTP adapter.", |
| ) |
| final_ids = variants[1] |
| if not final_ids: |
| raise _bad_response("The formatted scoring prompt is empty.") |
| if any( |
| ids != final_ids + [token] for ids, token in zip(variants[2:], token_ids) |
| ): |
| raise AdapterError( |
| "invalid_label_boundary", |
| "Labels are not single tokens at the assistant boundary; configure a compatible assistant_prefix.", |
| "options.assistant_prefix", |
| ) |
| return final_ids, final_text |
|
|
