diff --git a/AGENTS.md b/AGENTS.md index e7d7c7d2d..e1d430256 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -50,7 +50,7 @@ Source code lives in `src/nemo_safe_synthesizer/`: | `configurator/` | Pydantic-to-Click mapping, Parameter types, validators | | `data_processing/` | Holdout, actions, assembler, records, shared token budget (`budget.py`), shared column validators (`validation.py`) | | `evaluation/` | Evaluator, components (privacy, MI, AIA, PII replay), reports | -| `generation/` | GeneratorBackend, VllmBackend, regex manager, batch gen | +| `generation/` | GeneratorBackend (owns the concrete `generate()` batch-loop template method), VllmBackend, RemoteBackend (GPU-free; calls a vLLM/NIM OpenAI-compatible endpoint), regex manager, batch gen | | `holdout/` | Train/test splitting | | `llm/` | Model loading, metadata, memory management | | `pii_replacer/` | NER-based PII detection and replacement | diff --git a/src/nemo_safe_synthesizer/config/__init__.py b/src/nemo_safe_synthesizer/config/__init__.py index bcd17089f..d69917f02 100644 --- a/src/nemo_safe_synthesizer/config/__init__.py +++ b/src/nemo_safe_synthesizer/config/__init__.py @@ -9,7 +9,7 @@ from .differential_privacy import DifferentialPrivacyHyperparams from .evaluate import EvaluationParameters from .external_results import SafeSynthesizerSummary, SafeSynthesizerTiming -from .generate import GenerateParameters, StructuredGenerationParameters +from .generate import GenerateParameters, RemoteParameters, StructuredGenerationParameters from .internal_results import SafeSynthesizerResults from .job import SafeSynthesizerJobConfig from .parameters import SafeSynthesizerParameters @@ -26,6 +26,7 @@ "GenerateParameters", "PiiReplacerConfig", "PreflightParameters", + "RemoteParameters", "SafeSynthesizerJobConfig", "SafeSynthesizerParameters", "SafeSynthesizerResults", diff --git a/src/nemo_safe_synthesizer/config/generate.py b/src/nemo_safe_synthesizer/config/generate.py index 47b1e05aa..44543cc8d 100644 --- a/src/nemo_safe_synthesizer/config/generate.py +++ b/src/nemo_safe_synthesizer/config/generate.py @@ -26,11 +26,14 @@ StructuredGenerationSchemaMethod = Literal["auto", "regex", "json_schema", "structural_tag"] ResolvedStructuredGenerationSchemaMethod = Literal["regex", "json_schema", "structural_tag"] StructuredGenerationBackend = Literal["auto", "xgrammar", "guidance", "outlines", "lm-format-enforcer"] +RemoteDialect = Literal["vllm", "openai"] STRUCTURAL_TAG_COMPATIBLE_BACKENDS = frozenset({"auto", "xgrammar"}) __all__ = [ "GenerateParameters", + "RemoteDialect", + "RemoteParameters", "ResolvedStructuredGenerationSchemaMethod", "StructuredGenerationParameters", "StructuredGenerationBackend", @@ -176,6 +179,91 @@ def _validate_structural_tag_backend(self) -> Self: return self +class RemoteParameters(Parameters, BaseModel): + """Connection to an external vLLM OpenAI-compatible inference server. + + When set on :class:`GenerateParameters`, generation issues HTTP requests + to this endpoint instead of loading a local vLLM engine. The server must + already serve the base model with the fine-tuned LoRA adapter attached, + registered under ``model``. No GPU is used locally. + + Structured generation maps to vLLM's ``structured_outputs`` request field + (``regex`` / ``json`` / ``structural_tag``), so all schema methods are + supported against a vLLM 0.20+ server. + """ + + dialect: Annotated[ + RemoteDialect, + Field( + title="dialect", + description=( + "Which request fields to send. 'vllm' (default) adds vLLM's sampling extensions " + "(repetition_penalty, top_k, min_p, skip_special_tokens, include_stop_str_in_output, " + "ignore_eos). 'openai' sends only the universal OpenAI fields, for stricter servers " + "(e.g. NIM / TensorRT-LLM) that reject those extensions." + ), + ), + ] = "vllm" + + endpoint_url: Annotated[ + str, + Field( + title="endpoint_url", + description="Base URL of the OpenAI-compatible server, e.g. 'http://localhost:8000/v1'.", + ), + ] + + model: Annotated[ + str, + Field( + title="model", + description="Model name as registered on the server (the served base model or LoRA adapter name).", + ), + ] + + api_key_env: Annotated[ + str | None, + Field( + title="api_key_env", + description=( + "Name of the environment variable holding the bearer token for the endpoint. " + "When unset, no Authorization header is sent." + ), + ), + ] = None + + timeout_seconds: Annotated[ + float, + ValueValidator(value_func=lambda v: v > 0), + Field( + title="timeout_seconds", + description="Per-request timeout in seconds. Must be > 0.", + ), + ] = 300.0 + + max_concurrency: Annotated[ + int, + ValueValidator(value_func=lambda v: v >= 1), + Field( + title="max_concurrency", + description="Maximum number of concurrent in-flight requests per batch. Must be >= 1.", + ), + ] = 16 + + max_retries: Annotated[ + int, + ValueValidator(value_func=lambda v: v >= 0), + Field( + title="max_retries", + description=( + "Number of retry attempts per request for transient failures (connection drops, " + "timeouts, and HTTP 408/409/425/429/500/502/503/504), using exponential backoff with " + "full jitter and honoring a Retry-After header. 0 disables retries. Must be >= 0." + ), + ), + ] = 4 + + class GenerateParameters(Parameters, BaseModel): """Configuration parameters for synthetic data generation. @@ -260,6 +348,18 @@ class GenerateParameters(Parameters, BaseModel): default_factory=ValidationParameters, ) + remote: Annotated[ + RemoteParameters | None, + Field( + title="remote", + description=( + "When set, generate by calling an external vLLM OpenAI-compatible server instead of " + "loading a local vLLM engine. The server must already serve the trained LoRA adapter. " + "Not supported for time-series generation." + ), + ), + ] = None + attention_backend: Annotated[ str | None, Field( diff --git a/src/nemo_safe_synthesizer/config/parameters.py b/src/nemo_safe_synthesizer/config/parameters.py index d3fe890ea..abf38acb2 100644 --- a/src/nemo_safe_synthesizer/config/parameters.py +++ b/src/nemo_safe_synthesizer/config/parameters.py @@ -163,6 +163,18 @@ def check_timeseries_group_column(self) -> Self: ) return self + @model_validator(mode="after") + def check_remote_not_timeseries(self) -> Self: + """Reject remote generation for time-series datasets at config time. + + The remote backend has no equivalent of the grouped time-series + generation loop, so the combination is unsupported. Catching it here + fails fast instead of after training completes. + """ + if self.generation.remote is not None and self.time_series is not None and self.time_series.is_timeseries: + raise ParameterError("Remote generation is not supported for time-series datasets.") + return self + @classmethod @override def from_params(cls, **kwargs: object) -> "SafeSynthesizerParameters": diff --git a/src/nemo_safe_synthesizer/generation/backend.py b/src/nemo_safe_synthesizer/generation/backend.py index 97c64f906..d17b7a6b1 100644 --- a/src/nemo_safe_synthesizer/generation/backend.py +++ b/src/nemo_safe_synthesizer/generation/backend.py @@ -6,14 +6,18 @@ from __future__ import annotations import abc +import time from collections.abc import Callable from .. import utils from ..cli.artifact_structure import Workdir from ..config import SafeSynthesizerParameters -from ..generation.results import GenerateJobResults +from ..defaults import FIXED_RUNTIME_GENERATE_ARGS from ..llm.metadata import ModelMetadata -from ..observability import get_logger +from ..observability import get_logger, heartbeat +from .batch import Batch +from .processors import Processor, TabularDataProcessor +from .results import GenerateJobResults, GenerationBatches, GenerationStatus logger = get_logger(__name__) @@ -21,17 +25,24 @@ class GeneratorBackend(metaclass=abc.ABCMeta): """Abstract base class for generation backends. - Lifecycle: ``initialize`` -> ``prepare_params`` -> ``generate`` - [-> ``generate`` ...] -> ``teardown``. + Lifecycle: ``initialize`` -> ``generate`` [-> ``generate`` ...] -> + ``teardown``. + + ``generate`` is a concrete template method that owns the batch loop, + stopping conditions, and result aggregation -- machinery that is + identical across engines because only decoded text and token counts + flow out of the model call. Subclasses provide the engine-specific + pieces by implementing ``initialize``, ``prepare_params``, + ``_generate_batch``, ``_get_prompt_token_count``, and ``teardown``. + A subclass with a fundamentally different loop (e.g. the grouped + time-series flow) may override ``generate`` wholesale. ``teardown`` must be idempotent and safe to call multiple times. Callers should use ``try/finally`` to guarantee ``teardown`` runs even if ``generate`` raises. Each cleanup step should be isolated - so one failure doesn't prevent the next from running. - - Subclasses must implement ``initialize``, ``prepare_params``, - ``generate``, and ``teardown``. The ``_torn_down`` guard flag - pattern is recommended for teardown implementations. + so one failure doesn't prevent the next from running. The + ``_torn_down`` guard flag pattern is recommended for teardown + implementations. """ gen_method: Callable | None = None @@ -55,27 +66,26 @@ class GeneratorBackend(metaclass=abc.ABCMeta): workdir: Workdir """Working directory containing model artifacts.""" - @classmethod - def __subclasshook__(cls, subclass): - return ( - hasattr(subclass, "prepare_args") - and callable(subclass.prepare_params) - and hasattr(subclass, "load") - and callable(subclass.initialize) - and hasattr(subclass, "generate") - and callable(subclass.generate) - or NotImplemented - ) + prompt: str + """Templated generation prompt sent to the model for every record.""" + + columns: list[str] + """Schema column names, in order, used to assemble the result frame.""" + + processor: Processor + """Parser that turns raw model text into validated records.""" + + use_detailed_logs: bool = False + """Whether to emit verbose per-record error messages (may leak data).""" @abc.abstractmethod def initialize(self) -> None: - """Load the model and any required resources into memory. + """Acquire the resources the backend needs to serve generations. - Called once before the first ``generate()`` invocation. - Implementations should allocate GPU memory, instantiate the - inference engine (e.g. vLLM), load LoRA adapters, and configure - backend-specific settings such as attention backends or - structured-output support. + Called once before the first ``generate()`` invocation. What this + entails is backend-specific: a local engine (e.g. vLLM) allocates GPU + memory, instantiates the engine, and loads LoRA adapters, while a + remote backend opens an HTTP client and connection pool. After this method returns, the backend must be ready to accept ``prepare_params()`` and ``generate()`` calls. @@ -100,30 +110,176 @@ def prepare_params(self, **kwargs) -> None: """ @abc.abstractmethod + def _get_prompt_token_count(self) -> int: + """Return the templated prompt's tokenized length. + + Used to size the per-sample ``max_tokens`` budget so the prompt + plus completion stays within the model's context window. Return + ``0`` to disable the prompt-length clamp when no tokenizer is + available. + """ + + @abc.abstractmethod + def _generate_batch( + self, + num_prompts_per_batch: int, + batch: Batch, + **sampling_kwargs, + ) -> Batch: + """Run the engine on one batch of prompts and populate ``batch``. + + Implementations issue ``num_prompts_per_batch`` generations of + ``self.prompt``, then for each completion record its finish + reason and call ``batch.process(idx, text, completion_tokens=...)`` + with the decoded text and token count. The engine-native response + objects must not escape this method -- only text and counts flow + downstream. + + Args: + num_prompts_per_batch: Number of prompts to run this batch. + batch: Fresh batch carrying the configured processor. + + Returns: + The same ``batch``, populated with parsed records and stats. + """ + def generate( self, data_actions_fn: utils.DataActionsFn | None = None, ) -> GenerateJobResults: """Run the batch generation loop and return aggregated results. - Repeatedly prompts the model, processes each batch through the - configured - [`Processor`][nemo_safe_synthesizer.generation.processors.Processor], - and accumulates valid records until the target count is reached - or a stopping condition fires (e.g. too many consecutive invalid - batches). Progress and error statistics are logged after each - batch. + Repeatedly prompts the model via ``_generate_batch`` and processes + each batch through the configured + [`Processor`][nemo_safe_synthesizer.generation.processors.Processor] + until the target record count is reached or a stopping condition + fires (e.g. too many consecutive invalid batches). Progress and + error statistics are logged after each batch. + + Non-tabular processors need BOS/EOS delimiters in the raw text, so + generation keeps special tokens for those processors and strips + them only for ``TabularDataProcessor``. Native EOS stopping + remains enabled through ``ignore_eos=False``. Args: - data_actions_fn: Optional post-processing / validation - function applied to each batch of generated records. - Typically reverses training-time preprocessing and - enforces user-specified data constraints. + data_actions_fn: Optional post-processing / validation function + applied to each batch of generated records. Typically + reverses training-time preprocessing and enforces + user-specified data constraints. Returns: Results containing the generated DataFrame, validity statistics, and timing information. """ + generation_start = time.monotonic() + + need_special_token_outputs = not isinstance(self.processor, TabularDataProcessor) + sampling_kwargs = dict( + temperature=self.config.generation.temperature, + repetition_penalty=self.config.generation.repetition_penalty, + top_p=self.config.generation.top_p, + top_k=FIXED_RUNTIME_GENERATE_ARGS["top_k"], + min_p=FIXED_RUNTIME_GENERATE_ARGS["min_p"], + max_tokens=self.model_metadata.generation_max_tokens_for(self._get_prompt_token_count()), + skip_special_tokens=not need_special_token_outputs, + include_stop_str_in_output=need_special_token_outputs, + ignore_eos=False, + ) + + self.prepare_params(**sampling_kwargs) + + # The batches object collects batches and keeps track of the stopping condition. + batches = GenerationBatches( + target_num_records=self.config.generation.num_records, + invalid_fraction_threshold=self.config.generation.invalid_fraction_threshold, + patience=self.config.generation.patience, + data_actions_fn=data_actions_fn, + ) + + with heartbeat( + "Generation", + logger_name=__name__, + target_records=self.config.generation.num_records, + progress_note=("Long stretches with no new records are normal."), + ): + while batches.num_valid_records < self.config.generation.num_records: + # Generate a batch from prompts and process the responses. + num_prompts = batches.get_next_num_prompts() + start_time = time.perf_counter() + batch: Batch = self._generate_batch( + num_prompts_per_batch=num_prompts, + batch=Batch(processor=self.processor), + **sampling_kwargs, + ) + duration = time.perf_counter() - start_time + batches.add_batch(batch) + + # Log generation summary and progress. + batch.log_summary(detailed_errors=self.use_detailed_logs) + self._log_batch_timing_and_progress(batch=batch, duration=duration, batches=batches) + # Check if the generation job should stop. + if batches.status in [ + GenerationStatus.STOP_NO_RECORDS, + GenerationStatus.STOP_METRIC_REACHED, + ]: + break + + batches.job_complete() + batches.log_status() + + max_num_records = ( + self.config.generation.num_records + if self.config.data.group_training_examples_by is None and batches.status == GenerationStatus.COMPLETE + else None + ) + + self.elapsed_time = time.monotonic() - generation_start + self.gen_results = GenerateJobResults.from_batches( + batches=batches, + columns=self.columns, + max_num_records=max_num_records, + elapsed_time=self.elapsed_time, + ) + + return self.gen_results + + def _log_batch_timing_and_progress( + self, + batch: Batch, + duration: float, + batches: GenerationBatches, + ) -> None: + """Log batch timing and progress as a structured Rich table. + + Emits structured data via ``logger.user.info`` that is rendered + as a Rich ASCII table on the console and as key/value pairs in + JSON logs. + """ + records_per_second = 0 if duration == 0 else batch.num_valid_records / duration + + # Build structured data - processor renders as table for console + progress_data: dict[str, int | float] = { + "records_per_second": round(records_per_second, 2), + "duration_seconds": round(duration, 2), + "valid_records_generated": batches.num_valid_records, + "target_records": self.config.generation.num_records, + "progress_fraction": round(batches.num_valid_records / self.config.generation.num_records, 4), + } + if batch.total_completion_tokens > 0 and duration > 0: + progress_data["tokens_per_second"] = round(batch.total_completion_tokens / duration, 1) + progress_data["valid_tokens_per_second"] = round(batch.total_valid_record_tokens / duration, 1) + + # Pass structured data - processor renders for console, JSON keeps as-is + logger.user.info( + "", + extra={ + "ctx": { + "render_table": True, + "tabular_data": progress_data, + "title": "Batch Progress", + } + }, + ) @abc.abstractmethod def teardown(self) -> None: diff --git a/src/nemo_safe_synthesizer/generation/remote_backend.py b/src/nemo_safe_synthesizer/generation/remote_backend.py new file mode 100644 index 000000000..354254202 --- /dev/null +++ b/src/nemo_safe_synthesizer/generation/remote_backend.py @@ -0,0 +1,478 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Generation backend that calls a remote vLLM OpenAI-compatible server.""" + +from __future__ import annotations + +import json +import os +import random +import time +from concurrent.futures import ThreadPoolExecutor +from typing import Any + +import httpx + +from .. import utils +from ..cli.artifact_structure import Workdir +from ..config import SafeSynthesizerParameters +from ..config.generate import resolve_structured_generation_schema_method +from ..errors import GenerationError, InternalError, ParameterError +from ..llm.metadata import ModelMetadata +from ..observability import get_logger +from ..utils import load_json +from .backend import GeneratorBackend +from .batch import Batch +from .processors import Processor, create_processor +from .regex_manager import build_json_based_regex, build_json_structural_tag + +logger = get_logger(__name__) + +_COMPLETIONS_PATH = "/completions" + +# Status codes worth retrying: request-timeout/conflict/too-early, rate limiting, +# and the transient 5xx family. Other 4xx (400/401/403/404) are permanent for a +# fixed request body and fail fast instead. +_RETRYABLE_STATUS = frozenset({408, 409, 425, 429, 500, 502, 503, 504}) +_BACKOFF_BASE_SECONDS = 0.5 +_BACKOFF_MAX_SECONDS = 30.0 + + +def _parse_retry_after(response: httpx.Response) -> float | None: + """Return the ``Retry-After`` delay in seconds, or ``None`` if absent/unparseable. + + Only the numeric-seconds form is honored; the rarely-used HTTP-date form is + ignored so the caller falls back to computed backoff. + """ + value = response.headers.get("Retry-After") + if not value: + return None + try: + return max(0.0, float(value)) + except ValueError: + return None + + +def _backoff_delay(attempt: int, retry_after: float | None) -> float: + """Seconds to wait before retry ``attempt`` (0-indexed). + + Honors a server ``Retry-After`` when given; otherwise uses full-jitter + exponential backoff (``random in [0, base * 2**attempt]``), capped at + ``_BACKOFF_MAX_SECONDS``. Jitter spreads retries from the concurrent worker + pool so they don't stampede the server in lockstep. + """ + if retry_after is not None: + return min(retry_after, _BACKOFF_MAX_SECONDS) + capped = min(_BACKOFF_BASE_SECONDS * (2**attempt), _BACKOFF_MAX_SECONDS) + return random.uniform(0.0, capped) + + +def _coerce_token_count(value: object) -> int: + """Coerce a server-reported ``completion_tokens`` to a non-negative int. + + The ``usage`` block is advisory, so a missing, null, boolean, or malformed + value must not fail an otherwise-valid completion -- it degrades to ``0``. + """ + match value: + case bool(): # bool is an int subclass; treat as "not a real count" + return 0 + case int(): + return max(0, value) + case float() | str(): + try: + return max(0, int(float(value))) + except ValueError: + return 0 + case _: + return 0 + + +def _compact_json_completion(text: str) -> str: + """Collapse a pretty-printed JSON object onto a single line for the JSONL processor. + + Servers constrained with ``structured_outputs: {"json": schema}`` may + pretty-print the object across multiple lines. The line-oriented record + extractor matches ``{.+?}`` with ``.`` *not* spanning newlines, so a + multi-line object yields zero records. Re-encoding the parsed object with + no insignificant whitespace produces the compact single-line shape the + processor expects. + + Returns the input unchanged when it does not parse as a single JSON object + (e.g. already-compact JSONL, multiple objects, or non-JSON text), so this + is a safe no-op outside the pretty-printed ``json_schema`` case. + """ + stripped = text.strip() + if not stripped: + return text + try: + obj = json.loads(stripped) + except (ValueError, TypeError): + return text + if not isinstance(obj, dict): + return text + return json.dumps(obj, separators=(",", ":"), ensure_ascii=False) + + +class RemoteBackend(GeneratorBackend): + """Generation backend that calls an external vLLM OpenAI-compatible server. + + Unlike [`VllmBackend`][nemo_safe_synthesizer.generation.vllm_backend.VllmBackend], + this backend never loads a model locally: it issues HTTP requests to a + server that already serves the base model with the fine-tuned LoRA adapter + attached (registered under ``config.generation.remote.model``). It reuses + the shared batch loop in + [`GeneratorBackend.generate`][nemo_safe_synthesizer.generation.backend.GeneratorBackend.generate] + and implements only the HTTP-specific pieces, so no GPU, vLLM engine, or + CUDA dependency is required at generation time. + + Each record is one ``/v1/completions`` request with ``n=1``; the server's + reported ``usage.completion_tokens`` is therefore the exact per-completion + token count. Requests within a batch are issued concurrently up to + ``config.generation.remote.max_concurrency``. + + Structured generation maps to vLLM's ``structured_outputs`` request field + (``regex`` / ``json`` / ``structural_tag``), mirroring the offline + ``StructuredOutputsParams``, so all three schema methods -- including the + ``auto`` default -- are supported. + + Args: + config: Pipeline configuration. ``config.generation.remote`` must be set. + model_metadata: Model metadata (prompt template, instruction, schema). + workdir: Working directory containing the dataset schema. + **kwargs: Additional options. ``use_detailed_logs`` (bool) enables + verbose per-record error messages (disabled by default to avoid + leaking sensitive data). + """ + + def __init__( + self, + config: SafeSynthesizerParameters, + model_metadata: ModelMetadata, + workdir: Workdir, + **kwargs, + ): + self.config = config + self.model_metadata = model_metadata + self.workdir = workdir + self.remote = True + self.use_detailed_logs = kwargs.pop("use_detailed_logs", False) + + self.schema = load_json(workdir.schema_file) + self.columns = list(self.schema["properties"].keys()) + self.prompt = utils.create_schema_prompt( + self.columns, + instruction=model_metadata.instruction, + prompt_template=model_metadata.prompt_config.template, + ) + # No local tokenizer in the remote setup: per-record token counts come + # from the server's usage report, so the processor stays tokenizer-less. + self.processor: Processor = create_processor(self.schema, model_metadata, config) + + self._client: httpx.Client | None = None + self._pool: ThreadPoolExecutor | None = None + self._request_body: dict[str, Any] | None = None + self._prompt_token_count: int | None = None + # Set in ``_build_structured_outputs`` when the resolved schema method is + # ``json_schema``: such servers may pretty-print JSON, so completions are + # compacted to single-line JSONL before the processor sees them. + self._compact_json = False + self._torn_down = False + + @property + def _remote(self): + """Remote connection config, validated to be present.""" + remote = self.config.generation.remote + if remote is None: + raise InternalError("RemoteBackend requires `config.generation.remote` to be configured.") + return remote + + def initialize(self) -> None: + """Create the HTTP client and worker pool for the remote server. + + Does not contact the server; connection errors surface on the first + ``generate()`` request instead, with per-request context. + """ + self._torn_down = False + remote = self._remote + + headers: dict[str, str] = {} + if remote.api_key_env: + api_key = os.environ.get(remote.api_key_env) + if not api_key: + raise ParameterError(f"Remote endpoint API key env var {remote.api_key_env!r} is not set or is empty.") + headers["Authorization"] = f"Bearer {api_key}" + + self._client = httpx.Client( + base_url=remote.endpoint_url.rstrip("/"), + headers=headers, + timeout=remote.timeout_seconds, + ) + self._pool = ThreadPoolExecutor(max_workers=remote.max_concurrency) + logger.info( + "RemoteBackend ready: endpoint=%s model=%s max_concurrency=%d", + remote.endpoint_url, + remote.model, + remote.max_concurrency, + ) + + def _get_prompt_token_count(self) -> int: + """Return the templated prompt's token length, or ``0`` when no tokenizer is local. + + Uses ``model_metadata.tokenizer`` opportunistically -- it is present + on the train-then-generate path (loaded from the HF cache) but ``None`` + on the resume path (``from_metadata_json`` excludes it) and the typical + offline-remote setup where the model was never downloaded locally. The + remote backend never *forces* a tokenizer load, so it stays GPU- and + download-free. + + When the count is ``0`` the prompt-length clamp in + [`generation_max_tokens_for`][nemo_safe_synthesizer.llm.metadata.ModelMetadata.generation_max_tokens_for] + is disabled; the server enforces its own context window and the + per-sample ``max_tokens`` budget is sized from the training-time example + length, so it stays well within that window regardless. The count is + cached after the first call. + """ + if self._prompt_token_count is not None: + return self._prompt_token_count + tokenizer = self.model_metadata.tokenizer + if tokenizer is None: + return 0 + self._prompt_token_count = len(tokenizer.encode(self.prompt)) + return self._prompt_token_count + + def _build_structured_outputs(self) -> dict[str, Any]: + """Map structured-generation config to a vLLM ``structured_outputs`` request field. + + Mirrors the offline + [`StructuredOutputsParams`][vllm.sampling_params.StructuredOutputsParams] + the local backend builds, so all three schema methods are supported: + ``regex`` -> ``{"regex": ...}``, ``json_schema`` -> ``{"json": schema}``, + and ``structural_tag`` -> ``{"structural_tag": ...}`` (the XGrammar tag + is sent as its JSON-encoded string, which yields multi-record JSONL). + Returns an empty dict when structured generation is disabled. + + The legacy top-level ``guided_regex`` / ``guided_json`` fields are + silently ignored by vLLM 0.20+ servers, so the nested + ``structured_outputs`` field is used instead. + + For ``json_schema`` the constraint is compacted at two layers: on the + ``vllm`` dialect ``disable_any_whitespace`` makes xgrammar emit + single-line JSON at the source (saving completion tokens), and + ``self._compact_json`` is set so ``_generate_batch`` also collapses any + residual multi-line output -- a portable net for the ``openai`` dialect, + where that vLLM-only field can't be sent. The ``regex`` and + ``structural_tag`` methods already enforce the single-line shape, so no + post-processing is needed and ``auto`` (-> ``structural_tag``) sidesteps + the issue entirely. + """ + gen = self.config.generation + self._compact_json = False + structured_generation = gen.structured_generation + if not structured_generation.enabled: + return {} + + method = resolve_structured_generation_schema_method( + structured_generation.schema_method, + structured_generation.backend, + ) + pc = self.model_metadata.prompt_config + if method == "regex": + logger.info("Structured generation enabled; constraining output with a regex") + regex = build_json_based_regex(self.schema, self.config, bos_token=pc.bos_token, eos_token=pc.eos_token) + return {"structured_outputs": {"regex": regex}} + if method == "json_schema": + self._compact_json = True + json_constraint: dict[str, Any] = {"json": self.schema} + if self._remote.dialect == "vllm": + # Source fix for vLLM: xgrammar (its default backend) emits compact JSON + # when whitespace is disabled, so the server never pretty-prints across + # lines and no completion tokens are wasted on whitespace. This is a vLLM + # protocol extension -- strict OpenAI servers (the "openai" dialect, e.g. + # NIM/TRT-LLM) 400 on it -- so it is gated like the other vLLM extensions. + # ``_compact_json`` stays armed regardless as a portable safety net. + json_constraint["disable_any_whitespace"] = True + logger.info( + "Structured generation enabled; constraining output with a JSON schema " + "(remote completions compacted to single-line JSONL)" + ) + return {"structured_outputs": json_constraint} + if method == "structural_tag": + logger.info("Structured generation enabled; constraining output with an XGrammar structural tag") + tag = build_json_structural_tag(self.schema, self.config, bos_token=pc.bos_token, eos_token=pc.eos_token) + return {"structured_outputs": {"structural_tag": tag}} + + raise InternalError(f"Unhandled structured-generation schema method: {method!r}") + + def prepare_params(self, **kwargs) -> None: + """Build the reusable ``/v1/completions`` request body from sampling params. + + The request fields depend on ``config.generation.remote.dialect``: + + - ``"vllm"`` (default): the universal OpenAI fields (``temperature``, + ``top_p``, ``max_tokens``, ``n``) plus vLLM's protocol extensions + (``repetition_penalty``, ``top_k``, ``min_p``, ``skip_special_tokens``, + ``include_stop_str_in_output``, ``ignore_eos``). + - ``"openai"``: only the universal fields, for stricter servers (e.g. + NIM / TensorRT-LLM) that reject the vLLM extensions with a 400. + + The resolved sampling values are identical across dialects; only which + fields go on the wire differs. The prompt is constant across every + request in a run, so it is baked into the body here once rather than + merged per request. + """ + body: dict[str, Any] = { + "model": self._remote.model, + "prompt": self.prompt, + "n": 1, + "temperature": kwargs["temperature"], + "top_p": kwargs["top_p"], + "max_tokens": kwargs["max_tokens"], + } + if self._remote.dialect == "vllm": + body |= { + "repetition_penalty": kwargs["repetition_penalty"], + "top_k": kwargs["top_k"], + "min_p": kwargs["min_p"], + "skip_special_tokens": kwargs["skip_special_tokens"], + "include_stop_str_in_output": kwargs["include_stop_str_in_output"], + "ignore_eos": kwargs["ignore_eos"], + } + body |= self._build_structured_outputs() + self._request_body = body + + def _complete_one(self) -> tuple[str, int, str | None]: + """Issue one completion (with transient-failure retries) and parse the result. + + Returns ``(text, completion_tokens, finish_reason)``. + """ + return self._parse_completion(self._post_completion()) + + def _post_completion(self) -> httpx.Response: + """POST one completion request, retrying transient failures with backoff. + + Connection errors, timeouts, and retryable status codes + (``_RETRYABLE_STATUS``) are retried up to ``remote.max_retries`` times + with full-jitter exponential backoff, honoring a ``Retry-After`` header + when present. A non-retryable status (e.g. 400/401/404) fails + immediately via ``raise_for_status`` -- it would fail identically for + every record. ``GenerationError`` is raised once retries are exhausted. + """ + if self._client is None or self._request_body is None: + raise InternalError("RemoteBackend._post_completion() called before initialize()/prepare_params().") + + endpoint = self._remote.endpoint_url + max_retries = self._remote.max_retries + last_error = "no attempts made" + + for attempt in range(max_retries + 1): + retry_after: float | None = None + try: + response = self._client.post(_COMPLETIONS_PATH, json=self._request_body) + except httpx.HTTPError as exc: + last_error = f"request failed: {exc}" + else: + if response.status_code not in _RETRYABLE_STATUS: + return self._raise_for_status(response) + last_error = f"status {response.status_code}: {response.text[:200]}" + retry_after = _parse_retry_after(response) + + if attempt == max_retries: + break + delay = _backoff_delay(attempt, retry_after) + logger.warning( + "Remote endpoint %s transient failure (%s); retry %d/%d in %.1fs", + endpoint, + last_error, + attempt + 1, + max_retries, + delay, + ) + time.sleep(delay) + + raise GenerationError(f"Remote endpoint {endpoint} failed after {max_retries + 1} attempt(s): {last_error}") + + def _raise_for_status(self, response: httpx.Response) -> httpx.Response: + """Return ``response`` if OK, else raise ``GenerationError`` with a truncated body.""" + try: + response.raise_for_status() + except httpx.HTTPStatusError as exc: + raise GenerationError( + f"Remote endpoint {self._remote.endpoint_url} returned {exc.response.status_code}: " + f"{exc.response.text[:500]}" + ) from exc + return response + + def _parse_completion(self, response: httpx.Response) -> tuple[str, int, str | None]: + """Extract ``(text, completion_tokens, finish_reason)`` from a completion response. + + The ``usage`` token count is advisory and coerced defensively; only a + missing/empty ``choices`` array is fatal, since it means no completion + was produced. + """ + try: + data = response.json() + choice = data["choices"][0] + except (ValueError, KeyError, IndexError, TypeError) as exc: + # Truncate the body so a large or sensitive payload isn't echoed in full. + raise GenerationError( + f"Remote endpoint {self._remote.endpoint_url} returned an unexpected response shape: {exc}. " + f"Body: {response.text[:500]}" + ) from exc + + usage = data.get("usage") or {} + return ( + choice.get("text", ""), + _coerce_token_count(usage.get("completion_tokens")), + choice.get("finish_reason"), + ) + + def _generate_batch( + self, + num_prompts_per_batch: int, + batch: Batch, + **_sampling_kwargs, + ) -> Batch: + """Issue ``num_prompts_per_batch`` concurrent completions and process the responses. + + Sampling parameters are already baked into the request body by + ``prepare_params``; the trailing kwargs the shared loop forwards are + accepted and ignored. + """ + if self._pool is None: + raise InternalError("RemoteBackend._generate_batch() called before initialize().") + + futures = [self._pool.submit(self._complete_one) for _ in range(num_prompts_per_batch)] + for idx, future in enumerate(futures): + text, completion_tokens, finish_reason = future.result() + if self._compact_json: + text = _compact_json_completion(text) + batch.finish_reasons[str(finish_reason or "unknown")] += 1 + batch.process(idx, text, completion_tokens=completion_tokens) + return batch + + def teardown(self) -> None: + """Close the HTTP client and worker pool. Idempotent.""" + if self._torn_down: + return + self._torn_down = True + + if self._pool is not None: + try: + self._pool.shutdown(wait=False) + except Exception: + logger.debug("RemoteBackend pool shutdown failed during teardown", exc_info=True) + if self._client is not None: + try: + self._client.close() + except Exception: + logger.debug("RemoteBackend client close failed during teardown", exc_info=True) + self._pool = None + self._client = None + + def __del__(self) -> None: + """Clean up resources on garbage collection.""" + try: + self.teardown() + except Exception: + logger.debug("RemoteBackend teardown failed during garbage collection", exc_info=True) diff --git a/src/nemo_safe_synthesizer/generation/vllm_backend.py b/src/nemo_safe_synthesizer/generation/vllm_backend.py index d9fbc4c39..420230211 100644 --- a/src/nemo_safe_synthesizer/generation/vllm_backend.py +++ b/src/nemo_safe_synthesizer/generation/vllm_backend.py @@ -9,7 +9,6 @@ import logging import os import tempfile -import time from functools import partial from pathlib import Path from typing import Any, cast @@ -31,13 +30,13 @@ resolve_structured_generation_schema_method, structural_tag_backend_error_message, ) -from ..defaults import DEFAULT_SAMPLING_PARAMETERS, FIXED_RUNTIME_GENERATE_ARGS +from ..defaults import DEFAULT_SAMPLING_PARAMETERS from ..errors import InternalError, ParameterError from ..generation.backend import GeneratorBackend from ..generation.batch import Batch from ..generation.processors import EncodeOnlyTokenizer, Processor, TabularDataProcessor, create_processor from ..generation.regex_manager import build_json_based_regex, build_json_structural_tag -from ..generation.results import GenerateJobResults, GenerationBatches, GenerationStatus +from ..generation.results import GenerateJobResults from ..generation.vllm_observability import ( GenerationObservability, NvmlPeakSampler, @@ -46,6 +45,7 @@ read_vllm_runtime_metrics, ) from ..llm.metadata import ModelMetadata +from ..llm.model_host import ModelHost from ..llm.utils import ModelRef, cleanup_memory, get_max_vram from ..observability import get_logger, heartbeat from ..utils import all_equal_type, load_json @@ -216,7 +216,7 @@ def _install_noop_remote_cache_backends() -> None: _install_noop_remote_cache_backends() -class VllmBackend(GeneratorBackend): +class VllmBackend(GeneratorBackend, ModelHost[vLLM, PreTrainedTokenizerBase]): """Generation backend using vLLM for high-throughput inference. Loads the base model with a LoRA adapter via vLLM and generates @@ -275,6 +275,20 @@ def __init__( self.lora_req = LoRARequest("lora", 1, str(adapter_path)) if adapter_path else None self._torn_down = False + @property + def model(self) -> vLLM | None: + """Return the locally hosted vLLM engine.""" + return self.llm + + @property + def tokenizer(self) -> PreTrainedTokenizerBase | None: + """Return the tokenizer owned by the locally hosted engine.""" + if self.llm is None: + return None + # vLLM declares a wider tokenizer union, but supported NSS engines + # expose a Hugging Face tokenizer implementing this interface. + return cast(PreTrainedTokenizerBase, self.llm.get_tokenizer()) + def teardown(self) -> None: """Release GPU memory and distributed resources. Idempotent -- safe to call multiple times.""" if self._torn_down: @@ -351,7 +365,9 @@ def initialize(self, **kwargs) -> None: # asked for. self._engine_runtime_config = probe_engine_runtime_config(self.llm) - tokenizer: EncodeOnlyTokenizer = self.llm.get_tokenizer() + tokenizer = self.tokenizer + if tokenizer is None: + raise InternalError("VllmBackend.initialize() did not create a tokenizer.") self.processor = create_processor( self.schema, self.model_metadata, @@ -369,9 +385,9 @@ def _get_prompt_token_count(self) -> int: """ if self._prompt_token_count is not None: return self._prompt_token_count - if self.llm is None: + tokenizer = self.tokenizer + if tokenizer is None: return 0 - tokenizer = self.llm.get_tokenizer() self._prompt_token_count = len(tokenizer.encode(self.prompt)) return self._prompt_token_count @@ -639,68 +655,22 @@ def _generate_batch( return batch - def _log_batch_timing_and_progress( - self, - batch: Batch, - duration: float, - num_records: int, - num_valid_records: int, - batches: GenerationBatches, - ) -> None: - """Log batch timing and progress as a structured Rich table. - - Emits structured data via ``logger.user.info`` that is rendered - as a Rich ASCII table on the console and as key/value pairs in - JSON logs. - """ - records_per_second = 0 if duration == 0 else batch.num_valid_records / duration - - # Build structured data - processor renders as table for console - progress_data: dict[str, int | float] = { - "records_per_second": round(records_per_second, 2), - "duration_seconds": round(duration, 2), - "valid_records_generated": batches.num_valid_records, - "target_records": self.config.generation.num_records, - "progress_fraction": round(batches.num_valid_records / self.config.generation.num_records, 4), - } - if batch.total_completion_tokens > 0 and duration > 0: - progress_data["tokens_per_second"] = round(batch.total_completion_tokens / duration, 1) - progress_data["valid_tokens_per_second"] = round(batch.total_valid_record_tokens / duration, 1) - - # Pass structured data - processor renders for console, JSON keeps as-is - logger.user.info( - "", - extra={ - "ctx": { - "render_table": True, - "tabular_data": progress_data, - "title": "Batch Progress", - } - }, - ) - def generate( self, data_actions_fn: utils.DataActionsFn | None = None, ) -> GenerateJobResults: - """Generate synthetic tabular data in batches until the target count is reached. - - Iterates over generation batches, applying the processor to each - LLM output, until the configured ``num_records`` target is met or - a stopping condition fires. - - Non-tabular processors need BOS/EOS delimiters in the raw text, so - generation keeps special tokens for those processors and strips them - only for ``TabularDataProcessor``. Native EOS stopping remains enabled - through ``ignore_eos=False``. - - Wrapped in an :class:`NvmlPeakSampler` context so a - generation-complete observability event is emitted at end of - the call carrying peak device VRAM, host loadavg pre/post, vLLM's - kv_cache_usage_perc / prefix_cache_hit_rate / spec_accept_rate - (read at end-of-generation), and the engine's effective runtime - config probed at engine-init time. Sampler is degraded-mode when - NVML is unavailable; the event still emits with ``peak_vram_gb=None``. + """Generate synthetic tabular data, bracketed with vLLM observability. + + Delegates the batch loop to + [`GeneratorBackend.generate`][nemo_safe_synthesizer.generation.backend.GeneratorBackend.generate] + (the shared template method that owns batching, stopping conditions, and + result aggregation) and wraps it in an :class:`NvmlPeakSampler` context + so a generation-complete observability event is emitted at end + of the call carrying peak device VRAM, host loadavg pre/post, vLLM's + kv_cache_usage_perc / prefix_cache_hit_rate / spec_accept_rate (read at + end-of-generation), and the engine's effective runtime config probed at + engine-init time. The sampler degrades when NVML is unavailable; the + event still emits with ``peak_vram_gb=None``. Args: data_actions_fn: Optional post-processing / validation function @@ -713,7 +683,7 @@ def generate( sampler = NvmlPeakSampler() try: with sampler: - self._run_generation(data_actions_fn) + super().generate(data_actions_fn) finally: # Emit regardless of success: the sampler thread has been joined # by ``with`` exit, so ``sampler.peak_gb`` is the final peak, and @@ -723,88 +693,6 @@ def generate( return self.gen_results - def _run_generation(self, data_actions_fn: utils.DataActionsFn | None) -> None: - """Run the batch-generation loop until the target or a stop condition fires. - - Populates ``self.gen_results`` and ``self.elapsed_time``. Extracted - from :meth:`generate` so that method stays a thin observability - bracket around the loop. - """ - generation_start = time.monotonic() - need_special_token_outputs = not isinstance(self.processor, TabularDataProcessor) - sampling_kwargs = dict( - temperature=self.config.generation.temperature, - repetition_penalty=self.config.generation.repetition_penalty, - top_p=self.config.generation.top_p, - top_k=FIXED_RUNTIME_GENERATE_ARGS["top_k"], - min_p=FIXED_RUNTIME_GENERATE_ARGS["min_p"], - max_tokens=self.model_metadata.generation_max_tokens_for(self._get_prompt_token_count()), - skip_special_tokens=not need_special_token_outputs, - include_stop_str_in_output=need_special_token_outputs, - ignore_eos=False, - ) - - self.prepare_params(**sampling_kwargs) - - # The batches object collects batches and keeps track of the stopping condition. - batches = GenerationBatches( - target_num_records=self.config.generation.num_records, - invalid_fraction_threshold=self.config.generation.invalid_fraction_threshold, - patience=self.config.generation.patience, - data_actions_fn=data_actions_fn, - ) - - with heartbeat( - "Generation", - logger_name=__name__, - target_records=self.config.generation.num_records, - progress_note=("Long stretches with no new records are normal."), - ): - while batches.num_valid_records < self.config.generation.num_records: - # Generate a batch from prompts and process the responses. - num_prompts = batches.get_next_num_prompts() - start_time = time.perf_counter() - batch: Batch = self._generate_batch( - num_prompts_per_batch=num_prompts, - batch=Batch(processor=self.processor), - **sampling_kwargs, - ) - duration = time.perf_counter() - start_time - batches.add_batch(batch) - - # Log generation summary and progress. - batch.log_summary(detailed_errors=self.use_detailed_logs) - self._log_batch_timing_and_progress( - batch=batch, - duration=duration, - num_records=self.config.generation.num_records, - num_valid_records=batches.num_valid_records, - batches=batches, - ) - # Check if the generation job should stop. - if batches.status in [ - GenerationStatus.STOP_NO_RECORDS, - GenerationStatus.STOP_METRIC_REACHED, - ]: - break - - batches.job_complete() - batches.log_status() - - max_num_records = ( - self.config.generation.num_records - if self.config.data.group_training_examples_by is None and batches.status == GenerationStatus.COMPLETE - else None - ) - - self.elapsed_time = time.monotonic() - generation_start - self.gen_results = GenerateJobResults.from_batches( - batches=batches, - columns=self.columns, - max_num_records=max_num_records, - elapsed_time=self.elapsed_time, - ) - def _emit_generation_observability( self, sampler: NvmlPeakSampler, loadavg_pre: tuple[float, float, float] | None ) -> None: diff --git a/src/nemo_safe_synthesizer/llm/model_host.py b/src/nemo_safe_synthesizer/llm/model_host.py new file mode 100644 index 000000000..72fb982ab --- /dev/null +++ b/src/nemo_safe_synthesizer/llm/model_host.py @@ -0,0 +1,39 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Typed lifecycle contract for components that own a local model.""" + +from __future__ import annotations + +from abc import ABC, abstractmethod +from typing import Generic, TypeVar + +ModelT = TypeVar("ModelT") +TokenizerT = TypeVar("TokenizerT") + + +class ModelHost(ABC, Generic[ModelT, TokenizerT]): + """Own a local language model and tokenizer through teardown. + + This contract deliberately stops at model ownership. Tasks such as + synthetic-record generation and column classification retain their own + prompt construction, batching, and response parsing. + """ + + @property + @abstractmethod + def model(self) -> ModelT | None: + """Return the hosted model, or ``None`` before initialization.""" + + @property + @abstractmethod + def tokenizer(self) -> TokenizerT | None: + """Return the hosted tokenizer, or ``None`` before initialization.""" + + @abstractmethod + def initialize(self) -> None: + """Load the model and any resources required to use it.""" + + @abstractmethod + def teardown(self) -> None: + """Release the model and its resources; this must be idempotent.""" diff --git a/src/nemo_safe_synthesizer/sdk/library_builder.py b/src/nemo_safe_synthesizer/sdk/library_builder.py index 2943dd70f..ed851063c 100644 --- a/src/nemo_safe_synthesizer/sdk/library_builder.py +++ b/src/nemo_safe_synthesizer/sdk/library_builder.py @@ -22,6 +22,7 @@ from ..configurator.parameters import Parameters from ..errors import ParameterError from ..evaluation.evaluator import Evaluator +from ..generation.remote_backend import RemoteBackend from ..generation.timeseries_backend import TimeseriesBackend from ..generation.vllm_backend import VllmBackend from ..holdout.holdout import Holdout @@ -161,7 +162,8 @@ class SafeSynthesizer(ConfigBuilder): results = builder.results ``train()`` uses ``HuggingFaceBackend``. ``generate()`` chooses - ``TimeseriesBackend`` when ``config.time_series.is_timeseries`` is true and + ``RemoteBackend`` when ``config.generation.remote`` is set, + ``TimeseriesBackend`` when ``config.time_series.is_timeseries`` is true, and ``VllmBackend`` otherwise. Stepwise callers must call ``save_results()`` themselves after ``evaluate()``; ``run()`` does this automatically. @@ -544,9 +546,10 @@ def train(self) -> SafeSynthesizer: def generate(self) -> SafeSynthesizer: """Generate synthetic data using the trained model. - Selects the appropriate backend (``VllmBackend`` or - ``TimeseriesBackend``), initializes it, and generates - synthetic records. + Selects the appropriate backend -- ``RemoteBackend`` when + ``config.generation.remote`` is set, ``TimeseriesBackend`` for + time-series datasets, otherwise the local ``VllmBackend`` -- then + initializes it and generates synthetic records. Returns: Self for method chaining. @@ -565,8 +568,15 @@ def generate(self) -> SafeSynthesizer: trainer.teardown() assert self._workdir is not None - # Select backend based on time_series configuration - if self._nss_config.time_series and self._nss_config.time_series.is_timeseries: + # Select backend: a configured remote endpoint wins, then time-series, + # otherwise the local vLLM engine. The remote+time-series combination is + # rejected at config validation time (SafeSynthesizerParameters). + is_timeseries = bool(self._nss_config.time_series and self._nss_config.time_series.is_timeseries) + if self._nss_config.generation.remote is not None: + self.generator = RemoteBackend( + config=self._nss_config, model_metadata=self._llm_metadata, workdir=self._workdir + ) + elif is_timeseries: self.generator = TimeseriesBackend( config=self._nss_config, model_metadata=self._llm_metadata, workdir=self._workdir ) diff --git a/tests/generation/test_remote_backend.py b/tests/generation/test_remote_backend.py new file mode 100644 index 000000000..9e9027b43 --- /dev/null +++ b/tests/generation/test_remote_backend.py @@ -0,0 +1,632 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Unit tests for the RemoteBackend HTTP generation backend. + +These tests never contact a live server: ``httpx`` and the processor are +mocked so the suite stays fast, deterministic, and GPU-free. +""" + +from unittest.mock import MagicMock, patch + +import httpx +import pytest +from pydantic import ValidationError + +from nemo_safe_synthesizer.config import ( + DataParameters, + GenerateParameters, + RemoteParameters, + SafeSynthesizerParameters, + StructuredGenerationParameters, + TimeSeriesParameters, + TrainingHyperparams, +) +from nemo_safe_synthesizer.config.generate import RemoteDialect +from nemo_safe_synthesizer.errors import GenerationError, InternalError, ParameterError +from nemo_safe_synthesizer.generation.backend import GeneratorBackend +from nemo_safe_synthesizer.generation.batch import Batch +from nemo_safe_synthesizer.generation.processors import TabularDataProcessor, create_processor +from nemo_safe_synthesizer.generation.remote_backend import ( + RemoteBackend, + _backoff_delay, + _coerce_token_count, + _compact_json_completion, + _parse_retry_after, +) +from nemo_safe_synthesizer.llm.metadata import ModelMetadata + +MODULE = "nemo_safe_synthesizer.generation.remote_backend" + + +@pytest.fixture +def mock_schema(): + return {"properties": {"name": {"type": "string"}, "age": {"type": "integer"}}} + + +@pytest.fixture +def mock_model_metadata(): + metadata = MagicMock(spec=ModelMetadata) + metadata.instruction = "Generate data" + metadata.prompt_config = MagicMock() + metadata.prompt_config.template = "[INST] {instruction} {schema} [/INST]" + metadata.prompt_config.bos_token = "" + metadata.prompt_config.eos_token = "" + # The resume / offline-remote path carries no local tokenizer. + metadata.tokenizer = None + return metadata + + +def make_params( + *, + endpoint_url: str = "http://localhost:8000/v1", + model: str = "my-lora", + api_key_env: str | None = None, + max_concurrency: int = 16, + max_retries: int = 4, + dialect: RemoteDialect = "vllm", +) -> SafeSynthesizerParameters: + """Build params with a configured remote endpoint.""" + return SafeSynthesizerParameters( + data=DataParameters(group_training_examples_by=None, order_training_examples_by=None), + training=TrainingHyperparams(pretrained_model="test-model", lora_r=16), + generation=GenerateParameters( + num_records=10, + structured_generation=StructuredGenerationParameters(enabled=False), + remote=RemoteParameters( + endpoint_url=endpoint_url, + model=model, + api_key_env=api_key_env, + max_concurrency=max_concurrency, + max_retries=max_retries, + dialect=dialect, + ), + ), + ) + + +def make_backend(config, model_metadata, schema, processor=None) -> RemoteBackend: + """Construct a RemoteBackend with patched schema/prompt/processor helpers.""" + with ( + patch(f"{MODULE}.load_json", return_value=schema), + patch(f"{MODULE}.utils.create_schema_prompt", return_value="test prompt"), + patch(f"{MODULE}.create_processor", return_value=processor or MagicMock()), + ): + return RemoteBackend(config=config, model_metadata=model_metadata, workdir=MagicMock()) + + +def make_http_response(text: str, completion_tokens: int, finish_reason: str = "stop") -> MagicMock: + response = MagicMock(spec=httpx.Response) + response.status_code = 200 + response.json.return_value = { + "choices": [{"text": text, "finish_reason": finish_reason}], + "usage": {"completion_tokens": completion_tokens}, + } + response.raise_for_status.return_value = None + return response + + +def make_status_response(status_code: int, text: str = "", *, retry_after: str | None = None) -> MagicMock: + """A response with an explicit status code, for retry/error-path tests.""" + response = MagicMock(spec=httpx.Response) + response.status_code = status_code + response.text = text + response.headers = {"Retry-After": retry_after} if retry_after else {} + if status_code >= 400: + response.raise_for_status.side_effect = httpx.HTTPStatusError( + str(status_code), request=MagicMock(), response=response + ) + else: + response.raise_for_status.return_value = None + return response + + +class TestConstruction: + def test_sets_remote_flag_and_prompt(self, mock_model_metadata, mock_schema): + backend = make_backend(make_params(), mock_model_metadata, mock_schema) + assert backend.remote is True + assert backend.prompt == "test prompt" + assert backend.columns == ["name", "age"] + + def test_is_concrete_generator_backend(self, mock_model_metadata, mock_schema): + backend = make_backend(make_params(), mock_model_metadata, mock_schema) + assert isinstance(backend, GeneratorBackend) + + def test_prompt_token_count_is_zero_without_tokenizer(self, mock_model_metadata, mock_schema): + backend = make_backend(make_params(), mock_model_metadata, mock_schema) + assert backend._get_prompt_token_count() == 0 + + def test_prompt_token_count_uses_tokenizer_when_present(self, mock_model_metadata, mock_schema): + tokenizer = MagicMock() + tokenizer.encode.return_value = [1, 2, 3, 4, 5] + mock_model_metadata.tokenizer = tokenizer + backend = make_backend(make_params(), mock_model_metadata, mock_schema) + assert backend._get_prompt_token_count() == 5 + tokenizer.encode.assert_called_once_with("test prompt") + # Result is cached: the tokenizer is not re-invoked. + assert backend._get_prompt_token_count() == 5 + tokenizer.encode.assert_called_once() + + +class TestInitialize: + def test_creates_client_and_pool(self, mock_model_metadata, mock_schema): + backend = make_backend(make_params(), mock_model_metadata, mock_schema) + backend.initialize() + try: + client, pool = backend._client, backend._pool + assert client is not None + assert pool is not None + assert "Authorization" not in client.headers + finally: + backend.teardown() + + def test_api_key_header_set_from_env(self, mock_model_metadata, mock_schema, monkeypatch): + monkeypatch.setenv("MY_KEY", "secret-token") + backend = make_backend(make_params(api_key_env="MY_KEY"), mock_model_metadata, mock_schema) + backend.initialize() + try: + client = backend._client + assert client is not None + assert client.headers["Authorization"] == "Bearer secret-token" + finally: + backend.teardown() + + def test_missing_api_key_raises(self, mock_model_metadata, mock_schema, monkeypatch): + monkeypatch.delenv("MY_KEY", raising=False) + backend = make_backend(make_params(api_key_env="MY_KEY"), mock_model_metadata, mock_schema) + with pytest.raises(ParameterError, match="MY_KEY"): + backend.initialize() + + +class TestPrepareParams: + def _sampling_kwargs(self, **overrides): + kwargs = dict( + temperature=0.9, + top_p=1.0, + max_tokens=128, + repetition_penalty=1.0, + top_k=-1, + min_p=0, + skip_special_tokens=True, + include_stop_str_in_output=False, + ignore_eos=False, + ) + kwargs.update(overrides) + return kwargs + + def test_builds_request_body(self, mock_model_metadata, mock_schema): + backend = make_backend(make_params(), mock_model_metadata, mock_schema) + backend.prepare_params(**self._sampling_kwargs()) + body = backend._request_body + assert body is not None + assert body["model"] == "my-lora" + assert body["n"] == 1 + assert body["temperature"] == 0.9 + assert body["max_tokens"] == 128 + # vLLM protocol extensions are present. + assert body["repetition_penalty"] == 1.0 + assert body["top_k"] == -1 + assert body["min_p"] == 0 + # No structured generation -> no structured_outputs field. + assert "structured_outputs" not in body + + def test_openai_dialect_omits_vllm_extensions(self, mock_model_metadata, mock_schema): + backend = make_backend(make_params(dialect="openai"), mock_model_metadata, mock_schema) + backend.prepare_params(**self._sampling_kwargs()) + body = backend._request_body + assert body is not None + # Universal OpenAI fields are still present. + assert body["temperature"] == 0.9 + assert body["top_p"] == 1.0 + assert body["max_tokens"] == 128 + assert body["n"] == 1 + # The vLLM-only extensions that strict servers (NIM/TRT-LLM) reject are dropped. + for field in ( + "repetition_penalty", + "top_k", + "min_p", + "skip_special_tokens", + "include_stop_str_in_output", + "ignore_eos", + ): + assert field not in body + + def test_structured_outputs_json_vllm_dialect_disables_whitespace(self, mock_model_metadata, mock_schema): + config = make_params(dialect="vllm") + config.generation.structured_generation.enabled = True + config.generation.structured_generation.schema_method = "json_schema" + backend = make_backend(config, mock_model_metadata, mock_schema) + backend.prepare_params(**self._sampling_kwargs()) + body = backend._request_body + assert body is not None + # Source fix: vLLM/xgrammar emits compact JSON when whitespace is disabled. + assert body["structured_outputs"] == {"json": mock_schema, "disable_any_whitespace": True} + # Compaction stays armed as a safety net even with the source fix on. + assert backend._compact_json is True + + def test_structured_outputs_json_openai_dialect_omits_vllm_extension(self, mock_model_metadata, mock_schema): + config = make_params(dialect="openai") + config.generation.structured_generation.enabled = True + config.generation.structured_generation.schema_method = "json_schema" + backend = make_backend(config, mock_model_metadata, mock_schema) + backend.prepare_params(**self._sampling_kwargs()) + body = backend._request_body + assert body is not None + # The vLLM-only field would 400 on strict OpenAI servers, so it is dropped; + # the post-process compaction is the portable fallback. + assert body["structured_outputs"] == {"json": mock_schema} + assert backend._compact_json is True + + def test_compact_json_not_armed_for_other_methods(self, mock_model_metadata, mock_schema): + config = make_params() + config.generation.structured_generation.enabled = True + config.generation.structured_generation.schema_method = "regex" + config.generation.structured_generation.backend = "outlines" + backend = make_backend(config, mock_model_metadata, mock_schema) + with patch(f"{MODULE}.build_json_based_regex", return_value="REGEX"): + backend.prepare_params(**self._sampling_kwargs()) + assert backend._compact_json is False + + def test_compact_json_not_armed_without_structured_generation(self, mock_model_metadata, mock_schema): + backend = make_backend(make_params(), mock_model_metadata, mock_schema) + backend.prepare_params(**self._sampling_kwargs()) + assert backend._compact_json is False + + def test_structured_outputs_regex_when_regex_method(self, mock_model_metadata, mock_schema): + config = make_params() + config.generation.structured_generation.enabled = True + config.generation.structured_generation.schema_method = "regex" + config.generation.structured_generation.backend = "outlines" + backend = make_backend(config, mock_model_metadata, mock_schema) + with patch(f"{MODULE}.build_json_based_regex", return_value="REGEX") as build_regex: + backend.prepare_params(**self._sampling_kwargs()) + build_regex.assert_called_once() + body = backend._request_body + assert body is not None + assert body["structured_outputs"] == {"regex": "REGEX"} + + def test_structured_outputs_structural_tag_when_structural_tag_method(self, mock_model_metadata, mock_schema): + config = make_params() + config.generation.structured_generation.enabled = True + config.generation.structured_generation.schema_method = "structural_tag" + config.generation.structured_generation.backend = "xgrammar" + backend = make_backend(config, mock_model_metadata, mock_schema) + with patch(f"{MODULE}.build_json_structural_tag", return_value="TAG") as build_tag: + backend.prepare_params(**self._sampling_kwargs()) + build_tag.assert_called_once() + body = backend._request_body + assert body is not None + assert body["structured_outputs"] == {"structural_tag": "TAG"} + + +class TestCompleteOne: + def test_parses_text_tokens_and_finish_reason(self, mock_model_metadata, mock_schema): + backend = make_backend(make_params(), mock_model_metadata, mock_schema) + backend._client = MagicMock() + backend._client.post.return_value = make_http_response('{"name": "A", "age": 1}', 7, "stop") + backend._request_body = {"model": "my-lora", "prompt": "test prompt"} + text, tokens, finish = backend._complete_one() + assert text == '{"name": "A", "age": 1}' + assert tokens == 7 + assert finish == "stop" + # The prompt is baked into the request body by prepare_params and posted as-is. + _, kwargs = backend._client.post.call_args + assert kwargs["json"] == {"model": "my-lora", "prompt": "test prompt"} + + def test_complete_one_before_init_raises(self, mock_model_metadata, mock_schema): + backend = make_backend(make_params(), mock_model_metadata, mock_schema) + with pytest.raises(InternalError): + backend._complete_one() + + +class TestPostCompletionRetries: + """Transient-failure resilience: the network path must survive blips, not abort the run.""" + + @staticmethod + def _backend_and_client(mock_model_metadata, mock_schema, *, max_retries: int) -> tuple[RemoteBackend, MagicMock]: + backend = make_backend(make_params(max_retries=max_retries), mock_model_metadata, mock_schema) + client = MagicMock() + backend._client = client + backend._request_body = {"model": "m", "prompt": "p"} + return backend, client + + def test_retries_retryable_status_then_succeeds(self, mock_model_metadata, mock_schema): + backend, client = self._backend_and_client(mock_model_metadata, mock_schema, max_retries=3) + client.post.side_effect = [ + make_status_response(503, "busy"), + make_http_response('{"name": "A"}', 3, "stop"), + ] + with patch(f"{MODULE}.time.sleep") as sleep: + text, tokens, finish = backend._complete_one() + assert (text, tokens, finish) == ('{"name": "A"}', 3, "stop") + assert client.post.call_count == 2 + sleep.assert_called_once() + + def test_retries_transient_connection_error(self, mock_model_metadata, mock_schema): + backend, client = self._backend_and_client(mock_model_metadata, mock_schema, max_retries=3) + client.post.side_effect = [ + httpx.ConnectError("connection refused"), + make_http_response('{"name": "A"}', 1, "stop"), + ] + with patch(f"{MODULE}.time.sleep"): + backend._complete_one() + assert client.post.call_count == 2 + + def test_persistent_retryable_status_exhausts_attempts(self, mock_model_metadata, mock_schema): + backend, client = self._backend_and_client(mock_model_metadata, mock_schema, max_retries=2) + client.post.return_value = make_status_response(503, "down") + with patch(f"{MODULE}.time.sleep") as sleep, pytest.raises(GenerationError, match="after 3 attempt"): + backend._complete_one() + assert client.post.call_count == 3 # initial + 2 retries + assert sleep.call_count == 2 + + def test_non_retryable_status_fails_fast(self, mock_model_metadata, mock_schema): + backend, client = self._backend_and_client(mock_model_metadata, mock_schema, max_retries=5) + client.post.return_value = make_status_response(400, "bad request") + with patch(f"{MODULE}.time.sleep") as sleep, pytest.raises(GenerationError, match="400"): + backend._complete_one() + assert client.post.call_count == 1 # no retries on a permanent error + sleep.assert_not_called() + + def test_zero_retries_disables_retry(self, mock_model_metadata, mock_schema): + backend, client = self._backend_and_client(mock_model_metadata, mock_schema, max_retries=0) + client.post.return_value = make_status_response(503, "down") + with patch(f"{MODULE}.time.sleep") as sleep, pytest.raises(GenerationError, match="after 1 attempt"): + backend._complete_one() + assert client.post.call_count == 1 + sleep.assert_not_called() + + def test_honors_retry_after_header(self, mock_model_metadata, mock_schema): + backend, client = self._backend_and_client(mock_model_metadata, mock_schema, max_retries=1) + client.post.side_effect = [ + make_status_response(429, "slow down", retry_after="2.5"), + make_http_response('{"name": "A"}', 1, "stop"), + ] + with patch(f"{MODULE}.time.sleep") as sleep: + backend._complete_one() + sleep.assert_called_once_with(2.5) + + +class TestResilienceHelpers: + @pytest.mark.parametrize( + "value,expected", + [ + (5, 5), + (0, 0), + (-3, 0), + (None, 0), + (True, 0), + (False, 0), + ("7", 7), + (7.9, 7), + ("bad", 0), + ([], 0), + ({}, 0), + ], + ) + def test_coerce_token_count(self, value, expected): + assert _coerce_token_count(value) == expected + + def test_parse_retry_after_numeric_seconds(self): + response = MagicMock(spec=httpx.Response) + response.headers = {"Retry-After": "3"} + assert _parse_retry_after(response) == 3.0 + + def test_parse_retry_after_absent(self): + response = MagicMock(spec=httpx.Response) + response.headers = {} + assert _parse_retry_after(response) is None + + def test_parse_retry_after_http_date_ignored(self): + response = MagicMock(spec=httpx.Response) + response.headers = {"Retry-After": "Wed, 21 Oct 2025 07:28:00 GMT"} + assert _parse_retry_after(response) is None + + def test_backoff_delay_uses_retry_after_capped(self): + assert _backoff_delay(0, 999.0) == 30.0 + assert _backoff_delay(5, 3.0) == 3.0 + + def test_backoff_delay_full_jitter_within_bounds(self): + for attempt in range(7): + assert 0.0 <= _backoff_delay(attempt, None) <= 30.0 + + +class TestGenerateBatch: + def test_dispatches_n_requests_and_counts(self, mock_model_metadata, mock_schema): + backend = make_backend(make_params(), mock_model_metadata, mock_schema) + backend.initialize() + try: + with patch.object(backend, "_complete_one", return_value=("rec", 5, "stop")) as complete: + batch = backend._generate_batch(num_prompts_per_batch=3, batch=Batch(processor=MagicMock())) + assert complete.call_count == 3 + assert batch.finish_reasons["stop"] == 3 + assert batch.total_completion_tokens == 15 + finally: + backend.teardown() + + def test_generate_batch_before_init_raises(self, mock_model_metadata, mock_schema): + backend = make_backend(make_params(), mock_model_metadata, mock_schema) + with pytest.raises(InternalError): + backend._generate_batch(num_prompts_per_batch=1, batch=Batch(processor=MagicMock())) + + def test_compacts_pretty_printed_json_before_processing(self, mock_model_metadata, mock_schema): + backend = make_backend(make_params(), mock_model_metadata, mock_schema) + backend._compact_json = True + backend.initialize() + pretty = '{\n "name": "Alice",\n "age": 30\n}' + try: + with patch.object(backend, "_complete_one", return_value=(pretty, 8, "stop")): + batch = Batch(processor=MagicMock()) + with patch.object(batch, "process") as process: + backend._generate_batch(num_prompts_per_batch=1, batch=batch) + finally: + backend.teardown() + _, kwargs = process.call_args + # The multi-line object is collapsed to single-line JSONL the extractor can match. + assert process.call_args[0][1] == '{"name":"Alice","age":30}' + assert kwargs["completion_tokens"] == 8 + + def test_no_compaction_when_flag_disabled(self, mock_model_metadata, mock_schema): + backend = make_backend(make_params(), mock_model_metadata, mock_schema) + backend.initialize() + pretty = '{\n "name": "Alice"\n}' + try: + with patch.object(backend, "_complete_one", return_value=(pretty, 8, "stop")): + batch = Batch(processor=MagicMock()) + with patch.object(batch, "process") as process: + backend._generate_batch(num_prompts_per_batch=1, batch=batch) + finally: + backend.teardown() + # Text passes through untouched when compaction is not armed. + assert process.call_args[0][1] == pretty + + +class TestTeardown: + def test_idempotent(self, mock_model_metadata, mock_schema): + backend = make_backend(make_params(), mock_model_metadata, mock_schema) + backend.initialize() + pool = backend._pool + assert pool is not None + backend.teardown() + backend.teardown() # second call must be a no-op, not an error + assert backend._client is None + assert backend._pool is None + # underlying resources were released + assert pool._shutdown + + +class TestResponseEdgeCases: + def test_missing_usage_defaults_tokens_to_zero(self, mock_model_metadata, mock_schema): + backend = make_backend(make_params(), mock_model_metadata, mock_schema) + backend._client = MagicMock() + response = MagicMock(spec=httpx.Response) + response.status_code = 200 + response.json.return_value = {"choices": [{"text": "x", "finish_reason": "stop"}]} # no usage key + response.raise_for_status.return_value = None + backend._client.post.return_value = response + backend._request_body = {"model": "my-lora", "prompt": "p"} + _, tokens, _ = backend._complete_one() + assert tokens == 0 + + def test_null_completion_tokens_defaults_to_zero(self, mock_model_metadata, mock_schema): + backend = make_backend(make_params(), mock_model_metadata, mock_schema) + backend._client = MagicMock() + response = MagicMock(spec=httpx.Response) + response.status_code = 200 + # A null token count must degrade to 0, not crash the completion. + response.json.return_value = { + "choices": [{"text": "x", "finish_reason": "stop"}], + "usage": {"completion_tokens": None}, + } + response.raise_for_status.return_value = None + backend._client.post.return_value = response + backend._request_body = {"model": "my-lora", "prompt": "p"} + _, tokens, _ = backend._complete_one() + assert tokens == 0 + + def test_malformed_response_becomes_generation_error(self, mock_model_metadata, mock_schema): + backend = make_backend(make_params(), mock_model_metadata, mock_schema) + backend._client = MagicMock() + response = MagicMock(spec=httpx.Response) + response.status_code = 200 + response.json.return_value = {"unexpected": "shape"} # no "choices" + response.raise_for_status.return_value = None + response.text = '{"unexpected": "shape"}' + backend._client.post.return_value = response + backend._request_body = {"model": "my-lora", "prompt": "p"} + with pytest.raises(GenerationError, match="unexpected response shape"): + backend._complete_one() + + def test_none_finish_reason_counted_as_unknown(self, mock_model_metadata, mock_schema): + backend = make_backend(make_params(), mock_model_metadata, mock_schema) + backend.initialize() + try: + with patch.object(backend, "_complete_one", return_value=("rec", 1, None)): + batch = backend._generate_batch(num_prompts_per_batch=2, batch=Batch(processor=MagicMock())) + assert batch.finish_reasons["unknown"] == 2 + finally: + backend.teardown() + + +class TestGenerateEndToEnd: + """Exercise the shared GeneratorBackend.generate() template loop via RemoteBackend. + + Network is mocked at ``_complete_one``, but the real TabularDataProcessor, + Batch, and result aggregation run -- so this guards the base-class refactor. + """ + + def test_full_loop_with_real_processor(self, mock_model_metadata, mock_schema): + config = make_params() + config.generation.num_records = 3 + mock_model_metadata.generation_max_tokens_for.return_value = 128 + processor = create_processor(mock_schema, mock_model_metadata, config) + assert isinstance(processor, TabularDataProcessor) + + backend = make_backend(config, mock_model_metadata, mock_schema, processor=processor) + backend.initialize() + try: + with patch.object(backend, "_complete_one", return_value=('{"name": "Alice", "age": 30}', 8, "stop")): + results = backend.generate() + finally: + backend.teardown() + + assert results.df is not None + assert len(results.df) >= 3 + assert list(results.df.columns) == ["name", "age"] + + +class TestCompactJsonCompletion: + def test_collapses_multiline_object(self): + pretty = '{\n "name": "Alice",\n "age": 30\n}' + assert _compact_json_completion(pretty) == '{"name":"Alice","age":30}' + + def test_strips_surrounding_whitespace(self): + assert _compact_json_completion(' \n {"a": 1}\n ') == '{"a":1}' + + def test_preserves_non_ascii(self): + assert _compact_json_completion('{"city": "São Paulo"}') == '{"city":"São Paulo"}' + + def test_already_compact_object_unchanged(self): + assert _compact_json_completion('{"a":1,"b":2}') == '{"a":1,"b":2}' + + def test_non_json_returned_unchanged(self): + assert _compact_json_completion("not json at all") == "not json at all" + + def test_empty_returned_unchanged(self): + assert _compact_json_completion(" ") == " " + + def test_multiple_objects_returned_unchanged(self): + # Two JSONL records don't parse as one object; the line-oriented + # extractor already handles them, so leave the text untouched. + jsonl = '{"a": 1}\n{"a": 2}' + assert _compact_json_completion(jsonl) == jsonl + + def test_non_object_json_returned_unchanged(self): + # A bare array is valid JSON but not a record object; don't reshape it. + assert _compact_json_completion("[1, 2, 3]") == "[1, 2, 3]" + + +class TestRemoteParametersValidation: + def test_requires_endpoint_and_model(self): + with pytest.raises(ValidationError): + RemoteParameters.model_validate({}) + + def test_timeout_must_be_positive(self): + with pytest.raises(ValidationError): + RemoteParameters(endpoint_url="http://x/v1", model="m", timeout_seconds=0) + + def test_max_concurrency_at_least_one(self): + with pytest.raises(ValidationError): + RemoteParameters(endpoint_url="http://x/v1", model="m", max_concurrency=0) + + +class TestBackendSelectionValidation: + def test_remote_with_timeseries_rejected_at_config_time(self): + # The ParameterError raised in the model validator is wrapped by pydantic. + with pytest.raises(ValidationError, match="time-series"): + SafeSynthesizerParameters( + data=DataParameters(group_training_examples_by="g", order_training_examples_by=None), + training=TrainingHyperparams(pretrained_model="test-model", lora_r=16), + generation=GenerateParameters(remote=RemoteParameters(endpoint_url="http://x/v1", model="m")), + time_series=TimeSeriesParameters(is_timeseries=True, timestamp_column="t"), + ) diff --git a/tests/llm/test_model_host.py b/tests/llm/test_model_host.py new file mode 100644 index 000000000..d106e9ce7 --- /dev/null +++ b/tests/llm/test_model_host.py @@ -0,0 +1,49 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Tests for the local model host contract.""" + +from unittest.mock import MagicMock + +from nemo_safe_synthesizer.generation.vllm_backend import VllmBackend +from nemo_safe_synthesizer.llm.model_host import ModelHost + + +def test_vllm_backend_is_a_typed_local_model_host() -> None: + """The local generation backend exposes its model through the shared host boundary.""" + assert issubclass(VllmBackend, ModelHost) + + +def test_vllm_backend_model_delegates_to_loaded_engine() -> None: + """The shared model property returns the engine already owned by the backend.""" + engine = object() + backend = object.__new__(VllmBackend) + backend.llm = engine # type: ignore[assignment] + + assert backend.model is engine + + +def test_vllm_backend_model_is_none_before_initialization() -> None: + """Callers can inspect whether the host has loaded its model yet.""" + backend = object.__new__(VllmBackend) + backend.llm = None + + assert backend.model is None + + +def test_vllm_backend_tokenizer_delegates_to_loaded_engine() -> None: + """The host exposes the tokenizer owned by its loaded engine.""" + tokenizer = object() + backend = object.__new__(VllmBackend) + backend.llm = MagicMock() + backend.llm.get_tokenizer.return_value = tokenizer + + assert backend.tokenizer is tokenizer + + +def test_vllm_backend_tokenizer_is_none_before_initialization() -> None: + """The tokenizer follows the hosted model's lifecycle.""" + backend = object.__new__(VllmBackend) + backend.llm = None + + assert backend.tokenizer is None diff --git a/tests/smoke/README.md b/tests/smoke/README.md index 13fdacac0..80cee9a49 100644 --- a/tests/smoke/README.md +++ b/tests/smoke/README.md @@ -22,6 +22,18 @@ NSS_TELEMETRY_SMOKE_SEND=1 uv run --frozen pytest tests/smoke/test_telemetry_smo Set `NEMO_TELEMETRY_ENDPOINT` to point at a local or controlled endpoint if you do not want to contact the default NVIDIA telemetry endpoint. +The remote generation backend has two smoke layers: + +- `test_remote_generation_cpu.py` -- always runs in `mise run test:smoke`. Drives the full `RemoteBackend` path (real `httpx`, concurrency, retry, JSON compaction, processor) against an in-process loopback HTTP stub. No GPU, no network, no model. +- `test_remote_generation_live.py` -- opt-in, makes real billable calls to build.nvidia.com: + +```bash +NSS_REMOTE_SMOKE_SEND=1 NVIDIA_API_KEY=nvapi-... \ + uv run --frozen pytest tests/smoke/test_remote_generation_live.py -vvs -n0 +``` + +Override `NSS_REMOTE_SMOKE_ENDPOINT` / `NSS_REMOTE_SMOKE_MODEL` to target a different OpenAI-compatible server or model. The model must expose `/v1/completions` (the backend uses text completions, not chat); the test skips rather than fails when a model rejects that route. + ## When should I add a smoke test? If you're adding a new training backend, generation backend, evaluation diff --git a/tests/smoke/conftest.py b/tests/smoke/conftest.py index 6981da342..62ef47af0 100644 --- a/tests/smoke/conftest.py +++ b/tests/smoke/conftest.py @@ -15,7 +15,17 @@ from transformers import AutoTokenizer, LlamaConfig, LlamaForCausalLM, PreTrainedTokenizerBase from nemo_safe_synthesizer.cli.artifact_structure import Workdir +from nemo_safe_synthesizer.config import ( + DataParameters, + GenerateParameters, + RemoteParameters, + StructuredGenerationParameters, + TrainingHyperparams, +) +from nemo_safe_synthesizer.config.generate import RemoteDialect, StructuredGenerationSchemaMethod from nemo_safe_synthesizer.config.parameters import SafeSynthesizerParameters +from nemo_safe_synthesizer.defaults import DEFAULT_INSTRUCTION, PROMPT_TEMPLATE +from nemo_safe_synthesizer.llm.metadata import LLMPromptConfig, ModelMetadata from nemo_safe_synthesizer.sdk.library_builder import SafeSynthesizer @@ -187,6 +197,76 @@ def train_with_sdk(config: SafeSynthesizerParameters, data_df: pd.DataFrame, sav return nss +def build_remote_metadata() -> ModelMetadata: + """Tokenizer-free ``ModelMetadata`` for remote-backend smoke tests. + + The remote backend never loads a model, so this bypasses the HuggingFace + config/tokenizer load via ``model_construct`` and supplies only the fields + the backend reads: the instruction, the prompt template and BOS/EOS tokens + (for structured generation), and a context window for the ``max_tokens`` + clamp. ``tokenizer=None`` mirrors the offline-remote path, so the + prompt-length clamp is disabled and no download is forced. + """ + return ModelMetadata.model_construct( + model_name_or_path="remote-stub", + autoconfig=None, + base_max_seq_length=2048, + rope_scaling=None, + max_tokens_per_example=None, + tokenizer=None, + instruction=DEFAULT_INSTRUCTION, + prompt_config=LLMPromptConfig( + template=PROMPT_TEMPLATE, + add_bos_token_to_prompt=True, + add_eos_token_to_prompt=True, + bos_token="", + bos_token_id=1, + eos_token="", + eos_token_id=2, + ), + ) + + +def build_remote_config( + *, + endpoint_url: str, + model: str, + dialect: RemoteDialect = "vllm", + num_records: int = 5, + api_key_env: str | None = None, + max_retries: int = 4, + max_concurrency: int = 4, + use_structured_generation: bool = False, + structured_generation_schema_method: StructuredGenerationSchemaMethod = "auto", +) -> SafeSynthesizerParameters: + """Minimal tabular ``SafeSynthesizerParameters`` with a configured remote endpoint. + + No model is loaded at construction time, so ``pretrained_model`` is an + arbitrary placeholder. ``group_training_examples_by=None`` and the default + (non-time-series) ``time_series`` keep the backend on the plain tabular + processor path. + """ + return SafeSynthesizerParameters( + data=DataParameters(group_training_examples_by=None, order_training_examples_by=None), + training=TrainingHyperparams(pretrained_model="remote-stub", lora_r=16), + generation=GenerateParameters( + num_records=num_records, + structured_generation=StructuredGenerationParameters( + enabled=use_structured_generation, + schema_method=structured_generation_schema_method, + ), + remote=RemoteParameters( + endpoint_url=endpoint_url, + model=model, + dialect=dialect, + api_key_env=api_key_env, + max_retries=max_retries, + max_concurrency=max_concurrency, + ), + ), + ) + + @pytest.fixture(scope="session") def _patch_attn_eager() -> Generator[None, None, None]: """Override attn_implementation from 'flashinfer' (not a valid HF option) to 'sdpa'. diff --git a/tests/smoke/test_remote_generation_cpu.py b/tests/smoke/test_remote_generation_cpu.py new file mode 100644 index 000000000..a434c67f8 --- /dev/null +++ b/tests/smoke/test_remote_generation_cpu.py @@ -0,0 +1,197 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""CPU smoke tests for the remote generation backend against a local stub server. + +These run the full ``RemoteBackend`` path -- real ``httpx`` client, real +``ThreadPoolExecutor`` concurrency, the shared ``generate()`` batch loop, the +real ``TabularDataProcessor``, retry/backoff, and JSON compaction -- over a real +loopback socket. The server is an in-process stdlib HTTP server, so there is no +GPU, no model load, and no external network. The mocked unit suite in +``tests/generation/test_remote_backend.py`` covers fine-grained behavior; this +suite proves the pieces work together end-to-end over the wire. +""" + +from __future__ import annotations + +import json +import threading +from collections.abc import Callable, Iterator +from contextlib import contextmanager +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer +from pathlib import Path +from typing import Any +from unittest.mock import MagicMock + +import pytest + +from nemo_safe_synthesizer.cli.artifact_structure import Workdir +from nemo_safe_synthesizer.generation import remote_backend +from nemo_safe_synthesizer.generation.remote_backend import RemoteBackend + +from .conftest import build_remote_config, build_remote_metadata + +# Responder contract: (path, parsed_body) -> (status_code, json_payload, extra_headers | None) +Responder = Callable[[str, dict[str, Any]], tuple[int, dict[str, Any], dict[str, str] | None]] + +SCHEMA = {"properties": {"name": {"type": "string"}, "age": {"type": "integer"}}} + + +@contextmanager +def stub_completions_server(responder: Responder) -> Iterator[str]: + """Run a loopback HTTP server driven by ``responder``; yield its ``/v1`` base URL. + + Bound to an ephemeral port on 127.0.0.1 and served from a daemon thread, so + concurrent requests from the backend's worker pool are handled in parallel. + """ + + class Handler(BaseHTTPRequestHandler): + def do_POST(self) -> None: # noqa: N802 -- BaseHTTPRequestHandler API + length = int(self.headers.get("Content-Length") or 0) + raw = self.rfile.read(length) if length else b"" + try: + body = json.loads(raw) if raw else {} + except json.JSONDecodeError: + body = {} + status, payload, headers = responder(self.path, body) + data = json.dumps(payload).encode() + self.send_response(status) + self.send_header("Content-Type", "application/json") + for key, value in (headers or {}).items(): + self.send_header(key, value) + self.send_header("Content-Length", str(len(data))) + self.end_headers() + self.wfile.write(data) + + def log_message(self, format: str, *args: Any) -> None: # noqa: A002 -- match base signature + pass # silence per-request stderr noise + + server = ThreadingHTTPServer(("127.0.0.1", 0), Handler) + thread = threading.Thread(target=server.serve_forever, daemon=True) + thread.start() + try: + address = server.server_address + yield f"http://{address[0]}:{address[1]}/v1" + finally: + server.shutdown() + server.server_close() + thread.join(timeout=5) + + +def _completion(text: str, *, completion_tokens: int = 6) -> dict[str, Any]: + """A minimal OpenAI-style ``/completions`` response body wrapping ``text``.""" + return { + "choices": [{"text": text, "finish_reason": "stop"}], + "usage": {"completion_tokens": completion_tokens}, + } + + +def make_backend(config: Any, schema_path: Path) -> RemoteBackend: + """Build a ``RemoteBackend`` reading a real on-disk schema (everything else real).""" + workdir = MagicMock(spec=Workdir) + workdir.schema_file = schema_path + return RemoteBackend(config=config, model_metadata=build_remote_metadata(), workdir=workdir) + + +@pytest.fixture +def schema_path(tmp_path: Path) -> Path: + path = tmp_path / "schema.json" + path.write_text(json.dumps(SCHEMA)) + return path + + +def test_happy_path_generates_records(schema_path: Path) -> None: + """Full initialize -> generate loop yields the requested records over a real socket.""" + counter, lock = [0], threading.Lock() + + def responder(path: str, body: dict[str, Any]) -> tuple[int, dict[str, Any], None]: + assert path == "/v1/completions" + assert body["prompt"], "prompt should be baked into the request body" + assert body["n"] == 1 + with lock: + counter[0] += 1 + n = counter[0] + return 200, _completion(json.dumps({"name": f"person_{n}", "age": 20 + (n % 50)})), None + + with stub_completions_server(responder) as base_url: + config = build_remote_config(endpoint_url=base_url, model="stub-model", num_records=5) + backend = make_backend(config, schema_path) + try: + backend.initialize() + results = backend.generate() + finally: + backend.teardown() + backend.teardown() # idempotent + + assert results.num_valid_records >= 5 + assert list(results.df.columns) == ["name", "age"] + assert len(results.df) == 5 # truncated to num_records on completion + + +def test_recovers_from_transient_failures(schema_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: + """503 responses are retried with backoff until the server recovers.""" + monkeypatch.setattr(remote_backend, "_BACKOFF_BASE_SECONDS", 0.001) + monkeypatch.setattr(remote_backend, "_BACKOFF_MAX_SECONDS", 0.01) + counter, lock = [0], threading.Lock() + fail_first = 2 + + def responder(path: str, body: dict[str, Any]) -> tuple[int, dict[str, Any], dict[str, str] | None]: + with lock: + counter[0] += 1 + n = counter[0] + if n <= fail_first: + return 503, {"error": "server warming up"}, {"Retry-After": "0"} + return 200, _completion(json.dumps({"name": f"person_{n}", "age": 20 + (n % 50)})), None + + with stub_completions_server(responder) as base_url: + # Serialize requests so the transient failures land on the first record's attempts. + config = build_remote_config( + endpoint_url=base_url, model="stub-model", num_records=3, max_concurrency=1, max_retries=5 + ) + backend = make_backend(config, schema_path) + try: + backend.initialize() + results = backend.generate() + finally: + backend.teardown() + + assert results.num_valid_records >= 3 + assert counter[0] > 3, "expected retries beyond the initial 503 responses" + + +def test_compacts_pretty_printed_json(schema_path: Path) -> None: + """Multi-line JSON from a json_schema-constrained server is compacted before parsing. + + The line-oriented record extractor cannot match an object spanning newlines, + so without ``_compact_json`` the run would yield zero records. ``openai`` + dialect can't send the vLLM ``disable_any_whitespace`` source fix, making the + client-side compaction net the only thing that rescues the output. + """ + counter, lock = [0], threading.Lock() + + def responder(path: str, body: dict[str, Any]) -> tuple[int, dict[str, Any], None]: + with lock: + counter[0] += 1 + n = counter[0] + pretty = json.dumps({"name": f"person_{n}", "age": 20 + (n % 50)}, indent=2) + assert "\n" in pretty + return 200, _completion(pretty, completion_tokens=12), None + + with stub_completions_server(responder) as base_url: + config = build_remote_config( + endpoint_url=base_url, + model="stub-model", + num_records=3, + dialect="openai", + use_structured_generation=True, + structured_generation_schema_method="json_schema", + ) + backend = make_backend(config, schema_path) + try: + backend.initialize() + results = backend.generate() + finally: + backend.teardown() + + assert results.num_valid_records >= 3 + assert len(results.df) == 3 diff --git a/tests/smoke/test_remote_generation_live.py b/tests/smoke/test_remote_generation_live.py new file mode 100644 index 000000000..5b2f8a798 --- /dev/null +++ b/tests/smoke/test_remote_generation_live.py @@ -0,0 +1,82 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +r"""Opt-in live smoke test for the remote backend against build.nvidia.com. + +Disabled by default because it makes real, billable network calls. Enable with:: + + NSS_REMOTE_SMOKE_SEND=1 NVIDIA_API_KEY=nvapi-... \ + uv run --frozen pytest tests/smoke/test_remote_generation_live.py -vvs -n0 + +Overridable via env: ``NSS_REMOTE_SMOKE_ENDPOINT`` (default +``https://integrate.api.nvidia.com/v1``) and ``NSS_REMOTE_SMOKE_MODEL``. The +model must expose the OpenAI ``/v1/completions`` route (not all catalog models +do); the test skips rather than fails when the endpoint rejects that route. + +This validates the live HTTP transport (auth header, request body, response +parsing, retry) -- not record quality, since an instruct model prompted on the +raw completions route is not expected to emit clean JSONL. +""" + +from __future__ import annotations + +import json +import os +from pathlib import Path +from unittest.mock import MagicMock + +import pytest + +from nemo_safe_synthesizer.cli.artifact_structure import Workdir +from nemo_safe_synthesizer.errors import GenerationError +from nemo_safe_synthesizer.generation.remote_backend import RemoteBackend + +from .conftest import build_remote_config, build_remote_metadata + +LIVE_ENDPOINT = os.environ.get("NSS_REMOTE_SMOKE_ENDPOINT", "https://integrate.api.nvidia.com/v1") +LIVE_MODEL = os.environ.get("NSS_REMOTE_SMOKE_MODEL", "meta/llama-3.1-8b-instruct") +LIVE_API_KEY_ENV = "NVIDIA_API_KEY" + +pytestmark = pytest.mark.skipif( + os.environ.get("NSS_REMOTE_SMOKE_SEND") != "1", + reason="opt-in live remote smoke; set NSS_REMOTE_SMOKE_SEND=1 (and NVIDIA_API_KEY) to run", +) + +SCHEMA = {"properties": {"name": {"type": "string"}, "age": {"type": "integer"}}} + + +def test_live_build_nvidia_completions(tmp_path: Path) -> None: + if not os.environ.get(LIVE_API_KEY_ENV): + pytest.skip(f"{LIVE_API_KEY_ENV} not set") + + schema_path = tmp_path / "schema.json" + schema_path.write_text(json.dumps(SCHEMA)) + workdir = MagicMock(spec=Workdir) + workdir.schema_file = schema_path + + config = build_remote_config( + endpoint_url=LIVE_ENDPOINT, + model=LIVE_MODEL, + dialect="openai", # NIM is a strict OpenAI server; drop the vLLM-only fields + api_key_env=LIVE_API_KEY_ENV, + num_records=5, + max_concurrency=2, + max_retries=2, + ) + backend = RemoteBackend(config=config, model_metadata=build_remote_metadata(), workdir=workdir) + try: + backend.initialize() + try: + results = backend.generate() + except GenerationError as exc: + # A model that does not serve /v1/completions answers 400/404 on every + # request; that is an availability gap, not a transport bug, so skip. + if "400" in str(exc) or "404" in str(exc): + pytest.skip(f"model {LIVE_MODEL!r} did not accept /v1/completions: {exc}") + raise + finally: + backend.teardown() + + # Transport round-tripped and the batch loop ran. Record validity is not + # asserted: quality from an instruct model on a raw completions prompt varies. + assert results.num_prompts > 0 diff --git a/tools/vllm_debug.py b/tools/vllm_debug.py new file mode 100755 index 000000000..7c2bb23a0 --- /dev/null +++ b/tools/vllm_debug.py @@ -0,0 +1,473 @@ +#!/usr/bin/env -S uv run --script +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# /// script +# requires-python = ">=3.13" +# dependencies = [ +# "cyclopts>=3", +# "httpx>=0.27", +# "pydantic>=2", +# "rich>=13", +# "structlog>=24", +# ] +# /// +r"""vllm-debug: spin up a vLLM model and fire a few debug calls against it. + +A standalone companion to the VllmBackend / RemoteBackend code: serve a base +model (optionally with a trained LoRA adapter attached) and poke it over the +OpenAI-compatible API to see exactly what it emits. ``serve`` runs vLLM from the +project venv tuned for this A100 box (offline HF cache, FLASH_ATTN; CUDA graphs +on by default, ``--eager`` as an escape hatch); ``call`` and ``models`` are +lightweight and need only this script's own deps. + +Usage:: + + # Serve the local Nemotron text model on :8000 (Ctrl-C to stop) + uv run tools/vllm_debug.py serve + + # Serve a base model + a trained LoRA adapter, registered as model `lora` + uv run tools/vllm_debug.py serve meta-llama/Llama-3.2-1B \ + --adapter artifacts/.../train/adapter --max-lora-rank 32 + + # Print the launch command without starting it (no GPU touched) + uv run tools/vllm_debug.py serve --dry-run + + # One chat call against a running server, reasoning disabled + uv run tools/vllm_debug.py call "Say hello" --base-url http://localhost:8000/v1 + + # Debug structured generation: constrain a text completion to a JSON schema + uv run tools/vllm_debug.py call "a person" --mode text --json-schema schema.json + + # Constrain to a regex, repeat 3 times, emit machine-readable output + uv run tools/vllm_debug.py call "a person" --mode text --regex '\{.*\}' -n 3 --json + + # Health check / list served models + uv run tools/vllm_debug.py models --base-url http://localhost:8000/v1 + +Exit codes: 0 success, 1 a call failed, 125 bad input. +""" + +import json +import os +import sys +import time +from pathlib import Path +from typing import Annotated, Any, Literal, NoReturn, Self + +import cyclopts +import httpx +import structlog +from pydantic import BaseModel, Field +from rich.console import Console +from rich.table import Table + +DEFAULT_MODEL = "nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16" +DEFAULT_BASE_URL = "http://localhost:8000/v1" +REPO_ROOT = Path(__file__).resolve().parents[1] + +LogFormat = Literal["plain", "json"] +CallMode = Literal["chat", "text"] + +console = Console() +log = structlog.get_logger() + + +def configure_logging(fmt: LogFormat) -> None: + """Set up structlog to render to stderr in the chosen format.""" + structlog.configure( + processors=[ + structlog.stdlib.add_log_level, + structlog.dev.ConsoleRenderer() if fmt == "plain" else structlog.processors.JSONRenderer(), + ], + logger_factory=structlog.PrintLoggerFactory(file=sys.stderr), + ) + + +def bad_input(message: str) -> NoReturn: + """Log a usage error and exit with code 125.""" + log.error("bad-input", detail=message) + raise SystemExit(125) + + +# --------------------------------------------------------------------------- # +# serve +# --------------------------------------------------------------------------- # + + +class ServeConfig(BaseModel): + """Resolved arguments for launching a vLLM OpenAI-compatible server.""" + + model: str + port: int + host: str + max_model_len: int + gpu_util: float + max_num_seqs: int + attention_backend: str + enforce_eager: bool + prefix_caching: bool + trust_remote_code: bool + adapter: Path | None + lora_name: str + max_lora_rank: int + python: Path + + def argv(self) -> list[str]: + """Build the vLLM api_server command line.""" + cmd = [ + str(self.python), "-m", "vllm.entrypoints.openai.api_server", + "--model", self.model, + "--served-model-name", self.model, + "--port", str(self.port), + "--host", self.host, + "--max-model-len", str(self.max_model_len), + "--gpu-memory-utilization", str(self.gpu_util), + "--max-num-seqs", str(self.max_num_seqs), + "--attention-config", json.dumps({"backend": self.attention_backend}), + ] # fmt: skip + if self.trust_remote_code: + cmd.append("--trust-remote-code") + if self.enforce_eager: + cmd.append("--enforce-eager") + if self.prefix_caching: + cmd.append("--enable-prefix-caching") + if self.adapter is not None: + cmd += ["--enable-lora", "--lora-modules", f"{self.lora_name}={self.adapter}"] + cmd += ["--max-lora-rank", str(self.max_lora_rank)] + cmd.append("--no-enable-log-requests") + return cmd + + def server_env(self) -> dict[str, str]: + """Environment for the server: force the offline HF cache, keep FP8 off the Triton path.""" + return os.environ | { + "HF_HUB_OFFLINE": "1", + "TRANSFORMERS_OFFLINE": "1", + "VLLM_TEST_FORCE_FP8_MARLIN": "1", + } + + +def _default_python() -> Path: + """Path to the project venv's interpreter (where vLLM is installed).""" + return REPO_ROOT / ".venv" / "bin" / "python" + + +app = cyclopts.App(name="vllm-debug", help="Serve a vLLM model and fire debug calls against it.") + + +@app.command +def serve( + model: str = DEFAULT_MODEL, + *, + port: int = 8000, + host: str = "0.0.0.0", + adapter: Annotated[Path | None, cyclopts.Parameter(help="LoRA adapter dir to attach and serve")] = None, + lora_name: Annotated[str, cyclopts.Parameter(help="Model name to register the adapter under")] = "lora", + max_lora_rank: int = 32, + max_model_len: int = 8192, + gpu_util: Annotated[float, cyclopts.Parameter(name="--gpu-util")] = 0.90, + max_num_seqs: int = 128, + attention_backend: str = "FLASH_ATTN", + eager: Annotated[ + bool, + cyclopts.Parameter( + help="Force --enforce-eager, disabling torch.compile + CUDA graphs. " + "Only needed when FlashInfer is broken (see serve docstring); off by default." + ), + ] = False, + prefix_caching: bool = True, + trust_remote_code: bool = True, + python: Annotated[Path | None, cyclopts.Parameter(help="Interpreter to run vLLM (default: project venv)")] = None, + dry_run: Annotated[bool, cyclopts.Parameter(name="--dry-run", help="Print the command and exit")] = False, +) -> None: + """Launch a vLLM OpenAI-compatible server, optionally with a LoRA adapter attached. + + Replaces this process with the server (so Ctrl-C / logs behave normally). + Defaults are tuned for this A100 box: a small context window, high GPU + fraction, and the FLASH_ATTN attention backend. + + CUDA graphs are enabled by default (``--enforce-eager`` off) for throughput. + Pass ``--eager`` only if the venv's FlashInfer is broken -- vLLM's + torch.compile path imports ``flashinfer.comm`` unconditionally, so a broken + FlashInfer crashes engine init unless eager mode skips compilation. (A broken + FlashInfer usually means a partial package install; ``uv sync + --reinstall-package flashinfer-cubin`` repairs it.) + """ + if adapter is not None and not adapter.exists(): + bad_input(f"adapter path does not exist: {adapter}") + interpreter = python or _default_python() + cfg = ServeConfig( + model=model, port=port, host=host, max_model_len=max_model_len, gpu_util=gpu_util, + max_num_seqs=max_num_seqs, attention_backend=attention_backend, enforce_eager=eager, + prefix_caching=prefix_caching, trust_remote_code=trust_remote_code, adapter=adapter, + lora_name=lora_name, max_lora_rank=max_lora_rank, python=interpreter, + ) # fmt: skip + argv, env = cfg.argv(), cfg.server_env() + served_as = lora_name if adapter is not None else model + + if dry_run: + console.print("[bold]HF_HUB_OFFLINE=1 TRANSFORMERS_OFFLINE=1 VLLM_TEST_FORCE_FP8_MARLIN=1[/bold]") + console.print(" ".join(argv)) + console.print(f"[dim]→ call it with --base-url http://localhost:{port}/v1 --model {served_as}[/dim]") + return + + if not interpreter.exists(): + bad_input(f"interpreter not found: {interpreter} (pass --python or run `uv sync`)") + log.info("serving", model=model, port=port, adapter=str(adapter) if adapter else None, served_as=served_as) + os.execve(argv[0], argv, env) # replace this process with the server + + +# --------------------------------------------------------------------------- # +# call +# --------------------------------------------------------------------------- # + + +class TokenUsage(BaseModel): + """Token counts reported by the server's ``usage`` block.""" + + prompt_tokens: int = 0 + completion_tokens: int = 0 + total_tokens: int = 0 + + @classmethod + def from_openai(cls, usage: dict[str, Any]) -> Self: + def _i(value: Any) -> int: + try: + return max(0, int(value)) + except (TypeError, ValueError): + return 0 + + prompt, completion = _i(usage.get("prompt_tokens")), _i(usage.get("completion_tokens")) + return cls( + prompt_tokens=prompt, + completion_tokens=completion, + total_tokens=_i(usage.get("total_tokens")) or (prompt + completion), + ) + + +class CallResult(BaseModel): + """Outcome of a single debug call.""" + + content: str = "" + reasoning: str | None = None + finish_reason: str | None = None + usage: TokenUsage = Field(default_factory=TokenUsage) + latency_s: float = 0.0 + error: str | None = None + + +def _read_text_file(path: Path, flag: str) -> str: + """Read a file for ``flag``, exiting with a clean usage error on failure.""" + try: + return path.read_text() + except OSError as exc: + bad_input(f"{flag}: cannot read {path}: {exc}") + + +def _read_json_file(path: Path, flag: str) -> Any: + """Read and parse a JSON file for ``flag``, exiting cleanly on read/parse failure.""" + text = _read_text_file(path, flag) + try: + return json.loads(text) + except json.JSONDecodeError as exc: + bad_input(f"{flag}: {path} is not valid JSON: {exc}") + + +def build_structured_outputs(json_schema: Path | None, regex: str | None, structural_tag: Path | None) -> dict | None: + """Build the vLLM ``structured_outputs`` field from at most one constraint flag. + + Validates the at-most-one rule before touching the filesystem, so passing + two flags fails fast without a partial read. Missing or malformed files + exit via ``bad_input`` (code 125) rather than raising a traceback. + """ + candidates = (("json", json_schema), ("regex", regex), ("structural_tag", structural_tag)) + chosen = [(key, value) for key, value in candidates if value is not None] + if not chosen: + return None + if len(chosen) > 1: + bad_input("pass at most one of --json-schema / --regex / --structural-tag") + + key, value = chosen[0] + if key == "json": + return {"json": _read_json_file(value, "--json-schema")} + if key == "structural_tag": + return {"structural_tag": _read_text_file(value, "--structural-tag")} + return {"regex": value} + + +def build_call_payload( + prompt: str, + *, + mode: CallMode, + model: str, + max_tokens: int, + temperature: float, + top_p: float, + repetition_penalty: float, + think: bool, + structured_outputs: dict | None, +) -> dict[str, Any]: + """Assemble the request body for chat or text completion.""" + payload: dict[str, Any] = { + "model": model, + "max_tokens": max_tokens, + "temperature": temperature, + "top_p": top_p, + "repetition_penalty": repetition_penalty, + } + if mode == "chat": + payload["messages"] = [{"role": "user", "content": prompt}] + if not think: + payload["chat_template_kwargs"] = {"enable_thinking": False} + else: + payload["prompt"] = prompt + if structured_outputs is not None: + payload["structured_outputs"] = structured_outputs + return payload + + +def parse_call_response(mode: CallMode, data: dict[str, Any]) -> tuple[str, str | None, str | None]: + """Extract ``(content, reasoning, finish_reason)`` from a completion response. + + Raises ``ValueError`` on a response with no choices; ``do_call`` catches it + and reports it as a failed ``CallResult`` rather than crashing the run. + """ + choices = data.get("choices") or [] + if not choices: + raise ValueError(f"response contained no choices: {json.dumps(data)[:300]}") + choice = choices[0] + if mode == "chat": + message = choice.get("message") or {} + return message.get("content") or "", message.get("reasoning_content"), choice.get("finish_reason") + return choice.get("text") or "", None, choice.get("finish_reason") + + +def do_call(client: httpx.Client, payload: dict[str, Any], *, mode: CallMode, api_key: str) -> CallResult: + """Issue one completion request and capture text, usage, and latency.""" + path = "/chat/completions" if mode == "chat" else "/completions" + started = time.monotonic() + try: + response = client.post(path, json=payload, headers={"Authorization": f"Bearer {api_key}"}) + response.raise_for_status() + data = response.json() + content, reasoning, finish = parse_call_response(mode, data) + return CallResult( + content=content, + reasoning=reasoning, + finish_reason=finish, + usage=TokenUsage.from_openai(data.get("usage") or {}), + latency_s=round(time.monotonic() - started, 3), + ) + except Exception as exc: # noqa: BLE001 -- debug tool: report any failure as a result + return CallResult(latency_s=round(time.monotonic() - started, 3), error=f"{type(exc).__name__}: {exc}") + + +def render_calls(results: list[CallResult], *, json_output: bool) -> None: + """Print call results as rich panels, or as a JSON array with --json.""" + if json_output: + print(json.dumps([r.model_dump() for r in results], indent=2)) + return + table = Table(title="vllm-debug call", show_lines=True) + for col in ("#", "finish", "prompt tok", "completion tok", "latency s", "content"): + table.add_column(col, overflow="fold") + for i, r in enumerate(results, 1): + body = r.error and f"[red]{r.error}[/red]" or r.content + if r.reasoning: + body = f"[dim](reasoning: {len(r.reasoning)} chars hidden)[/dim]\n{body}" + table.add_row( + str(i), r.finish_reason or "-", str(r.usage.prompt_tokens), + str(r.usage.completion_tokens), f"{r.latency_s:.2f}", body, + ) # fmt: skip + console.print(table) + + +@app.command +def call( + prompt: Annotated[str, cyclopts.Parameter(help="Prompt text, or '-' to read stdin")], + *, + base_url: Annotated[str, cyclopts.Parameter(name="--base-url")] = DEFAULT_BASE_URL, + model: str = DEFAULT_MODEL, + mode: CallMode = "chat", + max_tokens: Annotated[int, cyclopts.Parameter(name="--max-tokens")] = 256, + temperature: float = 0.0, + top_p: Annotated[float, cyclopts.Parameter(name="--top-p")] = 1.0, + repetition_penalty: Annotated[float, cyclopts.Parameter(name="--repetition-penalty")] = 1.0, + think: Annotated[bool, cyclopts.Parameter(help="Allow reasoning CoT (chat mode; default off)")] = False, + json_schema: Annotated[ + Path | None, cyclopts.Parameter(name="--json-schema", help="Constrain to JSON schema") + ] = None, + regex: Annotated[str | None, cyclopts.Parameter(help="Constrain output to a regex")] = None, + structural_tag: Annotated[ + Path | None, cyclopts.Parameter(name="--structural-tag", help="XGrammar tag file") + ] = None, + n: Annotated[int, cyclopts.Parameter(help="Number of times to repeat the call")] = 1, + api_key: Annotated[str, cyclopts.Parameter(name="--api-key")] = "EMPTY", + request_timeout: Annotated[float, cyclopts.Parameter(name="--request-timeout")] = 600.0, + json_output: Annotated[bool, cyclopts.Parameter(name="--json")] = False, +) -> None: + """Send one or more debug calls to a running server and show text + token usage. + + Structured generation maps to vLLM's ``structured_outputs`` field; pass at + most one of --json-schema / --regex / --structural-tag. + """ + if n < 1: + bad_input("-n must be >= 1") + text = sys.stdin.read() if prompt == "-" else prompt + structured = build_structured_outputs(json_schema, regex, structural_tag) + payload = build_call_payload( + text, mode=mode, model=model, max_tokens=max_tokens, temperature=temperature, + top_p=top_p, repetition_penalty=repetition_penalty, think=think, structured_outputs=structured, + ) # fmt: skip + + with httpx.Client(base_url=base_url.rstrip("/"), timeout=httpx.Timeout(request_timeout)) as client: + results = [do_call(client, payload, mode=mode, api_key=api_key) for _ in range(n)] + render_calls(results, json_output=json_output) + if any(r.error for r in results): + raise SystemExit(1) + + +# --------------------------------------------------------------------------- # +# models +# --------------------------------------------------------------------------- # + + +@app.command +def models( + *, + base_url: Annotated[str, cyclopts.Parameter(name="--base-url")] = DEFAULT_BASE_URL, + api_key: Annotated[str, cyclopts.Parameter(name="--api-key")] = "EMPTY", + json_output: Annotated[bool, cyclopts.Parameter(name="--json")] = False, +) -> None: + """List the models a running server is serving (also a health check).""" + try: + with httpx.Client(base_url=base_url.rstrip("/"), timeout=httpx.Timeout(30.0)) as client: + response = client.get("/models", headers={"Authorization": f"Bearer {api_key}"}) + response.raise_for_status() + entries = response.json().get("data", []) + except Exception as exc: # noqa: BLE001 + log.error("models-failed", base_url=base_url, detail=f"{type(exc).__name__}: {exc}") + raise SystemExit(1) from exc + + if json_output: + print(json.dumps(entries, indent=2)) + return + table = Table(title=f"served models @ {base_url}") + table.add_column("id") + table.add_column("max_model_len") + table.add_column("root", overflow="fold") + for entry in entries: + table.add_row(entry.get("id", "?"), str(entry.get("max_model_len", "?")), entry.get("root", "")) + console.print(table) + + +@app.meta.default +def _launcher( + *tokens: Annotated[str, cyclopts.Parameter(show=False, allow_leading_hyphen=True)], + log_format: Annotated[LogFormat, cyclopts.Parameter(name="--log-format")] = "plain", +) -> None: + """Configure logging, then dispatch to the requested subcommand.""" + configure_logging(log_format) + app(tokens) + + +if __name__ == "__main__": + app.meta()