diff --git a/components/src/dynamo/frontend/vllm_processor.py b/components/src/dynamo/frontend/vllm_processor.py index f026ea3d1499..377a3910634d 100644 --- a/components/src/dynamo/frontend/vllm_processor.py +++ b/components/src/dynamo/frontend/vllm_processor.py @@ -376,6 +376,40 @@ async def _generate_and_stream( for output in vllm_out.request_outputs[0].outputs: choice = post.process_output(output) if choice: + # ── RL logprobs injection ────────────────────── + # The vLLM worker sends log_probs/top_logprobs in + # the engine_response dict. Since we can't easily + # construct LogprobsLists for EngineCoreOutput, we + # inject them directly into the choice here. + worker_log_probs = engine_response.get("log_probs") + worker_top_logprobs = engine_response.get("top_logprobs") + if worker_log_probs is not None and choice.get("logprobs") is None: + oai_logprobs_content = [] + new_tids = engine_response.get("token_ids", []) + for i, lp in enumerate(worker_log_probs): + # Always populate token/bytes so consumers never see a + # missing key. If top_logprobs is absent or the token + # string cannot be resolved we fall back to the numeric + # ID as a string — better than a KeyError / silent None. + tid_str = str(new_tids[i]) if i < len(new_tids) else "" + entry: dict = { + "logprob": lp, + "token": tid_str, + "bytes": None, + } + # Resolve the human-readable token string and top_logprobs + # from the engine's top_logprobs table when available. + if worker_top_logprobs and i < len(worker_top_logprobs): + tops = worker_top_logprobs[i] + entry["top_logprobs"] = tops + if i < len(new_tids): + for tp in tops: + if tp.get("token_id") == new_tids[i]: + entry["token"] = tp.get("token", tid_str) + break + oai_logprobs_content.append(entry) + choice["logprobs"] = {"content": oai_logprobs_content} + choices.append(choice) if choices: @@ -389,6 +423,11 @@ async def _generate_and_stream( if usage := engine_response.get("completion_usage"): dynamo_out["usage"] = usage + # ── RL: pass output token IDs for nvext.completion_token_ids ── + new_token_ids = engine_response.get("token_ids", []) + if new_token_ids: + dynamo_out["_completion_token_ids"] = new_token_ids + yield dynamo_out finally: if vllm_preproc.request_id in self.output_processor.request_states: diff --git a/components/src/dynamo/vllm/handlers.py b/components/src/dynamo/vllm/handlers.py index d7d70e8ec10d..ad22ef147b37 100644 --- a/components/src/dynamo/vllm/handlers.py +++ b/components/src/dynamo/vllm/handlers.py @@ -625,6 +625,280 @@ async def wake_up(self, body: dict) -> dict: logger.error(f"Failed to wake up engine: {e}") return {"status": "error", "message": str(e)} + # ── RL weight lifecycle engine routes ────────────────────────────── + # Signatures kept compatible with SGLang's merged #6094 routes so + # a single admin coordinator can talk to either backend. + + async def pause_generation(self, body: dict) -> dict: + """Pause the engine: drain in-flight requests, keep model loaded. + + Called by RL admin coordinator before weight updates. + Uses engine_client.pause_generation() directly -- does NOT sleep + (no GPU memory release) and does NOT unregister from discovery. + """ + body = body or {} + try: + await self.engine_client.pause_generation() + logger.info("[RL] Engine paused (generation quiesced)") + return {"status": "ok", "message": "Engine paused"} + except Exception as e: + logger.error(f"[RL] Failed to pause: {e}") + return {"status": "error", "message": str(e)} + + async def resume_generation(self, body: dict) -> dict: + """Resume the engine after a weight update.""" + body = body or {} + try: + await self.engine_client.resume_generation() + logger.info("[RL] Engine resumed") + return {"status": "ok", "message": "Engine resumed"} + except Exception as e: + logger.error(f"[RL] Failed to resume: {e}") + return {"status": "error", "message": str(e)} + + async def flush_cache(self, body: dict) -> dict: + """Invalidate prefix/KV cache. Called after weight updates.""" + body = body or {} + try: + await self.engine_client.reset_prefix_cache() + logger.info("[RL] Prefix cache flushed") + return {"status": "ok", "message": "Cache flushed"} + except Exception as e: + logger.error(f"[RL] Failed to flush cache: {e}") + return {"status": "error", "message": str(e)} + + async def update_weights_from_path(self, body: dict) -> dict: + """Load weights from a filesystem path (safetensors/torch checkpoint). + + Expects body: {"path": "/path/to/weights", "version": "step_N"} + The caller is responsible for pausing/resuming around this call. + """ + body = body or {} + path = body.get("path") + version = body.get("version", "unknown") + if not path: + return {"status": "error", "message": "Missing 'path' in body"} + try: + # Use vLLM's built-in reload_weights via collective RPC. + # This calls Worker.reload_weights() -> GPUModelRunner.reload_weights() + # which handles loading safetensors from a directory using vLLM's + # model loader with proper layerwise reload. + await self.engine_client.collective_rpc( + "reload_weights", + kwargs={"weights_path": path}, + ) + self._weight_version = version + logger.info(f"[RL] Weights loaded from {path} (version={version})") + return { + "status": "ok", + "message": f"Weights loaded from {path}", + "version": version, + } + except Exception as e: + logger.error(f"[RL] Failed to load weights from {path}: {e}") + return {"status": "error", "message": str(e)} + + async def get_weight_version(self, body: dict) -> dict: + """Return the current weight version tag.""" + return {"version": getattr(self, "_weight_version", "initial")} + + async def load_lora_adapter(self, body: dict) -> dict: + """Load (or hot-swap) a LoRA adapter from a filesystem path. + + Expects body: {"lora_name": str, "lora_path": "/path/to/adapter_dir"} + + The adapter directory must contain ``adapter_model.safetensors`` and + ``adapter_config.json`` -- the standard PEFT output layout that Prime-RL + writes each training step. + + Unlike :meth:`load_lora` (which downloads from a URI via ``LoRAManager`` + streaming a gRPC response), this method is the RL admin equivalent used + for training-loop weight updates: + + * Reads the adapter directly from the given filesystem path (no URI / + no network fetch, no LoRAManager needed). + * Hot-swaps if ``lora_name`` is already loaded (remove old id then + re-add) so every training step replaces the same logical adapter. + * Resets the prefix cache after a hot-swap so stale KV entries keyed + to the previous adapter weights do not poison subsequent rollouts. + * Publishes a ModelDeploymentCard the first time a new ``lora_name`` is + loaded. Prime-RL switches its request ``model`` field to the LoRA + name after load (``scheduler.py``: ``self.model_name = self.lora_name``) + so the frontend needs an MDC entry to route ``r16-a32`` → this worker. + On subsequent hot-swaps the MDC is already published and we skip + re-registration. + """ + body = body or {} + lora_name = body.get("lora_name") + lora_path = body.get("lora_path") + if not lora_name: + return {"status": "error", "message": "Missing 'lora_name' in body"} + if not lora_path: + return {"status": "error", "message": "Missing 'lora_path' in body"} + try: + lock = self._get_lora_lock(lora_name) + async with lock: + lora_id = lora_name_to_id(lora_name) + is_hot_swap = lora_name in self.loaded_loras + + # Hot-swap: vLLM's add_lora is a no-op when the lora_int_id is + # already registered, so we must remove the previous adapter + # first. remove_lora is best-effort on a fresh add. + if is_hot_swap: + old_id = self.loaded_loras[lora_name].id + try: + await self.engine_client.remove_lora(old_id) + # Invalidate the cache entry immediately after remove succeeds. + # If add_lora below fails, this prevents a stale entry pointing + # at an adapter the engine no longer holds from poisoning future + # rollouts with wrong importance ratios (Tier-1 RL correctness risk). + self.loaded_loras.pop(lora_name, None) + except Exception as e: + logger.warning( + f"[RL] remove_lora({lora_name}, id={old_id}) failed during hot-swap: {e}" + ) + + await self.engine_client.add_lora( + LoRARequest( + lora_name=lora_name, + lora_int_id=lora_id, + lora_path=lora_path, + ) + ) + self.loaded_loras[lora_name] = LoRAInfo(id=lora_id, path=lora_path) + + # Invalidate KV cache on hot-swap so stale prefix entries keyed + # to the previous LoRA weights can't contaminate new rollouts. + if is_hot_swap: + try: + await self.engine_client.reset_prefix_cache() + except Exception as e: + # ERROR not WARNING: a failed cache reset means subsequent requests + # sharing a prefix with an old rollout can reuse KV state computed + # under the previous adapter — causing silent logprobs mismatch. + logger.error( + f"[RL] reset_prefix_cache after LoRA swap failed — KV cache may " + f"be contaminated with stale entries from the old adapter. " + f"Rollouts on this worker are unreliable until the next successful " + f"swap: {e}" + ) + + # Publish an MDC for the LoRA on first load so Dynamo's frontend + # can route requests with model= to this worker. + # Mirror the logic in load_lora() (URI variant). Skip on hot-swap + # since the MDC was already published on the first load. + if not is_hot_swap and self.generate_endpoint is not None: + try: + runtime_config = ModelRuntimeConfig() + runtime_config.tool_call_parser = self.config.dyn_tool_call_parser + runtime_config.reasoning_parser = self.config.dyn_reasoning_parser + await register_model( + model_input=ModelInput.Tokens, + model_type=ModelType.Chat | ModelType.Completions, + endpoint=self.generate_endpoint, + model_path=self.config.model, + kv_cache_block_size=self.config.engine_args.block_size, + runtime_config=runtime_config, + user_data={"lora_adapter": True, "lora_id": lora_id}, + lora_name=lora_name, + base_model_path=self.config.model, + ) + logger.info( + f"[RL] Published LoRA '{lora_name}' ModelDeploymentCard" + ) + except Exception as e: + # Rollback: remove the LoRA from the engine to keep state consistent. + logger.exception( + f"[RL] Failed to publish LoRA '{lora_name}' MDC: {e}; rolling back add_lora" + ) + try: + await self.engine_client.remove_lora(lora_id) + except Exception as rollback_err: + # The adapter is now leaked in the engine: it is registered but + # unreachable via loaded_loras (we pop it below). Log at ERROR + # so this doesn't go unnoticed in production. + logger.error( + f"[RL] Rollback remove_lora({lora_name}, id={lora_id}) failed " + f"— adapter is leaked in the engine: {rollback_err}" + ) + self.loaded_loras.pop(lora_name, None) + return { + "status": "error", + "message": f"Failed to register LoRA '{lora_name}' in discovery registry: {e}", + "lora_name": lora_name, + } + + logger.info( + f"[RL] LoRA adapter {'hot-swapped' if is_hot_swap else 'loaded'}: " + f"name={lora_name} id={lora_id} path={lora_path}" + ) + return { + "status": "ok", + "message": f"LoRA adapter '{lora_name}' loaded from {lora_path}", + "lora_name": lora_name, + "lora_id": lora_id, + "hot_swap": is_hot_swap, + } + except Exception as e: + logger.exception( + f"[RL] Failed to load LoRA adapter '{lora_name}' from {lora_path}: {e}" + ) + return {"status": "error", "message": str(e)} + + async def unload_lora_adapter(self, body: dict) -> dict: + """Unload a LoRA adapter previously loaded via :meth:`load_lora_adapter`. + + Expects body: {"lora_name": str} + + Idempotent: unloading an already-absent LoRA returns status=ok so + callers can safely retry without special-casing the not-found path. + """ + body = body or {} + lora_name = body.get("lora_name") + if not lora_name: + return {"status": "error", "message": "Missing 'lora_name' in body"} + try: + lock = self._get_lora_lock(lora_name) + async with lock: + lora = self.loaded_loras.get(lora_name) + if lora is None: + return { + "status": "ok", + "message": f"LoRA adapter '{lora_name}' not loaded (no-op)", + "lora_name": lora_name, + } + lora_id = lora.id + await self.engine_client.remove_lora(lora_id) + del self.loaded_loras[lora_name] + + # Unregister the MDC published on load so the frontend stops + # routing `model=` requests to this worker. + if self.generate_endpoint is not None: + try: + await unregister_model( + endpoint=self.generate_endpoint, + lora_name=lora_name, + ) + except Exception as e: + logger.warning( + f"[RL] Failed to unregister LoRA '{lora_name}' MDC (adapter already removed from engine): {e}" + ) + + logger.info( + f"[RL] LoRA adapter unloaded: name={lora_name} id={lora_id}" + ) + return { + "status": "ok", + "message": f"LoRA adapter '{lora_name}' unloaded", + "lora_name": lora_name, + "lora_id": lora_id, + } + except Exception as e: + logger.exception( + f"[RL] Failed to unload LoRA adapter '{lora_name}': {e}" + ) + return {"status": "error", "message": str(e)} + @abstractmethod def generate(self, request: RequestT, context: Context) -> AsyncIterator[ResponseT]: raise NotImplementedError diff --git a/components/src/dynamo/vllm/publisher.py b/components/src/dynamo/vllm/publisher.py index 315904b78a8d..f25d227332cb 100644 --- a/components/src/dynamo/vllm/publisher.py +++ b/components/src/dynamo/vllm/publisher.py @@ -57,6 +57,10 @@ def record( *args: object, **kwargs: object, ) -> None: + # scheduler_stats can be None right after a weight reload / cache reset. + if scheduler_stats is None: + return + active_decode_blocks = int(self.num_gpu_block * scheduler_stats.kv_cache_usage) self.inner.publish(self.dp_rank, kv_used_blocks=active_decode_blocks) diff --git a/components/src/dynamo/vllm/worker_factory.py b/components/src/dynamo/vllm/worker_factory.py index f20b19c78440..6c100c510728 100644 --- a/components/src/dynamo/vllm/worker_factory.py +++ b/components/src/dynamo/vllm/worker_factory.py @@ -335,8 +335,21 @@ async def _create_decode_worker( runtime.register_engine_route("sleep", handler.sleep) runtime.register_engine_route("wake_up", handler.wake_up) runtime.register_engine_route("scale_elastic_ep", handler.scale_elastic_ep) + # RL weight lifecycle routes (parity with SGLang #6094) + runtime.register_engine_route("pause_generation", handler.pause_generation) + runtime.register_engine_route("resume_generation", handler.resume_generation) + runtime.register_engine_route("flush_cache", handler.flush_cache) + runtime.register_engine_route("update_weights_from_path", handler.update_weights_from_path) + runtime.register_engine_route("get_weight_version", handler.get_weight_version) + # RL LoRA adapter routes: filesystem-native hot-swap used by Prime-RL + # every training step to broadcast new adapter weights into the engine. + runtime.register_engine_route("load_lora_adapter", handler.load_lora_adapter) + runtime.register_engine_route("unload_lora_adapter", handler.unload_lora_adapter) logger.info( - "Registered engine routes: /engine/sleep, /engine/wake_up, /engine/scale_elastic_ep" + "Registered engine routes: sleep, wake_up, scale_elastic_ep, " + "pause_generation, resume_generation, flush_cache, " + "update_weights_from_path, get_weight_version, " + "load_lora_adapter, unload_lora_adapter" ) # Parse endpoint types from --endpoint-types flag @@ -570,8 +583,21 @@ async def _create_prefill_worker( runtime.register_engine_route("sleep", handler.sleep) runtime.register_engine_route("wake_up", handler.wake_up) runtime.register_engine_route("scale_elastic_ep", handler.scale_elastic_ep) + # RL weight lifecycle routes (parity with SGLang #6094) + runtime.register_engine_route("pause_generation", handler.pause_generation) + runtime.register_engine_route("resume_generation", handler.resume_generation) + runtime.register_engine_route("flush_cache", handler.flush_cache) + runtime.register_engine_route("update_weights_from_path", handler.update_weights_from_path) + runtime.register_engine_route("get_weight_version", handler.get_weight_version) + # RL LoRA adapter routes: filesystem-native hot-swap used by Prime-RL + # every training step to broadcast new adapter weights into the engine. + runtime.register_engine_route("load_lora_adapter", handler.load_lora_adapter) + runtime.register_engine_route("unload_lora_adapter", handler.unload_lora_adapter) logger.info( - "Registered engine routes: /engine/sleep, /engine/wake_up, /engine/scale_elastic_ep" + "Registered engine routes: sleep, wake_up, scale_elastic_ep, " + "pause_generation, resume_generation, flush_cache, " + "update_weights_from_path, get_weight_version, " + "load_lora_adapter, unload_lora_adapter" ) # Wait for self-benchmark to complete before registering. diff --git a/container/deps/vllm/install_vllm.sh b/container/deps/vllm/install_vllm.sh index c63f274b1d84..93c3fbb9d60b 100755 --- a/container/deps/vllm/install_vllm.sh +++ b/container/deps/vllm/install_vllm.sh @@ -12,7 +12,7 @@ set -euo pipefail -VLLM_VER="0.19.0" +VLLM_VER="0.19.1" VLLM_REF="v${VLLM_VER}" DEVICE="cuda" @@ -275,4 +275,35 @@ if [ "$DEVICE" = "cuda" ]; then # TODO we will be able to specify which pplx and deepep commit we want in future TORCH_CUDA_ARCH_LIST="$TORCH_CUDA_ARCH_LIST" bash install_python_libraries.sh fi + +# --------------------------------------------------------------------------- +# prime-rl inference-side vLLM plugin (pinned tag). +# +# Registers the ``vllm.general_plugins`` entry-point that applies prime-rl's +# monkey patches (LoRA adapter load, DP engine pause/resume deadlock, Qwen 3.5 +# LoRA, etc.) automatically in every vLLM worker process -- including spawned +# subprocesses. Required for prime-rl / Dynamo RL training integration. +# +# Override at build time: --build-arg PRIME_RL_REF=v0.5.1.dev101 +# --no-deps: prime-rl's full dep tree includes trainer + wandb; Dynamo only +# needs the inference-side plugin and worker-extension classes. +# Python version: prime-rl pins requires-python = "~=3.12.0"; Dynamo containers +# are Python 3.12, so no version override is needed. For 3.11 local +# dev venvs use the regular pip (not uv) with --ignore-requires-python. +# --------------------------------------------------------------------------- +PRIME_RL_REF="${PRIME_RL_REF:-v0.5.1.dev101}" +echo "\n=== Installing prime-rl vLLM plugin (ref=${PRIME_RL_REF}) ===" +uv pip install --no-deps \ + "prime-rl @ git+https://github.com/PrimeIntellect-ai/prime-rl@${PRIME_RL_REF}" + +# Sanity-check: confirm vllm.general_plugins entry-point is registered. +python3 - <<'PY_SANITY' +from importlib.metadata import entry_points +names = [ep.name for ep in entry_points(group="vllm.general_plugins")] +assert "prime_rl" in names, ( + f"prime-rl plugin NOT registered; vllm.general_plugins={names}" +) +print(f"✓ prime-rl plugin registered (vllm.general_plugins={names})") +PY_SANITY + echo "\n✅ All installations completed successfully!" diff --git a/docs/Dynamo-RL-api-draft.md b/docs/Dynamo-RL-api-draft.md new file mode 100644 index 000000000000..ae8113ed3073 --- /dev/null +++ b/docs/Dynamo-RL-api-draft.md @@ -0,0 +1,973 @@ +# Dynamo RL API Draft + +**Branch:** `bis/parity-tokenize-tcp` (HEAD: `d837fbd67`) + +Commit `70f84570b` (the current auto-enable-token-ids commit on HEAD) is an +equivalent cherry-pick of the earlier `19d1bf13d` referenced in prior drafts — +same subject, same patch semantics, different parent tree after rebase onto +`origin/main`. + +--- + +## Table of Contents + +1. [Overview](#1-overview) +2. [Architecture](#2-architecture) +3. [Configuration](#3-configuration) +4. [API Reference](#4-api-reference) + - 4.1 Chat Completions (RL-enhanced) + - 4.2 Token-In / Token-Out (TITO) + - 4.3 Tokenization + - 4.4 Fleet Control (`/v1/rl/*`) +5. [Data Flow](#5-data-flow) +6. [Key Data Structures](#6-key-data-structures) +7. [Worker Engine Routes (Internal)](#7-worker-engine-routes-internal) +8. [Known Limitations](#8-known-limitations) +9. [Validation Results](#9-validation-results) + +--- + +## 1. Overview + +This document describes the RL training API surface on the Dynamo serving stack for integration with prime-rl. The Dynamo frontend (Rust) exposes: + +- An `/v1/rl/*` router for the full RL control-plane lifecycle (pause/resume, weight updates, readiness checks) +- Automatic token-level data injection (`prompt_token_ids`, `completion_token_ids`) in chat completion responses +- `/v1/tokenize` and `/v1/detokenize` endpoints +- A `/v1/chat/completions/tokens` TITO endpoint for pre-tokenized prompt bypass + +Zero Python in the inference or admin data path. The Rust frontend handles all HTTP API surface while vLLM workers expose engine routes for weight lifecycle operations on the GPU. + +### Endpoint Summary + +| Capability | Endpoint | Purpose | +|------------|----------|---------| +| Inference | `POST /v1/chat/completions` | Generate rollouts; responses include `prompt_token_ids` + `choice.token_ids` | +| TITO inference | `POST /v1/chat/completions/tokens` | Pre-tokenized prompt bypass (turn 2+ in multi-turn RL) | +| Tokenization | `POST /v1/tokenize` | Consistent tokenization using the model's chat template | +| Detokenization | `POST /v1/detokenize` | Token IDs back to text | +| Pause fleet | `POST /v1/rl/pause` | Drain in-flight requests before weight update | +| Resume fleet | `POST /v1/rl/resume` | Resume generation after weight update | +| Update weights | `POST /v1/rl/update_weights` | Atomic flush + reload from checkpoint directory | +| Load LoRA adapter | `POST /v1/rl/load_lora_adapter` | Hot-load/swap a PEFT-style adapter from filesystem path | +| Unload LoRA adapter | `POST /v1/rl/unload_lora_adapter` | Remove a previously loaded adapter by name | +| Weight version | `GET /v1/rl/weight_version` | Query current weight version across workers | +| Health | `GET /v1/rl/health` | Lightweight frontend health check | +| Readiness | `GET /v1/rl/ready` | Deep check: are workers reachable and healthy? | + +### What Changed vs. Stock Dynamo + +All changes are on `bis/parity-tokenize-tcp` (18 commits, 26 files, +3030/-41 — the diff counts include this doc). Nothing touches Dynamo's core serving pipeline (NATS, scheduler, KV cache, disaggregation). The changes are additive: + +- **Rust frontend** (`lib/llm/`): New routes, response post-processing, tokenization endpoints, LoRA hot-swap admin routes +- **vLLM worker** (`components/`): 7 engine route handlers (5 weight-lifecycle + 2 LoRA), publisher crash guard +- **Deps** (`container/`): default `VLLM_VER` bumped 0.19.0 → 0.19.1; prime-rl plugin installed via `pip` so `vllm.general_plugins` patches apply at engine start +- **Compat fixes**: `/v1/tokenize` and `/v1/detokenize` adapted to upstream `DecodeResult`-returning decoder (main commit `2cabf4414`, #8022) + +--- + +## 2. Architecture + +### Component Topology + +```mermaid +flowchart TD + subgraph prime_rl["prime-rl"] + orch["Orchestrator
(prime_rl.orchestrator)"] + trainer["Trainer
(prime_rl.trainer.rl.train)
torchrun --nproc-per-node=N"] + end + + subgraph dynamo["Dynamo Serving Stack"] + subgraph frontend["Frontend Pod (Rust, port 8000)"] + cc["/v1/chat/completions
+ prompt_token_ids
+ choice.token_ids"] + tito["/v1/chat/completions/tokens
(TITO)"] + tok["/v1/tokenize   /v1/detokenize"] + rl["/v1/rl/*
health, ready, pause, resume,
update_weights, weight_version"] + end + subgraph worker["vLLM Worker Pod (Python, system port 9090)"] + eng["/engine/*
pause_generation
resume_generation
flush_cache
update_weights_from_path
get_weight_version"] + gpu["GPU
Model Weights"] + end + end + + subgraph storage["Shared Storage (PVC)"] + pvc["prime-rl-shared-data
safetensors checkpoints"] + end + + orch -- "rollouts
POST /v1/chat/completions" --> cc + orch -- "TITO turn 2+
POST /v1/chat/completions/tokens" --> tito + orch -- "POST /v1/tokenize" --> tok + orch -- "weight lifecycle
pause / update_weights / resume" --> rl + rl -- "HTTP fan-out
(concurrent to all workers)" --> eng + eng --> gpu + trainer -- "write checkpoint" --> pvc + eng -- "reload_weights
(collective_rpc)" --> pvc +``` + +### Key Design Decisions + +1. **Single entry point.** Prime-RL points both `base_url` and `admin_base_url` at the Dynamo frontend. No separate admin service to deploy. + +2. **Fan-out in Rust.** The `/v1/rl/*` handlers fan out to all vLLM workers via `DYN_RL_WORKER_SYSTEM_URLS`. This supports DP>1 without Prime-RL needing to discover workers. The frontend returns HTTP 200 only when all workers respond OK, and HTTP 502 otherwise with per-worker details. + +3. **Token IDs as a response extension.** When `DYN_ENABLE_RL=true`, `prompt_token_ids` and `choices[i].token_ids` are injected into every non-streaming response automatically. No client-side configuration needed. + +4. **Backward compatible.** All new response fields use `#[serde(skip_serializing_if = "Option::is_none")]`. Clients that don't set `DYN_ENABLE_RL` see standard OpenAI-compatible responses with no extra fields. + +--- + +## 3. Configuration + +### Environment Variables (Frontend) + +| Variable | Default | Description | +|----------|---------|-------------| +| `DYN_ENABLE_RL` | `false` | Master switch. Mounts `/v1/rl/*` routes, auto-injects token IDs in chat completion responses, mounts TITO endpoint. | +| `DYN_RL_WORKER_SYSTEM_URLS` | `http://localhost:8081` | Comma-separated list of vLLM worker system HTTP base URLs for fan-out. | + +### Environment Variables (Worker) + +| Variable | Default | Description | +|----------|---------|-------------| +| `DYN_SYSTEM_PORT` | `8081` (local) / `9090` (k8s) | Worker's system HTTP port where engine routes are registered. | + +### Prime-RL Configuration (`orch.toml`) + +```toml +max_steps = 20 +seq_len = 512 +batch_size = 16 +rollouts_per_example = 4 +use_token_client = false + +[model] +name = "PrimeIntellect/Qwen3-0.6B-Reverse-Text-SFT" + +[sampling] +max_tokens = 64 + +[[env]] +id = "reverse-text" + +[client] +# Point BOTH base_url and admin_base_url at the Dynamo frontend. +# admin_base_url uses /v1/rl because Prime-RL strips trailing /v1 +# from admin URLs, but /v1/rl is preserved. +base_url = ["http://:8000/v1"] +admin_base_url = ["http://:8000/v1/rl"] +skip_model_check = true + +[weight_broadcast] +type = "filesystem" + +[experimental] +# Disable prefix cache salt until Dynamo supports it. +# verifiers dev6+ defaults use_prefix_cache_salt=True; current image returns 400. +use_prefix_cache_salt = false +``` + +**Important:** Do NOT set `send_return_token_ids = true` in `[sampling]`. The Rust frontend handles token ID injection automatically when `DYN_ENABLE_RL=true`. Sending `return_token_ids=true` in the request causes the OpenAI SDK to parse the response and strip unknown fields. + +### Kubernetes (DGD) + +```yaml +# Frontend pod env +- name: DYN_ENABLE_RL + value: "true" +- name: DYN_RL_WORKER_SYSTEM_URLS + value: "http://prime-rl-dynamo-vllmworker..svc.cluster.local:9090" +``` + +### Launch Commands (Local) + +```bash +# Frontend with RL routes enabled +DYN_ENABLE_RL=true \ +DYN_RL_WORKER_SYSTEM_URLS=http://localhost:8081 \ + python -m dynamo.frontend + +# vLLM Worker +CUDA_VISIBLE_DEVICES=0 \ +DYN_SYSTEM_PORT=8081 \ + python -m dynamo.vllm \ + --model PrimeIntellect/Qwen3-0.6B-Reverse-Text-SFT \ + --served-model-name PrimeIntellect/Qwen3-0.6B-Reverse-Text-SFT \ + --enforce-eager \ + --max-model-len 2048 \ + --gpu-memory-utilization 0.5 +``` + +--- + +## 4. API Reference + +All endpoints live on the Dynamo Rust frontend (default port 8000). Unless noted, request/response formats follow the OpenAI API specification. + +### 4.1 Chat Completions (RL-enhanced) + +``` +POST /v1/chat/completions +``` + +Standard OpenAI chat completions with RL extensions. When `DYN_ENABLE_RL=true`, every non-streaming response is automatically enriched with token IDs for the trainer. + +#### Request + +Standard OpenAI `ChatCompletionRequest`. Two additional fields are accepted and silently consumed (never forwarded to the vLLM worker): + +| Field | Type | Default | Mapped to | +|-------|------|---------|-----------| +| `tokens` | `u32[]` | `null` | `nvext.token_data` (tokenizer bypass) | +| `return_token_ids` | `bool` | `null` | `nvext.extra_fields: ["token_ids", "completion_token_ids"]` + `logprobs: true` | + +When `DYN_ENABLE_RL=true`, `return_token_ids` is implicitly `true` for every request. + +#### Sample Request + +```bash +curl -s -X POST http://localhost:8000/v1/chat/completions \ + -H 'Content-Type: application/json' \ + -d '{ + "model": "PrimeIntellect/Qwen3-0.6B-Reverse-Text-SFT", + "messages": [ + {"role": "user", "content": "Reverse this: hello world"} + ], + "max_tokens": 64, + "temperature": 1.0 + }' +``` + +#### Sample Response (Non-Streaming, with `DYN_ENABLE_RL=true`) + +```json +{ + "id": "chatcmpl-abc123", + "object": "chat.completion", + "model": "PrimeIntellect/Qwen3-0.6B-Reverse-Text-SFT", + "choices": [{ + "index": 0, + "message": {"role": "assistant", "content": "dlrow olleh"}, + "finish_reason": "stop", + "logprobs": { + "content": [ + {"token": "dl", "logprob": -0.523, "top_logprobs": []}, + {"token": "row", "logprob": -0.102, "top_logprobs": []}, + {"token": " ol", "logprob": -0.834, "top_logprobs": []}, + {"token": "leh", "logprob": -0.211, "top_logprobs": []} + ] + }, + "token_ids": [67, 1245, 893, 15] + }], + "prompt_token_ids": [151644, 8948, 198, 151645, 198, 151644, 872, 198, + 49, 1075, 513, 420, 25, 24748, 1879, 198, 151645, + 198, 151644, 77091, 198], + "usage": {"prompt_tokens": 21, "completion_tokens": 4, "total_tokens": 25}, + "nvext": { + "completion_token_ids": [67, 1245, 893, 15] + } +} +``` + +#### Response Field Reference + +| Field | JSON path | Description | +|-------|-----------|-------------| +| `prompt_token_ids` | `response.prompt_token_ids` | Token IDs from tokenizing the prompt messages through the model's chat template. Generated by the Rust frontend's tokenizer after the response is fully received. | +| `token_ids` | `response.choices[i].token_ids` | Completion token IDs generated by the engine. Promoted from `nvext.completion_token_ids`. | +| `completion_token_ids` | `response.nvext.completion_token_ids` | Canonical Dynamo location for output token IDs. Accumulated across all SSE chunks by `DeltaGenerator`. | + +**Why `token_ids` appears in two locations:** Prime-RL's verifiers library reads `response.prompt_token_ids` and `choices[i].token_ids` (top-level on the choice object). Dynamo natively emits output token IDs in `nvext.completion_token_ids`. The Rust post-processor promotes the latter to the former for compatibility. Both contain the same values. + +**Invariant:** `len(completion_token_ids) == len(logprobs.content)` -- the output token IDs are in exact 1:1 correspondence with the logprob entries. + +#### Sample Response (Streaming / SSE -- final chunk only) + +Intermediate chunks carry `delta.content` only. Token IDs appear exclusively on the **final chunk** (the one with a non-null `finish_reason`): + +``` +data: {"id":"chatcmpl-abc123","object":"chat.completion.chunk", + "choices":[{"index":0,"delta":{},"finish_reason":"stop", + "nvext":{"completion_token_ids":[67,1245,893,15]}}], + "prompt_token_ids":[151644,8948,198,151645,198,151644,872,198, + 49,1075,513,420,25,24748,1879,198,151645, + 198,151644,77091,198]} + +data: [DONE] +``` + +#### RL Post-Processing Pipeline + +For non-streaming requests, the handler performs the following after the backend response is fully aggregated: + +```mermaid +flowchart LR + A["Request arrives
DYN_ENABLE_RL=true"] --> B["Save messages
for later tokenization"] + B --> C["Inject nvext.extra_fields:
[token_ids, completion_token_ids]
Force logprobs=true"] + C --> D["Standard pipeline
(preprocessor, backend,
delta generator, aggregator)"] + D --> E["Aggregate response
(nvext.completion_token_ids
accumulated in delta.rs)"] + E --> F["rl_tokenize_prompt()
messages -> prompt_token_ids
via model chat template"] + F --> G["rl_promote_token_ids()
nvext.completion_token_ids
-> choices[i].token_ids"] + G --> H["Return enriched
JSON response"] +``` + +--- + +### 4.2 Token-In / Token-Out (TITO) + +``` +POST /v1/chat/completions/tokens +``` + +Dedicated endpoint for Prime-RL's pre-tokenized prompt flow (multi-turn RL, turn 2+). The orchestrator sends raw token IDs instead of text messages, bypassing the frontend's tokenizer entirely. This avoids redundant encode/decode round-trips and ensures token-level alignment. + +#### Sample Request + +```bash +curl -s -X POST http://localhost:8000/v1/chat/completions/tokens \ + -H 'Content-Type: application/json' \ + -d '{ + "model": "PrimeIntellect/Qwen3-0.6B-Reverse-Text-SFT", + "messages": [{"role": "user", "content": "(token-in mode)"}], + "tokens": [151644, 8948, 198, 151645, 198, 151644, 872, 198, + 49, 1075, 513, 420, 25, 24748, 1879, 198, 151645, + 198, 151644, 77091, 198, 67, 1245, 893, 15], + "max_tokens": 64 + }' +``` + +| Field | Required | Description | +|-------|----------|-------------| +| `tokens` | **Yes** | Pre-tokenized prompt IDs. Must be non-empty. Injected as `nvext.token_data`. | +| `messages` | Yes (can be placeholder) | A placeholder `{"role": "user", "content": "(token-in mode)"}` is auto-injected if empty. | + +#### Behavior + +1. Extracts `tokens`, returns 400 if missing or empty +2. Injects into `nvext.token_data` (triggers tokenizer bypass in `preprocessor.rs`) +3. Adds `extra_fields: ["token_ids", "completion_token_ids"]` +4. Forces `logprobs = true` +5. Delegates to the standard `chat_completions()` pipeline (zero HTTP proxy) + +#### Sample Response + +Same shape as section 4.1. The response includes `prompt_token_ids` and `choices[i].token_ids`. + +--- + +### 4.3 Tokenization + +``` +POST /v1/tokenize +POST /v1/detokenize +``` + +Consistent tokenization using the model's tokenizer and chat template, running entirely in Rust. These are critical for RL: prompt token IDs in the chat completion response must match what the tokenizer produces for the same messages. Both endpoints use the same tokenizer instance that the frontend uses for its own request preprocessing. + +#### Sample: Tokenize (Chat variant) + +```bash +curl -s -X POST http://localhost:8000/v1/tokenize \ + -H 'Content-Type: application/json' \ + -d '{ + "model": "PrimeIntellect/Qwen3-0.6B-Reverse-Text-SFT", + "messages": [ + {"role": "user", "content": "Reverse this: hello world"} + ], + "add_generation_prompt": true, + "add_special_tokens": true + }' +``` + +**Response:** + +```json +{ + "count": 21, + "max_model_len": 2048, + "tokens": [151644, 8948, 198, 151645, 198, 151644, 872, 198, + 49, 1075, 513, 420, 25, 24748, 1879, 198, 151645, + 198, 151644, 77091, 198], + "token_strs": null +} +``` + +#### Tokenize Request Fields (Chat variant) + +| Field | Type | Default | Description | +|-------|------|---------|-------------| +| `model` | `string?` | auto-resolve | Model name for tokenizer lookup | +| `messages` | `ChatMessage[]` | -- | Messages to tokenize through the model's chat template | +| `add_generation_prompt` | `bool` | `true` | Append generation prompt (e.g., `<\|im_start\|>assistant\n`) | +| `add_special_tokens` | `bool` | `true` | Add BOS/EOS tokens | +| `return_token_strs` | `bool` | `false` | Include human-readable string representation of each token | +| `chat_template` | `string?` | `null` | Override the model's default chat template (Jinja2) | +| `chat_template_kwargs` | `object?` | `null` | Extra template variables | +| `continue_final_message` | `bool` | `false` | Continue last message instead of starting a new turn | + +The chat variant renders messages through the model's chat template before tokenizing, so the token count is the exact number of tokens that a corresponding chat completion request would consume. + +#### Sample: Tokenize (Completion variant) + +```bash +curl -s -X POST http://localhost:8000/v1/tokenize \ + -H 'Content-Type: application/json' \ + -d '{ + "model": "PrimeIntellect/Qwen3-0.6B-Reverse-Text-SFT", + "prompt": "hello world", + "add_special_tokens": true, + "return_token_strs": true + }' +``` + +**Response:** + +```json +{ + "count": 3, + "max_model_len": 2048, + "tokens": [9707, 1879, 3], + "token_strs": ["hello", " world", ""] +} +``` + +#### Tokenize Response Fields + +| Field | Type | Description | +|-------|------|-------------| +| `count` | `int` | Number of tokens | +| `max_model_len` | `int` | Model's configured maximum context length | +| `tokens` | `list[int]` | Token ID list | +| `token_strs` | `list[str]?` | Human-readable token strings; only present if `return_token_strs: true` | + +#### Sample: Detokenize + +```bash +curl -s -X POST http://localhost:8000/v1/detokenize \ + -H 'Content-Type: application/json' \ + -d '{ + "model": "PrimeIntellect/Qwen3-0.6B-Reverse-Text-SFT", + "tokens": [9707, 1879, 3] + }' +``` + +**Response:** + +```json +{"prompt": "hello world"} +``` + +--- + +### 4.4 Fleet Control (`/v1/rl/*`) + +All `/v1/rl/*` routes are mounted only when `DYN_ENABLE_RL=true`. They fan out to vLLM worker system ports defined by `DYN_RL_WORKER_SYSTEM_URLS`. + +In prime-rl's config: + +```toml +[client] +base_url = ["http://:8000/v1"] +admin_base_url = ["http://:8000/v1/rl"] +``` + +--- + +#### `GET /v1/rl/health` + +Lightweight liveness probe. Returns immediately as long as the frontend process is running. Used by prime-rl's `check_health()` on the admin client. + +```bash +curl -s http://localhost:8000/v1/rl/health +``` + +```json +{"status": "ok"} +``` + +--- + +#### `GET /v1/rl/ready` + +Composite readiness probe. Polls `/health` on every configured worker system URL concurrently. Returns 200 only when all workers respond with HTTP 2xx. + +```bash +curl -s http://localhost:8000/v1/rl/ready +``` + +```json +// All workers ready (200) +{ + "status": "ready", + "workers": [ + {"url": "http://localhost:8081", "healthy": true} + ] +} + +// Not all workers ready (503) +{ + "status": "not_ready", + "workers_ready": 0, + "workers_total": 1, + "workers": [ + {"url": "http://localhost:8081", "healthy": false, "error": "connection refused"} + ] +} +``` + +--- + +#### `POST /v1/rl/pause` + +Quiesces generation on all workers. Each worker calls `engine_client.pause_generation()` which drains in-flight requests without unloading the model from GPU memory. + +```bash +curl -s -X POST http://localhost:8000/v1/rl/pause -H 'Content-Type: application/json' -d '{}' +``` + +```json +// Success (200) +{ + "status": "ok", + "workers": [ + {"status": "ok", "message": "Engine paused"} + ] +} + +// Failure (502) +{ + "status": "error", + "workers": [ + {"status": "ok", "message": "Engine paused"}, + {"status": "error", "message": "timeout"} + ] +} +``` + +--- + +#### `POST /v1/rl/resume` + +Resumes generation on all workers after a weight update. + +```bash +curl -s -X POST http://localhost:8000/v1/rl/resume -H 'Content-Type: application/json' -d '{}' +``` + +```json +{ + "status": "ok", + "workers": [ + {"status": "ok", "message": "Engine resumed"} + ] +} +``` + +--- + +#### `POST /v1/rl/update_weights` + +Atomic weight-loading sequence: `flush_cache` then `update_weights_from_path` on all workers concurrently. The frontend performs the two-phase fan-out internally; prime-rl does not need to call these separately. + +For filesystem-backed weight broadcast, the trainer writes safetensors files to a shared PVC directory and passes that path here. The worker calls `engine_client.collective_rpc("reload_weights", kwargs={"weights_path": path})` which triggers vLLM's layerwise in-place weight reload on every GPU worker. + +For NCCL-based weight broadcast (`weight_dir: null`), Dynamo returns 200 immediately -- the actual weight transfer happens out-of-band via NCCL and Dynamo does not participate. + +```bash +# Filesystem mode +curl -s -X POST http://localhost:8000/v1/rl/update_weights \ + -H 'Content-Type: application/json' \ + -d '{"weight_dir": "/data/outputs/run_default/broadcasts/step_5"}' + +# NCCL mode (Dynamo no-op) +curl -s -X POST http://localhost:8000/v1/rl/update_weights \ + -H 'Content-Type: application/json' \ + -d '{"weight_dir": null}' +``` + +```json +// Success (200) +{ + "status": "ok", + "version": "step_5", + "workers": [ + {"status": "ok", "message": "Weights loaded from /data/outputs/...", "version": "step_5"} + ] +} + +// Failure at flush_cache stage (502) +{ + "status": "error", + "stage": "flush_cache", + "workers": [...] +} + +// Failure at update_weights stage (502) +{ + "status": "error", + "stage": "update_weights_from_path", + "workers": [...] +} +``` + +The `version` string is derived from the basename of `weight_dir` (e.g., `step_5` from `/data/outputs/run_default/broadcasts/step_5`). This version is stored in the worker and retrievable via `/v1/rl/weight_version`. + +--- + +#### `POST /v1/rl/load_lora_adapter` + +Hot-load or hot-swap a LoRA adapter from a filesystem path. The adapter directory must contain PEFT-style `adapter_model.safetensors` and `adapter_config.json` -- the default output layout of prime-rl's LoRA trainer. + +This is the RL-native LoRA path, distinct from Dynamo's URI-based `load_lora` gRPC endpoint (which downloads from S3/file URIs via `LoRAManager`). The admin route is optimized for the training loop: no URI fetch, no MDC churn on hot-swap. + +- **First call for a given `lora_name`**: `add_lora` in the engine, publish a ModelDeploymentCard so subsequent inference requests with `model=` route to this worker. +- **Subsequent calls (hot-swap)**: `remove_lora(old_id)` → `add_lora` with new weights → `reset_prefix_cache`. The MDC is left in place since it already points at this worker. + +Pair with `/v1/rl/pause` and `/v1/rl/resume` for full drain-swap-resume semantics. + +```bash +curl -s -X POST http://localhost:8000/v1/rl/load_lora_adapter \ + -H 'Content-Type: application/json' \ + -d '{"lora_name": "r16-a32", "lora_path": "/data/outputs/run_default/broadcasts/step_5"}' +``` + +```json +// Success (200) +{ + "status": "ok", + "workers": [ + { + "status": "ok", + "message": "LoRA adapter 'r16-a32' loaded from /data/outputs/...", + "lora_name": "r16-a32", + "lora_id": 788776416, + "hot_swap": false + } + ] +} + +// Missing / empty field (400) +{ + "status": "error", + "message": "Expected body: {\"lora_name\": str, \"lora_path\": str} (both required, non-empty)" +} + +// Worker-side failure (502) -- e.g. bad adapter file, rank mismatch, vLLM not --enable-lora +{ + "status": "error", + "workers": [{"status": "error", "message": "..."}] +} +``` + +**vLLM worker requirements**: the engine must be started with `--enable-lora --max-lora-rank R --max-loras N`, with `R` ≥ the adapter rank and `N` ≥ the number of distinct `lora_name` values you expect to have loaded at once. For Prime-RL's single-adapter training loop, `--max-loras 1` is sufficient. + +--- + +#### `POST /v1/rl/unload_lora_adapter` + +Remove a previously loaded LoRA adapter by name. Idempotent: unloading an already-absent adapter returns `status: ok` so callers can retry safely. + +Unregisters the adapter's ModelDeploymentCard so the frontend stops routing `model=` requests to this worker. + +```bash +curl -s -X POST http://localhost:8000/v1/rl/unload_lora_adapter \ + -H 'Content-Type: application/json' \ + -d '{"lora_name": "r16-a32"}' +``` + +```json +// Success (200) +{ + "status": "ok", + "workers": [ + { + "status": "ok", + "message": "LoRA adapter 'r16-a32' unloaded", + "lora_name": "r16-a32", + "lora_id": 788776416 + } + ] +} + +// Already absent -- still ok (200) +{ + "status": "ok", + "workers": [{"status": "ok", "message": "LoRA adapter 'r16-a32' not loaded (no-op)", "lora_name": "r16-a32"}] +} +``` + +--- + +#### `GET /v1/rl/weight_version` + +Returns the currently loaded weight version from all workers. Useful for debugging weight update races or confirming that all workers converged to the same checkpoint. + +```bash +curl -s http://localhost:8000/v1/rl/weight_version +``` + +```json +// All workers consistent (200) +{ + "status": "ok", + "version": "step_5", + "workers": [ + {"version": "step_5"}, + {"version": "step_5"} + ] +} + +// Workers inconsistent (200, with warning) +{ + "status": "inconsistent", + "versions": ["step_4", "step_5"], + "workers": [ + {"version": "step_4"}, + {"version": "step_5"} + ] +} +``` + +Returns HTTP 200 even when versions are inconsistent -- the `status` field distinguishes the cases. A 502 is only returned for network-level failures. + +--- + +## 5. Data Flow + +### 5.1 Rollout (Inference) Path + +```mermaid +sequenceDiagram + participant Orch as prime-rl Orchestrator + participant FE as Dynamo Frontend (Rust) + participant Worker as vLLM Worker (GPU) + + Orch->>FE: POST /v1/chat/completions
{messages, max_tokens, ...} + Note over FE: DYN_ENABLE_RL=true:
inject nvext.extra_fields
= ["token_ids", "completion_token_ids"]
force logprobs=true
save messages for tokenization + FE->>Worker: forward request (TCP/NATS) + Worker-->>FE: SSE chunks
(delta.content + delta.token_ids per chunk) + Note over FE: DeltaGenerator accumulates
completion_token_ids across chunks + Worker-->>FE: final chunk
(finish_reason + nvext.completion_token_ids) + Note over FE: Post-process:
1. rl_tokenize_prompt(messages)
-> response.prompt_token_ids
2. Promote nvext.completion_token_ids
-> choices[i].token_ids + FE-->>Orch: Enriched response:
prompt_token_ids + choices[i].token_ids
+ nvext.completion_token_ids +``` + +### 5.2 Weight Update Path + +```mermaid +sequenceDiagram + participant Trainer as prime-rl Trainer + participant PVC as Shared Storage + participant Orch as prime-rl Orchestrator + participant FE as Dynamo Frontend (Rust) + participant W1 as vLLM Worker 1 + participant W2 as vLLM Worker 2 + + Trainer->>PVC: write checkpoint
/data/outputs/.../step_N/*.safetensors + Trainer->>Orch: notify weight update ready
(internal IPC) + Orch->>FE: POST /v1/rl/pause + FE->>W1: POST /engine/pause_generation + FE->>W2: POST /engine/pause_generation + W1-->>FE: {status: ok} + W2-->>FE: {status: ok} + FE-->>Orch: {status: ok} + Orch->>FE: POST /v1/rl/update_weights
{weight_dir: /data/outputs/.../step_N} + FE->>W1: POST /engine/flush_cache + FE->>W2: POST /engine/flush_cache + W1-->>FE: {status: ok} + W2-->>FE: {status: ok} + FE->>W1: POST /engine/update_weights_from_path
{path: ..., version: step_N} + FE->>W2: POST /engine/update_weights_from_path
{path: ..., version: step_N} + Note over W1,W2: collective_rpc(reload_weights)
vLLM GPUModelRunner.reload_weights()
in-place layer-by-layer load from safetensors + W1-->>FE: {status: ok, version: step_N} + W2-->>FE: {status: ok, version: step_N} + FE-->>Orch: {status: ok, version: step_N} + Orch->>FE: POST /v1/rl/resume + FE->>W1: POST /engine/resume_generation + FE->>W2: POST /engine/resume_generation + W1-->>FE: {status: ok} + W2-->>FE: {status: ok} + FE-->>Orch: {status: ok} + Note over Orch: Continue rollouts
with updated weights +``` + +### 5.3 TITO (Tokens-In, Tokens-Out) Path + +```mermaid +sequenceDiagram + participant Client as Orchestrator (turn 2+) + participant FE as Dynamo Frontend (Rust) + participant PP as Preprocessor + participant Worker as vLLM Worker + + Client->>FE: POST /v1/chat/completions/tokens
{tokens: [9707, 1879, 3], max_tokens: 64} + Note over FE: Extract tokens field
Inject as nvext.token_data
Force extra_fields, logprobs + FE->>PP: Request with nvext.token_data + Note over PP: token_data present:
skip tokenization,
use provided IDs directly + PP->>Worker: Token IDs sent as-is + Worker-->>FE: Completion response
(with completion_token_ids) + FE-->>Client: Enriched response
(prompt_token_ids + token_ids) +``` + +--- + +## 6. Key Data Structures + +### `NvExtResponse` (Rust -- response side) + +Serialized as the `nvext` field in each SSE chunk or the unary response body: + +``` +NvExtResponse { + worker_id?: WorkerIdInfo -- prefill/decode worker IDs for disaggregated serving + timing?: TimingInfo -- request timing (enabled via extra_fields: ["timing"]) + token_ids?: Vec -- GAIE Stage 1: tokenized prompt for Stage 2 reuse + routed_experts?: serde_json::Value -- SGLang-specific expert routing payload + completion_token_ids?: Vec -- RL: generated output token IDs (final chunk only) +} +``` + +The `completion_token_ids` field is populated automatically for all requests when `DYN_ENABLE_RL=true`, or when the client sends `nvext.extra_fields: ["completion_token_ids"]`. + +### `NvCreateChatCompletionRequest` (Rust -- request side) + +New fields relevant to RL: + +| Field | Serialized | Description | +|-------|-----------|-------------| +| `tokens` | No (`skip_serializing`) | Pre-tokenized prompt token IDs (TITO path via `/v1/chat/completions/tokens`) | +| `return_token_ids` | No (`skip_serializing`) | prime-rl compat field; accepted but ignored on the standard endpoint -- use `DYN_ENABLE_RL` or `nvext.extra_fields` instead | + +Both fields are stripped before the request is forwarded to the vLLM worker, preventing 400 errors from the vLLM OpenAI-compatible API. + +### `NvCreateChatCompletionResponse` (Rust -- response side) + +``` +NvCreateChatCompletionResponse { + inner: CreateChatCompletionResponse -- standard OpenAI response fields + nvext?: serde_json::Value -- NvExtResponse serialized as JSON + prompt_token_ids?: Vec -- RL: tokenized prompt IDs (DYN_ENABLE_RL only) +} +``` + +### `DeltaGenerator` (Rust -- streaming pipeline) + +Manages per-request streaming state. Accumulates output token IDs across chunks: + +``` +DeltaGenerator { + ... + accumulated_completion_token_ids: Vec -- grows per chunk +} +``` + +- **Activation:** `options.enable_completion_token_ids` is set to `true` when `extra_fields` includes `"completion_token_ids"` (auto-set when `DYN_ENABLE_RL=true`). +- **Accumulation:** On each postprocessor output chunk, appends `delta.token_ids` to the accumulator. +- **Emission:** On the final chunk (`finish_reason` is set), the full list is emitted in `nvext.completion_token_ids` and the accumulator is cleared. + +### `NvExt` (Rust -- request-side NVIDIA extensions) + +Relevant fields for RL: + +``` +NvExt { + token_data?: Vec -- Pre-tokenized prompt IDs (TITO / EPP bypass) + extra_fields?: Vec -- Request extra response fields, e.g. ["completion_token_ids"] + backend_instance_id?: u64 -- Targeted routing to a specific worker + ... +} +``` + +### Tokenization Types + +``` +TokenizeRequest = Completion { prompt, model?, add_special_tokens? } + | Chat { messages, model?, add_generation_prompt?, chat_template?, ... } + +TokenizeResponse { count: int, max_model_len: int, tokens: Vec, token_strs?: Vec } + +DetokenizeRequest { model?: String, tokens: Vec } +DetokenizeResponse { prompt: String } +``` + +### Post-Processing Helpers + +**`rl_tokenize_prompt(state, model, messages) -> Option>`** + +Tokenizes the original prompt messages using the model's chat template and tokenizer: resolves model card from `state`, gets the tokenizer instance, builds a `PromptFormatter` from the model deployment card, renders messages through the chat template (same logic as the preprocessor), tokenizes the rendered string, and returns the token IDs. + +**`rl_promote_token_ids_in_response(json_val)`** + +Copies `response.nvext.completion_token_ids` to `response.choices[i].token_ids` for each choice. Bridges Dynamo's `nvext` convention with the field paths that Prime-RL/verifiers expects. + +--- + +## 7. Worker Engine Routes (Internal) + +Five engine route handlers are registered on each vLLM worker's system HTTP port (default `8081` local / `9090` in k8s). These are **internal** routes called by the Rust frontend's `/v1/rl/*` handlers -- not called directly by Prime-RL. + +| Route | Method | vLLM API called | Description | +|-------|--------|-----------------|-------------| +| `/engine/pause_generation` | POST | `engine_client.pause_generation()` | Drain in-flight requests, keep model loaded in GPU memory | +| `/engine/resume_generation` | POST | `engine_client.resume_generation()` | Resume accepting inference requests | +| `/engine/flush_cache` | POST | `engine_client.reset_prefix_cache()` | Invalidate prefix/KV cache (required before weight reload) | +| `/engine/update_weights_from_path` | POST | `engine_client.collective_rpc("reload_weights", ...)` | Load weights from filesystem (safetensors checkpoint) | +| `/engine/get_weight_version` | POST | `self._weight_version` | Return current weight version string | + +Both decode and prefill worker types register all 5 routes. Route signatures are compatible with SGLang's merged `#6094` routes for backend interoperability. + +### Registration (worker_factory.py) + +```python +runtime.register_engine_route("pause_generation", handler.pause_generation) +runtime.register_engine_route("resume_generation", handler.resume_generation) +runtime.register_engine_route("flush_cache", handler.flush_cache) +runtime.register_engine_route("update_weights_from_path", handler.update_weights_from_path) +runtime.register_engine_route("get_weight_version", handler.get_weight_version) +``` + +### publisher.py Crash Guard + +The `DynamoStatLoggerPublisher.record()` method includes a guard for `scheduler_stats is None`. This prevents an `AttributeError` crash during the transient window right after a weight reload when the vLLM engine's stats logger fires before the engine core has re-initialized its scheduler stats. + +--- + +## 8. Known Limitations + +| Limitation | Workaround | Notes | +|-----------|-----------|-------| +| `cache_salt` not supported -- returns 400 for requests with `cache_salt` in body | Set `[experimental] use_prefix_cache_salt = false` in prime-rl `orch.toml` | verifiers dev6+ defaults `use_prefix_cache_salt=True` | +| `prompt_token_ids` only injected for non-streaming responses | Use non-streaming mode for RL rollouts (the default) | Streaming final-chunk injection is planned | +| Weight version `"initial"` before first update | Do not depend on version string for correctness; use `/v1/rl/ready` for readiness | | +| NCCL weight broadcast is a no-op on Dynamo side | Use `type = "filesystem"` in `[weight_broadcast]` for all current deployments | | +| ~~`VLLM_USE_V1=0` required~~ **Resolved on vLLM 0.19.1** | Set `VLLM_USE_V1=1` on images shipping `VLLM_VER≥0.19.1` (current default). Keep `VLLM_USE_V1=0` only on legacy 0.18.x images where Meta-tensor crash with `--enforce-eager` still reproduces. | Verified under Run D (Qwen3.5-35B-A3B-FP8, 12 workers, batch=64) with V1 enabled and CUDA graphs for LoRA decode. | +| Filesystem weight broadcast scales poorly for large models | Acceptable for 0.6B (257ms load); marginal at 7B (~25s); impractical at 70B (~5 min) | RDMA pull transfer planned | + +--- + +## 9. Validation Results + +### Local (2x A6000, Qwen3-0.6B-Reverse-Text-SFT, 20 steps) + +| Metric | vLLM Baseline | Dynamo | Delta | +|--------|:-------------:|:------:|:-----:| +| Steps completed | 20/20 | 20/20 | -- | +| Peak reward | 0.798 | **0.825** | +3.4% | +| Final reward | 0.716 | **0.724** | +1.2% | +| `is_masked/mean` | 1.2% | **0.13%** | -92% (better) | +| Mismatch KL (final) | 0.0075 | **0.0056** | -25% (better) | +| Weight update cycles | 19 | 19 | -- | +| Mean weight cycle time | -- | 257.5ms | pause: 3.2ms, load: 249ms, resume: 4.8ms | + +W&B: https://wandb.ai/test232/prime-rl-parity-apr17 + +### Kubernetes (GB200, same model, 20 steps) + +| Metric | K8s | +|--------|:---:| +| Steps completed | 20/20 | +| Reward at step 13 | 0.714 (climbing) | +| Mismatch KL (steps 0-3) | 0.0007 - 0.0009 | +| Pods | 4 | +| All RL routes verified | Yes | + +The Rust RL API produces **better token alignment than native vLLM** (0.13% masked vs 1.2%). diff --git a/lib/llm/src/audit/stream.rs b/lib/llm/src/audit/stream.rs index 2663cdd7eeb0..4985139d37c8 100644 --- a/lib/llm/src/audit/stream.rs +++ b/lib/llm/src/audit/stream.rs @@ -101,6 +101,7 @@ where service_tier: None, }, nvext: None, + prompt_token_ids: None, } }) }), @@ -138,6 +139,7 @@ where service_tier: None, }, nvext: None, + prompt_token_ids: None, }; let _ = tx.send(fallback.clone()); final_response_to_one_chunk_stream(fallback) @@ -160,6 +162,7 @@ where service_tier: None, }, nvext: None, + prompt_token_ids: None, } }) }); diff --git a/lib/llm/src/entrypoint/input/text.rs b/lib/llm/src/entrypoint/input/text.rs index 1c0138fd34b3..6c7046602d7e 100644 --- a/lib/llm/src/entrypoint/input/text.rs +++ b/lib/llm/src/entrypoint/input/text.rs @@ -116,6 +116,8 @@ async fn main_loop( chat_template_args: None, media_io_kwargs: None, unsupported_fields: Default::default(), + tokens: None, + return_token_ids: None, }; // Call the model diff --git a/lib/llm/src/http/service/openai.rs b/lib/llm/src/http/service/openai.rs index 85e5d28339f3..c67f7bbfa6ef 100644 --- a/lib/llm/src/http/service/openai.rs +++ b/lib/llm/src/http/service/openai.rs @@ -55,6 +55,10 @@ use crate::protocols::openai::{ embeddings::{NvCreateEmbeddingRequest, NvCreateEmbeddingResponse}, images::{NvCreateImageRequest, NvImagesResponse}, responses::{NvCreateResponse, NvResponse, ResponseParams, chat_completion_to_response}, + tokenization::{ + DetokenizeRequest, DetokenizeResponse, TokenizeCompletionRequest, TokenizeRequest, + TokenizeResponse, + }, videos::{NvCreateVideoRequest, NvVideosResponse}, }; use crate::protocols::unified::UnifiedRequest; @@ -202,6 +206,21 @@ impl ErrorMessage { /// Not Implemented Error /// Return this error when the client requests a feature that is not yet implemented. /// This should be used for features that are planned but not available. + /// Bad Request Error + /// Return this error when the client sends an invalid request. + pub fn bad_request(msg: &str) -> ErrorResponse { + let code = StatusCode::BAD_REQUEST; + let error_type = map_error_code_to_error_type(code); + ( + code, + Json(ErrorMessage { + message: msg.to_string(), + error_type, + code: code.as_u16(), + }), + ) + } + pub fn not_implemented_error(msg: T) -> ErrorResponse { tracing::error!("Not Implemented error: {msg}"); let code = StatusCode::NOT_IMPLEMENTED; @@ -875,6 +894,69 @@ async fn handler_chat_completions( request.nvext = apply_header_routing_overrides(request.nvext.take(), &headers); + // RL field promotion: wire `tokens` and `return_token_ids` when provided on the standard + // chat completions endpoint. This eliminates the need for the rl_admin Python proxy to + // intercept and rewrite these fields. + // + // If `return_token_ids` is true, request completion_token_ids in the response. + // Auto-enable when DYN_ENABLE_RL is set -- ensures token IDs flow even if the + // client forgets to request them. + let rl_want_token_ids = request + .return_token_ids + .take() + .unwrap_or_else(|| dynamo_runtime::config::env_is_truthy("DYN_ENABLE_RL")); + if rl_want_token_ids { + tracing::info!("RL: want_token_ids=true, will promote nvext.extra_fields"); + } + { + // If `tokens` is provided, inject into nvext.token_data (pre-tokenized prompt path). + let token_data = request.tokens.take(); + + if token_data.is_some() || rl_want_token_ids { + let mut nvext = request.nvext.take().unwrap_or_default(); + + if let Some(ids) = token_data { + if !ids.is_empty() { + nvext.token_data = Some(ids); + // Ensure messages is non-empty for model lookup / chat template + if request.inner.messages.is_empty() { + use dynamo_protocols::types::{ + ChatCompletionRequestMessage, ChatCompletionRequestUserMessage, + ChatCompletionRequestUserMessageContent, + }; + request + .inner + .messages + .push(ChatCompletionRequestMessage::User( + ChatCompletionRequestUserMessage { + content: ChatCompletionRequestUserMessageContent::Text( + "(token-in mode)".to_string(), + ), + name: None, + }, + )); + } + } + } + + if rl_want_token_ids { + let mut extra_fields = nvext.extra_fields.take().unwrap_or_default(); + for field in &["token_ids", "completion_token_ids"] { + if !extra_fields.contains(&field.to_string()) { + extra_fields.push(field.to_string()); + } + } + nvext.extra_fields = Some(extra_fields); + // Also force logprobs on when RL is requesting token IDs + if request.inner.logprobs.is_none() { + request.inner.logprobs = Some(true); + } + } + + request.nvext = Some(nvext); + } + } + // create the context for the request let request_id = get_or_create_request_id(&headers); let streaming = request.inner.stream.unwrap_or(false); @@ -1159,6 +1241,26 @@ async fn chat_completions( // todo - decide on default let streaming = request.inner.stream.unwrap_or(false); + // RL: save messages for post-response prompt tokenization (needed for prompt_token_ids). + let rl_saved_messages = if dynamo_runtime::config::env_is_truthy("DYN_ENABLE_RL") && !streaming + { + Some(request.inner.messages.clone()) + } else { + None + }; + + // RL: for TITO requests the caller (handler_chat_completions_tokens) injects a + // placeholder message so Dynamo can select a chat template, but then saves the + // real token IDs in nvext.token_data. Capture them now — before the request is + // consumed by engine.generate() — so the post-processing step can use them + // directly as prompt_token_ids instead of re-tokenizing the placeholder. + let rl_tito_token_ids: Option> = + if dynamo_runtime::config::env_is_truthy("DYN_ENABLE_RL") && !streaming { + request.nvext.as_ref().and_then(|nv| nv.token_data.clone()) + } else { + None + }; + // Apply template values first to resolve the model before creating metrics guards if let Some(template) = template { if request.inner.model.is_empty() { @@ -1370,6 +1472,41 @@ async fn chat_completions( if ctx.is_killed() { inflight_guard.mark_error(ErrorType::Cancelled); } + + // RL post-processing: when DYN_ENABLE_RL is active, promote + // token IDs to the top-level locations that Prime-RL / verifiers expects: + // response.prompt_token_ids (from tokenizing the prompt) + // response.choices[i].token_ids (from nvext.completion_token_ids) + let response = if let Some(ref messages) = rl_saved_messages { + let mut response = response; + // For TITO requests, nvext.token_data IS the prompt — use those IDs + // directly. Falling back to rl_tokenize_prompt would re-tokenize the + // placeholder message injected by handler_chat_completions_tokens and + // return the wrong IDs. + response.prompt_token_ids = + rl_tito_token_ids.or_else(|| rl_tokenize_prompt(&state, &model, messages)); + match serde_json::to_value(&response) { + Ok(mut json_val) => { + rl_promote_token_ids_in_response(&mut json_val); + return Ok(Json(json_val).into_response()); + } + Err(e) => { + // This path means choice.token_ids will NOT be promoted — Prime-RL + // will see None for completion token IDs and may silently drop the + // rollout or crash. Log at error so data-loss does not go unnoticed. + tracing::error!( + request_id, + "rl_promote_token_ids: serde_json serialization failed — \ + choice.token_ids will NOT be promoted to top-level; \ + Prime-RL rollout may be dropped or corrupt: {e}" + ); + } + } + response + } else { + response + }; + Ok(Json(response).into_response()) } } @@ -1842,6 +1979,179 @@ pub(crate) fn check_ready(_state: &Arc) -> Result<(), ErrorRe Ok(()) } +// ── Tokenize / Detokenize ──────────────────────────────────────────── + +fn bad_request>(message: T) -> ErrorResponse { + let code = StatusCode::BAD_REQUEST; + ( + code, + Json(ErrorMessage { + message: message.into(), + error_type: map_error_code_to_error_type(code), + code: code.as_u16(), + }), + ) +} + +fn resolve_tokenizer_model_name( + state: &Arc, + requested_model: Option<&str>, +) -> Result { + if let Some(model) = requested_model { + if state.manager().has_model_any(model) { + return Ok(model.to_string()); + } + return Err(ErrorMessage::model_not_found()); + } + let served_models = state.manager().model_display_names(); + if served_models.len() == 1 { + return Ok(served_models.into_iter().next().unwrap()); + } + Err(bad_request( + "Model must be specified when more than one model is served.", + )) +} + +fn resolve_model_card( + state: &Arc, + requested_model: Option<&str>, +) -> Result<(String, crate::model_card::ModelDeploymentCard), ErrorResponse> { + let model = resolve_tokenizer_model_name(state, requested_model)?; + let card = state + .manager() + .get_model_cards() + .into_iter() + .find(|card| card.display_name == model) + .ok_or_else(|| { + ErrorMessage::internal_server_error(&format!( + "Tokenizer metadata is not available for model '{}'", + model + )) + })?; + Ok((model, card)) +} + +async fn tokenize( + State(state): State>, + Json(request): Json, +) -> Result { + check_ready(&state)?; + + let (_, card) = resolve_model_card(&state, request.model())?; + let tokenizer = card + .tokenizer() + .map_err(|err| ErrorMessage::from_anyhow(err, "Failed to load tokenizer"))?; + + let (tokens, token_strs) = match request { + TokenizeRequest::Completion(TokenizeCompletionRequest { + prompt, + add_special_tokens, + return_token_strs, + .. + }) => { + let encoding = tokenizer + .encode_with_special_tokens(&prompt, add_special_tokens) + .map_err(|err| ErrorMessage::from_anyhow(err, "Failed to tokenize prompt"))?; + let token_ids = encoding.token_ids().to_vec(); + let token_strs = if return_token_strs { + Some(tokenizer.convert_ids_to_tokens(&token_ids).map_err(|err| { + ErrorMessage::from_anyhow(err, "Failed to resolve token strings") + })?) + } else { + None + }; + (token_ids, token_strs) + } + TokenizeRequest::Chat(request) => { + let model = request + .model + .clone() + .unwrap_or_else(|| card.display_name.clone()); + // Render the chat messages to a prompt string via the model's chat template + let formatter = crate::preprocessor::prompt::PromptFormatter::from_mdc(&card) + .map_err(|err| ErrorMessage::from_anyhow(err, "Failed to build chat formatter"))?; + let inner_request = dynamo_protocols::types::CreateChatCompletionRequest { + model, + messages: request.messages.clone(), + tools: request.tools.clone(), + ..Default::default() + }; + let wrapped = + crate::protocols::openai::chat_completions::NvCreateChatCompletionRequest { + inner: inner_request, + common: Default::default(), + nvext: None, + chat_template_args: Some(request.merged_chat_template_kwargs()), + media_io_kwargs: None, + tokens: None, + return_token_ids: None, + unsupported_fields: Default::default(), + }; + let prompt = match formatter { + crate::preprocessor::prompt::PromptFormatter::OAI(f) => f.render(&wrapped), + } + .map_err(|err| ErrorMessage::from_anyhow(err, "Failed to render chat prompt"))?; + + let encoding = tokenizer + .encode_with_special_tokens(&prompt, request.add_special_tokens) + .map_err(|err| { + ErrorMessage::from_anyhow(err, "Failed to tokenize rendered chat prompt") + })?; + let token_ids = encoding.token_ids().to_vec(); + let token_strs = if request.return_token_strs { + Some(tokenizer.convert_ids_to_tokens(&token_ids).map_err(|err| { + ErrorMessage::from_anyhow(err, "Failed to resolve token strings") + })?) + } else { + None + }; + (token_ids, token_strs) + } + }; + + Ok(Json(TokenizeResponse { + count: tokens.len(), + max_model_len: card.context_length, + tokens, + token_strs, + }) + .into_response()) +} + +async fn detokenize( + State(state): State>, + Json(request): Json, +) -> Result { + check_ready(&state)?; + + let (_, card) = resolve_model_card(&state, request.model.as_deref())?; + let tokenizer = card + .tokenizer() + .map_err(|err| ErrorMessage::from_anyhow(err, "Failed to load tokenizer"))?; + let prompt: String = tokenizer + .decode(&request.tokens, false) + .map_err(|err| ErrorMessage::from_anyhow(err, "Failed to detokenize prompt"))? + .into(); + + Ok(Json(DetokenizeResponse { prompt }).into_response()) +} + +pub fn tokenization_router(state: Arc) -> (Vec, Router) { + let tokenize_path = "/v1/tokenize"; + let detokenize_path = "/v1/detokenize"; + let docs = vec![ + RouteDoc::new(axum::http::Method::POST, tokenize_path), + RouteDoc::new(axum::http::Method::POST, detokenize_path), + ]; + let router = Router::new() + .route(tokenize_path, post(tokenize)) + .route(detokenize_path, post(detokenize)) + .layer(middleware::from_fn(smart_json_error_middleware)) + .layer(axum::extract::DefaultBodyLimit::max(get_body_limit())) + .with_state(state); + (docs, router) +} + /// openai compatible format /// Example: /// { @@ -1953,6 +2263,137 @@ pub fn chat_completions_router( (vec![doc], router) } +/// Create an Axum [`Router`] for the RL TITO (Token-In / Token-Out) endpoint. +/// +/// This endpoint accepts Prime-RL's `tokens` field (pre-tokenized prompt), +/// translates it to `nvext.token_data`, forces logprobs on, and delegates +/// to the standard chat_completions handler -- all in Rust, eliminating the +/// Python rl-admin proxy from the hot inference path. +/// +/// If no path is provided, the default path is `/v1/chat/completions/tokens` +pub fn chat_completions_tokens_router( + state: Arc, + template: Option, + path: Option, +) -> (Vec, Router) { + let path = path.unwrap_or("/v1/chat/completions/tokens".to_string()); + let doc = RouteDoc::new(axum::http::Method::POST, &path); + let router = Router::new() + .route(&path, post(handler_chat_completions_tokens)) + .layer(middleware::from_fn(smart_json_error_middleware)) + .layer(axum::extract::DefaultBodyLimit::max(get_body_limit())) + .with_state((state, template)); + (vec![doc], router) +} + +/// Handler for TITO (Token-In / Token-Out) chat completions. +/// +/// Accepts Prime-RL's request format which includes a `tokens` field containing +/// pre-tokenized prompt token IDs. The handler: +/// 1. Extracts the `tokens` field +/// 2. Injects them as `nvext.token_data` (Dynamo's native pre-tokenized input) +/// 3. Requests `token_ids` and `completion_token_ids` in the response via `nvext.extra_fields` +/// 4. Forces `logprobs = true` (RL always needs logprobs) +/// 5. Ensures `messages` is non-empty (Dynamo requires it for chat template selection) +/// 6. Delegates to the standard `chat_completions()` internal function (zero HTTP proxy) +async fn handler_chat_completions_tokens( + State((state, template)): State<(Arc, Option)>, + headers: HeaderMap, + Json(mut request): Json, +) -> Result { + check_ready(&state)?; + + // Extract the tokens field (Prime-RL's TITO input) + let tokens = request.tokens.take(); + // Clear return_token_ids (not supported by Dynamo, avoid confusion) + request.return_token_ids = None; + + if let Some(token_ids) = tokens { + if token_ids.is_empty() { + return Err(ErrorMessage::bad_request( + "TITO endpoint requires non-empty 'tokens' field", + )); + } + + // Inject tokens into nvext.token_data + let mut nvext = request.nvext.take().unwrap_or_default(); + nvext.token_data = Some(token_ids); + + // Request token echo and completion token IDs in response + let mut extra_fields = nvext.extra_fields.take().unwrap_or_default(); + for field in &["token_ids", "completion_token_ids"] { + if !extra_fields.contains(&field.to_string()) { + extra_fields.push(field.to_string()); + } + } + nvext.extra_fields = Some(extra_fields); + request.nvext = Some(nvext); + + // Force logprobs on (RL always needs them) + if request.inner.logprobs.is_none() { + request.inner.logprobs = Some(true); + } + + // Ensure messages is non-empty (Dynamo requires it for model lookup / chat template) + if request.inner.messages.is_empty() { + use dynamo_protocols::types::{ + ChatCompletionRequestMessage, ChatCompletionRequestUserMessage, + ChatCompletionRequestUserMessageContent, + }; + request + .inner + .messages + .push(ChatCompletionRequestMessage::User( + ChatCompletionRequestUserMessage { + content: ChatCompletionRequestUserMessageContent::Text( + "(token-in mode)".to_string(), + ), + name: None, + }, + )); + } + } else { + return Err(ErrorMessage::bad_request( + "Missing 'tokens' field for TITO endpoint. \ + Use /v1/chat/completions for message-based requests.", + )); + } + + // Apply header routing overrides + request.nvext = apply_header_routing_overrides(request.nvext.take(), &headers); + + // Delegate to the standard chat completions flow (no HTTP proxy!) + let request_id = get_or_create_request_id(&headers); + let streaming = request.inner.stream.unwrap_or(false); + let cancellation_labels = CancellationLabels { + model: request.inner.model.clone(), + endpoint: Endpoint::ChatCompletions.to_string(), + request_type: if streaming { "stream" } else { "unary" }.to_string(), + }; + let request = Context::with_id(request, request_id); + let context = request.context(); + + let (mut connection_handle, stream_handle) = create_connection_monitor( + context.clone(), + Some(state.metrics_clone()), + cancellation_labels, + ) + .await; + + let response = + tokio::spawn(chat_completions(state, template, request, stream_handle).in_current_span()) + .await + .map_err(|e| { + ErrorMessage::internal_server_error(&format!( + "Failed to await TITO chat completions task: {:?}", + e, + )) + })?; + + connection_handle.disarm(); + response +} + /// Create an Axum [`Router`] for the OpenAI API Embeddings endpoint /// If not path is provided, the default path is `/v1/embeddings` pub fn embeddings_router( @@ -2534,6 +2975,511 @@ pub fn audios_router( (vec![doc], router) } +// ────────────────────────────────────────────────────────────────────────── +// RL Admin router: /v1/rl/* +// ────────────────────────────────────────────────────────────────────────── + +/// Environment variable for comma-separated worker system HTTP URLs. +/// Defaults to `http://localhost:8081` when not set. +const DYN_RL_WORKER_SYSTEM_URLS_ENV: &str = "DYN_RL_WORKER_SYSTEM_URLS"; + +/// Shared state for the RL admin router. +#[derive(Clone)] +struct RlState { + /// Worker system HTTP base URLs (e.g. `http://localhost:8081`). + /// Set via `DYN_RL_WORKER_SYSTEM_URLS` (comma-separated list). + worker_system_urls: Vec, + /// Shared HTTP client for all fan-out calls to worker system ports. + http_client: reqwest::Client, +} + +impl RlState { + fn from_env() -> Self { + let worker_system_urls = std::env::var(DYN_RL_WORKER_SYSTEM_URLS_ENV) + .unwrap_or_else(|_| "http://localhost:8081".to_string()) + .split(',') + .map(|s| s.trim().to_string()) + .filter(|s| !s.is_empty()) + .collect::>(); + tracing::info!( + "RL admin router configured with {} worker(s): {:?}", + worker_system_urls.len(), + worker_system_urls + ); + Self { + worker_system_urls, + http_client: reqwest::Client::builder() + .timeout(std::time::Duration::from_secs(600)) + .build() + .expect("Failed to create RL router HTTP client"), + } + } + + /// Call a single engine route on one worker. Returns the JSON body. + async fn call_engine_route( + &self, + url: &str, + route: &str, + body: &serde_json::Value, + ) -> serde_json::Value { + let endpoint = format!("{url}/engine/{route}"); + match self.http_client.post(&endpoint).json(body).send().await { + Ok(resp) => { + let status = resp.status(); + match resp.json::().await { + Ok(v) => v, + Err(e) => serde_json::json!({ + "status": "error", + "message": format!("Failed to decode response from {endpoint}: {e}"), + "http_status": status.as_u16() + }), + } + } + Err(e) => serde_json::json!({ + "status": "error", + "message": format!("Request to {endpoint} failed: {e}") + }), + } + } + + /// Fan out an engine route call to all configured workers concurrently. + async fn fan_out(&self, route: &str, body: serde_json::Value) -> Vec { + let futures: Vec<_> = self + .worker_system_urls + .iter() + .map(|url| self.call_engine_route(url, route, &body)) + .collect(); + futures::future::join_all(futures).await + } + + /// Returns true only if all results have `status: "ok"`. + fn all_ok(results: &[serde_json::Value]) -> bool { + results + .iter() + .all(|r| r.get("status").and_then(|s| s.as_str()) == Some("ok")) + } +} + +/// `GET /v1/rl/ready` — composite readiness check: worker health via system port. +async fn rl_ready(State(state): State>) -> impl IntoResponse { + let futures: Vec<_> = state + .worker_system_urls + .iter() + .map(|url| { + let client = state.http_client.clone(); + let health_url = format!("{url}/health"); + async move { + client + .get(&health_url) + .send() + .await + .map(|r| r.status().is_success()) + .unwrap_or(false) + } + }) + .collect(); + let results = futures::future::join_all(futures).await; + let all_ready = !results.is_empty() && results.iter().all(|ok| *ok); + if all_ready { + (StatusCode::OK, Json(serde_json::json!({"status": "ready"}))) + } else { + ( + StatusCode::SERVICE_UNAVAILABLE, + Json(serde_json::json!({ + "status": "not_ready", + "workers_ready": results.iter().filter(|ok| **ok).count(), + "workers_total": results.len() + })), + ) + } +} + +/// `POST /v1/rl/pause` — fan out `pause_generation` to all workers. +async fn rl_pause(State(state): State>) -> impl IntoResponse { + let results = state + .fan_out("pause_generation", serde_json::json!({})) + .await; + if RlState::all_ok(&results) { + tracing::info!("RL pause: all {} worker(s) paused", results.len()); + ( + StatusCode::OK, + Json(serde_json::json!({"status": "ok", "workers": results})), + ) + } else { + tracing::warn!("RL pause: some workers failed: {:?}", results); + ( + StatusCode::BAD_GATEWAY, + Json(serde_json::json!({"status": "error", "workers": results})), + ) + } +} + +/// `POST /v1/rl/resume` — fan out `resume_generation` to all workers. +async fn rl_resume(State(state): State>) -> impl IntoResponse { + let results = state + .fan_out("resume_generation", serde_json::json!({})) + .await; + if RlState::all_ok(&results) { + tracing::info!("RL resume: all {} worker(s) resumed", results.len()); + ( + StatusCode::OK, + Json(serde_json::json!({"status": "ok", "workers": results})), + ) + } else { + tracing::warn!("RL resume: some workers failed: {:?}", results); + ( + StatusCode::BAD_GATEWAY, + Json(serde_json::json!({"status": "error", "workers": results})), + ) + } +} + +/// `POST /v1/rl/update_weights` — atomic `flush_cache → update_weights_from_path` across all workers. +/// +/// Expected body: `{"weight_dir": "/path/to/checkpoint"}` or `{"weight_dir": null}` for NCCL mode. +/// +/// The sequence per worker is: `flush_cache → update_weights_from_path`. +/// The pause/resume envelope is left to Prime-RL, which can call `/v1/rl/pause` and +/// `/v1/rl/resume` explicitly for full drain-and-swap semantics. +async fn rl_update_weights( + State(state): State>, + body: axum::extract::Json, +) -> impl IntoResponse { + let weight_dir = body + .get("weight_dir") + .and_then(|v| v.as_str()) + .map(str::to_string); + + if weight_dir.is_none() { + tracing::info!("RL update_weights: weight_dir=null (NCCL mode, no-op on Dynamo side)"); + return ( + StatusCode::OK, + Json(serde_json::json!({"status": "ok", "message": "NCCL mode, no-op on Dynamo side"})), + ); + } + + let weight_dir = weight_dir.unwrap(); + tracing::info!("RL update_weights: weight_dir={weight_dir}"); + + // Step 1: flush_cache across all workers + let flush_results = state.fan_out("flush_cache", serde_json::json!({})).await; + if !RlState::all_ok(&flush_results) { + tracing::warn!("RL update_weights: flush_cache failed: {:?}", flush_results); + return ( + StatusCode::BAD_GATEWAY, + Json( + serde_json::json!({"status": "error", "stage": "flush_cache", "workers": flush_results}), + ), + ); + } + + // Step 2: update_weights_from_path across all workers + let version = std::path::Path::new(&weight_dir) + .file_name() + .and_then(|n| n.to_str()) + .unwrap_or("unknown") + .to_string(); + let load_body = serde_json::json!({"path": weight_dir, "version": version}); + let load_results = state.fan_out("update_weights_from_path", load_body).await; + if RlState::all_ok(&load_results) { + tracing::info!( + "RL update_weights: all {} worker(s) updated weights to {weight_dir}", + load_results.len() + ); + ( + StatusCode::OK, + Json(serde_json::json!({"status": "ok", "workers": load_results})), + ) + } else { + tracing::warn!( + "RL update_weights: update_weights_from_path failed: {:?}", + load_results + ); + ( + StatusCode::BAD_GATEWAY, + Json( + serde_json::json!({"status": "error", "stage": "update_weights_from_path", "workers": load_results}), + ), + ) + } +} + +/// `POST /v1/rl/load_lora_adapter` — hot-load/swap a LoRA adapter from a filesystem path. +/// +/// Expected body: `{"lora_name": "r16-a32.0", "lora_path": "/path/to/adapter_dir"}` +/// +/// The adapter directory must contain PEFT-style `adapter_model.safetensors` and +/// `adapter_config.json`. This is the RL-specific LoRA path used by Prime-RL every +/// training step (separate from Dynamo's URI-based `load_lora` gRPC endpoint which +/// downloads adapters from S3/file URIs and publishes a new ModelDeploymentCard). +/// +/// Hot-swap semantics: calling with a `lora_name` that is already loaded removes +/// the previous adapter and loads the new one under the same deterministic int ID, +/// then resets the prefix cache so stale KV entries don't poison new rollouts. +/// +/// Pair with `/v1/rl/pause` and `/v1/rl/resume` for a full drain-swap-resume cycle. +async fn rl_load_lora_adapter( + State(state): State>, + body: axum::extract::Json, +) -> impl IntoResponse { + let lora_name = body.get("lora_name").and_then(|v| v.as_str()); + let lora_path = body.get("lora_path").and_then(|v| v.as_str()); + + let (lora_name, lora_path) = match (lora_name, lora_path) { + (Some(n), Some(p)) if !n.is_empty() && !p.is_empty() => (n.to_string(), p.to_string()), + _ => { + return ( + StatusCode::BAD_REQUEST, + Json(serde_json::json!({ + "status": "error", + "message": "Expected body: {\"lora_name\": str, \"lora_path\": str} (both required, non-empty)" + })), + ); + } + }; + + tracing::info!("RL load_lora_adapter: lora_name={lora_name} lora_path={lora_path}"); + let results = state + .fan_out( + "load_lora_adapter", + serde_json::json!({"lora_name": lora_name, "lora_path": lora_path}), + ) + .await; + + if RlState::all_ok(&results) { + tracing::info!( + "RL load_lora_adapter: all {} worker(s) loaded LoRA '{lora_name}' from {lora_path}", + results.len() + ); + ( + StatusCode::OK, + Json(serde_json::json!({"status": "ok", "workers": results})), + ) + } else { + tracing::warn!("RL load_lora_adapter: some workers failed: {:?}", results); + ( + StatusCode::BAD_GATEWAY, + Json(serde_json::json!({"status": "error", "workers": results})), + ) + } +} + +/// `POST /v1/rl/unload_lora_adapter` — remove a previously loaded LoRA adapter by name. +/// +/// Expected body: `{"lora_name": "r16-a32.0"}` +/// +/// Idempotent: unloading an already-absent LoRA returns `status: ok` so callers +/// can retry safely without special-casing not-found. +async fn rl_unload_lora_adapter( + State(state): State>, + body: axum::extract::Json, +) -> impl IntoResponse { + let lora_name = body + .get("lora_name") + .and_then(|v| v.as_str()) + .map(str::to_string); + + let lora_name = match lora_name { + Some(n) if !n.is_empty() => n, + _ => { + return ( + StatusCode::BAD_REQUEST, + Json(serde_json::json!({ + "status": "error", + "message": "Expected body: {\"lora_name\": str} (required, non-empty)" + })), + ); + } + }; + + tracing::info!("RL unload_lora_adapter: lora_name={lora_name}"); + let results = state + .fan_out( + "unload_lora_adapter", + serde_json::json!({"lora_name": lora_name}), + ) + .await; + + if RlState::all_ok(&results) { + tracing::info!( + "RL unload_lora_adapter: all {} worker(s) unloaded LoRA '{lora_name}'", + results.len() + ); + ( + StatusCode::OK, + Json(serde_json::json!({"status": "ok", "workers": results})), + ) + } else { + tracing::warn!("RL unload_lora_adapter: some workers failed: {:?}", results); + ( + StatusCode::BAD_GATEWAY, + Json(serde_json::json!({"status": "error", "workers": results})), + ) + } +} + +/// `GET /v1/rl/weight_version` — query weight version from all workers. +async fn rl_weight_version(State(state): State>) -> impl IntoResponse { + let results = state + .fan_out("get_weight_version", serde_json::json!({})) + .await; + + // Collect distinct versions and check for consistency + let versions: Vec<_> = results + .iter() + .filter_map(|r| { + r.get("version") + .and_then(|v| v.as_str()) + .map(str::to_string) + }) + .collect(); + + let unique: std::collections::HashSet<&str> = versions.iter().map(String::as_str).collect(); + if unique.len() == 1 { + ( + StatusCode::OK, + Json(serde_json::json!({ + "status": "ok", + "version": unique.into_iter().next().unwrap_or(""), + "workers": results + })), + ) + } else { + ( + StatusCode::OK, + Json(serde_json::json!({ + "status": "inconsistent", + "versions": unique.into_iter().collect::>(), + "workers": results + })), + ) + } +} + +/// Promote token IDs from the Dynamo `nvext` response object to the top-level +/// locations that Prime-RL / verifiers expects: +/// +/// response.nvext.completion_token_ids → response.choices[i].token_ids +/// +/// Tokenize chat messages using the model's tokenizer and return prompt token IDs. +/// Used by the RL post-processing path to populate `response.prompt_token_ids`. +fn rl_tokenize_prompt( + state: &Arc, + model: &str, + messages: &[dynamo_protocols::types::ChatCompletionRequestMessage], +) -> Option> { + if messages.is_empty() { + return None; + } + let (_, card) = resolve_model_card(state, Some(model)).ok()?; + let tokenizer = card.tokenizer().ok()?; + let formatter = crate::preprocessor::prompt::PromptFormatter::from_mdc(&card).ok()?; + let inner_request = dynamo_protocols::types::CreateChatCompletionRequest { + model: model.to_string(), + messages: messages.to_vec(), + ..Default::default() + }; + let wrapped = crate::protocols::openai::chat_completions::NvCreateChatCompletionRequest { + inner: inner_request, + common: Default::default(), + nvext: None, + chat_template_args: None, + media_io_kwargs: None, + tokens: None, + return_token_ids: None, + unsupported_fields: Default::default(), + }; + let prompt = match formatter { + crate::preprocessor::prompt::PromptFormatter::OAI(f) => f.render(&wrapped), + } + .ok()?; + let encoding = tokenizer.encode_with_special_tokens(&prompt, true).ok()?; + Some(encoding.token_ids().to_vec()) +} + +/// This lets Prime-RL read `choice.token_ids` without knowing about the `nvext` +/// extension structure. Called on non-streaming responses when RL token ID mode +/// is active. +fn rl_promote_token_ids_in_response(json_val: &mut serde_json::Value) { + // Move completion_token_ids from response-level nvext to each choice. + // Prime-RL / verifiers expects: + // response.choices[i].token_ids (not response.nvext.completion_token_ids) + let has_nvext = json_val.get("nvext").is_some(); + let has_completion_ids = json_val + .get("nvext") + .and_then(|nv| nv.get("completion_token_ids")) + .is_some(); + + tracing::debug!( + has_nvext, + has_completion_ids, + "rl_promote_token_ids_in_response: inspecting response" + ); + + if let Some(nvext) = json_val.get("nvext") { + if let Some(completion_ids) = nvext.get("completion_token_ids").cloned() { + let n = completion_ids.as_array().map(|a| a.len()).unwrap_or(0); + tracing::info!( + n_completion_ids = n, + "rl_promote: copying completion_token_ids to choices[].token_ids" + ); + if let Some(choices) = json_val.get_mut("choices").and_then(|c| c.as_array_mut()) { + for choice in choices.iter_mut() { + if let Some(obj) = choice.as_object_mut() { + obj.insert("token_ids".to_string(), completion_ids.clone()); + } + } + } + } + } +} + +/// `GET /v1/rl/health` — lightweight health check for Prime-RL admin client. +/// +/// Prime-RL's `check_health()` calls `GET /health` on the admin client. When +/// `admin_base_url = ["http://dynamo:8000/v1/rl"]` the request arrives here. +/// Returns 200 OK if the frontend process is running (no deep probe needed — +/// the frontend's own `/health` endpoint handles that separately). +async fn rl_health() -> impl IntoResponse { + (StatusCode::OK, Json(serde_json::json!({"status": "ok"}))) +} + +/// Create an Axum [`Router`] for the RL admin endpoints at `/v1/rl/*`. +/// +/// Worker system URLs are read from the `DYN_RL_WORKER_SYSTEM_URLS` environment +/// variable (comma-separated, defaults to `http://localhost:8081`). +/// +/// Exposed only when `DYN_ENABLE_RL=true` or `HttpServiceConfig.enable_rl` is set. +/// +/// Prime-RL usage: set `admin_base_url = ["http://dynamo-frontend:8000/v1/rl"]` +/// in the orchestrator config. Prime-RL strips the trailing `/v1` suffix only +/// if present, so `/v1/rl` is preserved and all routes resolve correctly. +pub fn rl_router() -> (Vec, Router) { + let rl_state = Arc::new(RlState::from_env()); + let docs = vec![ + RouteDoc::new(axum::http::Method::GET, "/v1/rl/health"), + RouteDoc::new(axum::http::Method::GET, "/v1/rl/ready"), + RouteDoc::new(axum::http::Method::POST, "/v1/rl/pause"), + RouteDoc::new(axum::http::Method::POST, "/v1/rl/resume"), + RouteDoc::new(axum::http::Method::POST, "/v1/rl/update_weights"), + RouteDoc::new(axum::http::Method::POST, "/v1/rl/load_lora_adapter"), + RouteDoc::new(axum::http::Method::POST, "/v1/rl/unload_lora_adapter"), + RouteDoc::new(axum::http::Method::GET, "/v1/rl/weight_version"), + ]; + let router = Router::new() + .route("/v1/rl/health", get(rl_health)) + .route("/v1/rl/ready", get(rl_ready)) + .route("/v1/rl/pause", post(rl_pause)) + .route("/v1/rl/resume", post(rl_resume)) + .route("/v1/rl/update_weights", post(rl_update_weights)) + .route("/v1/rl/load_lora_adapter", post(rl_load_lora_adapter)) + .route("/v1/rl/unload_lora_adapter", post(rl_unload_lora_adapter)) + .route("/v1/rl/weight_version", get(rl_weight_version)) + .layer(middleware::from_fn(smart_json_error_middleware)) + .with_state(rl_state); + (docs, router) +} + #[cfg(test)] mod tests { diff --git a/lib/llm/src/http/service/service_v2.rs b/lib/llm/src/http/service/service_v2.rs index 1eeaefdff912..fe5d8ea6ba13 100644 --- a/lib/llm/src/http/service/service_v2.rs +++ b/lib/llm/src/http/service/service_v2.rs @@ -246,6 +246,12 @@ pub struct HttpServiceConfig { #[builder(default = "false")] enable_anthropic_endpoints: bool, + /// When true, expose the RL admin routes at `/v1/rl/*` (pause, resume, + /// update_weights, weight_version, ready). Worker system URLs are read + /// from `DYN_RL_WORKER_SYSTEM_URLS` (comma-separated, default `http://localhost:8081`). + #[builder(default = "false")] + enable_rl: bool, + #[builder(default = "None")] request_template: Option, @@ -518,7 +524,7 @@ impl HttpServiceConfigBuilder { }; // System routes (health, metrics, models) — debug-level spans - let system_routes = vec![ + let mut system_routes = vec![ metrics::router( registry, var(HTTP_SVC_METRICS_PATH_ENV).ok(), @@ -532,10 +538,16 @@ impl HttpServiceConfigBuilder { } else { super::openai::list_models_router(state.clone(), var(HTTP_SVC_MODELS_PATH_ENV).ok()) }, + super::openai::tokenization_router(state.clone()), super::health::health_check_router(state.clone(), var(HTTP_SVC_HEALTH_PATH_ENV).ok()), super::health::live_check_router(state.clone(), var(HTTP_SVC_LIVE_PATH_ENV).ok()), super::busy_threshold::busy_threshold_router(state.clone(), None), ]; + // RL admin routes: enabled when builder flag is set OR when DYN_ENABLE_RL env var is truthy. + if config.enable_rl || env_is_truthy("DYN_ENABLE_RL") { + tracing::info!("RL admin routes enabled at /v1/rl/*"); + system_routes.push(super::openai::rl_router()); + } let mut system_router = axum::Router::new(); for (route_docs, route) in system_routes { system_router = system_router.merge(route); @@ -600,6 +612,15 @@ impl HttpServiceConfigBuilder { request_template.clone(), var(HTTP_SVC_CHAT_PATH_ENV).ok(), ); + // RL TITO (Token-In / Token-Out) endpoint -- mounted alongside chat completions. + // Accepts Prime-RL's `tokens` field, translates to nvext.token_data, and delegates + // to the standard chat completions pipeline. Eliminates the Python rl-admin proxy. + let (tito_docs, tito_route) = super::openai::chat_completions_tokens_router( + state.clone(), + request_template.clone(), + None, + ); + let (cmpl_docs, cmpl_route) = super::openai::completions_router(state.clone(), var(HTTP_SVC_CMP_PATH_ENV).ok()); let (embed_docs, embed_route) = @@ -612,8 +633,13 @@ impl HttpServiceConfigBuilder { request_template.clone(), var(HTTP_SVC_RESPONSES_PATH_ENV).ok(), ); + // Merge TITO route and docs into the chat route (shares enable/disable flag) + let chat_route = chat_route.merge(tito_route); + let mut combined_chat_docs = chat_docs; + combined_chat_docs.extend(tito_docs); + let mut endpoint_routes = HashMap::new(); - endpoint_routes.insert(EndpointType::Chat, (chat_docs, chat_route)); + endpoint_routes.insert(EndpointType::Chat, (combined_chat_docs, chat_route)); endpoint_routes.insert(EndpointType::Completion, (cmpl_docs, cmpl_route)); endpoint_routes.insert(EndpointType::Embedding, (embed_docs, embed_route)); endpoint_routes.insert(EndpointType::Images, (images_docs, images_route)); diff --git a/lib/llm/src/preprocessor.rs b/lib/llm/src/preprocessor.rs index 296ec7e5fb0b..4d766f381efd 100644 --- a/lib/llm/src/preprocessor.rs +++ b/lib/llm/src/preprocessor.rs @@ -572,23 +572,23 @@ impl OpenAIPreprocessor { let token_data = request.nvext().and_then(|ext| ext.token_data.as_ref()); - let (tokens_vec, skip_token_annotation) = if has_backend_instance_id { - if let Some(tokens) = token_data { - tracing::trace!( - "Using provided tokens from EPP: {} ids", - tokens.len() - ); - // need ownership for the builder, so clone. - (tokens.clone(), true) - } else { - tracing::warn!( - "backend_instance_id provided but no token_data; tokenizing prompt" - ); - let encoding = self.encode_with_timing(&prompt, tracker)?; - (encoding.token_ids().to_vec(), false) - } + // Use token_data when provided (TITO / EPP / RL), + // regardless of backend_instance_id. + let (tokens_vec, skip_token_annotation) = if let Some(tokens) = + token_data + { + tracing::trace!( + "Using provided token_data: {} ids (backend_instance_id={has_backend_instance_id})", + tokens.len() + ); + (tokens.clone(), has_backend_instance_id) + } else if has_backend_instance_id { + tracing::warn!( + "backend_instance_id provided but no token_data; tokenizing prompt" + ); + let encoding = self.encode_with_timing(&prompt, tracker)?; + (encoding.token_ids().to_vec(), false) } else { - // No backend_instance_id provided, continue the normal flow. let encoding = self.encode_with_timing(&prompt, tracker)?; (encoding.token_ids().to_vec(), false) }; diff --git a/lib/llm/src/protocols/anthropic/types.rs b/lib/llm/src/protocols/anthropic/types.rs index 5214ee10b0a7..4104191a61ae 100644 --- a/lib/llm/src/protocols/anthropic/types.rs +++ b/lib/llm/src/protocols/anthropic/types.rs @@ -141,6 +141,8 @@ impl TryFrom for NvCreateChatCompletionRequest { }, media_io_kwargs: None, unsupported_fields: Default::default(), + tokens: None, + return_token_ids: None, }) } } diff --git a/lib/llm/src/protocols/openai.rs b/lib/llm/src/protocols/openai.rs index 42ef621f8797..4d022ac01c83 100644 --- a/lib/llm/src/protocols/openai.rs +++ b/lib/llm/src/protocols/openai.rs @@ -20,6 +20,7 @@ pub mod images; pub mod models; pub mod nvext; pub mod responses; +pub mod tokenization; pub mod tools; pub mod validate; pub mod videos; @@ -90,6 +91,10 @@ pub(crate) trait OpenAIOutputOptionsProvider { fn get_skip_special_tokens(&self) -> Option; fn get_formatted_prompt(&self) -> Option; + + fn get_return_tokens_as_token_ids(&self) -> Option { + None + } } impl SamplingOptionsProvider for T { @@ -203,7 +208,6 @@ impl OutputOptionsProvider for T { let prompt_logprobs = self.get_prompt_logprobs(); let skip_special_tokens = self.get_skip_special_tokens(); let formatted_prompt = self.get_formatted_prompt(); - Ok(common::OutputOptions { logprobs, prompt_logprobs, diff --git a/lib/llm/src/protocols/openai/chat_completions.rs b/lib/llm/src/protocols/openai/chat_completions.rs index 8a77038d5834..6e200e53bd67 100644 --- a/lib/llm/src/protocols/openai/chat_completions.rs +++ b/lib/llm/src/protocols/openai/chat_completions.rs @@ -59,6 +59,18 @@ pub struct NvCreateChatCompletionRequest { #[serde(default, skip_serializing_if = "Option::is_none")] pub media_io_kwargs: Option, + /// RL: Pre-tokenized prompt tokens from Prime-RL's TITO interface. + /// On the standard `/v1/chat/completions` endpoint this field is accepted but ignored + /// (use `/v1/chat/completions/tokens` for TITO mode where tokens are authoritative). + /// Accepting it here avoids 400 errors when Prime-RL sends it without the rl-admin proxy. + #[serde(default, skip_serializing)] + pub tokens: Option>, + + /// RL: Prime-RL requests token IDs in the response via this field. + /// Accepted but ignored on standard chat completions (use `nvext.extra_fields` instead). + #[serde(default, skip_serializing)] + pub return_token_ids: Option, + /// Catch-all for unsupported fields - checked during validation #[serde(flatten, default, skip_serializing)] pub unsupported_fields: std::collections::HashMap, @@ -72,6 +84,10 @@ pub struct NvCreateChatCompletionResponse { pub inner: dynamo_protocols::types::CreateChatCompletionResponse, #[serde(skip_serializing_if = "Option::is_none")] pub nvext: Option, + /// RL: Prompt token IDs for Prime-RL/verifiers alignment. + /// Populated when `DYN_ENABLE_RL=true` or `return_token_ids=true`. + #[serde(skip_serializing_if = "Option::is_none")] + pub prompt_token_ids: Option>, } /// A response structure for streamed chat completions, embedding OpenAI's diff --git a/lib/llm/src/protocols/openai/chat_completions/aggregator.rs b/lib/llm/src/protocols/openai/chat_completions/aggregator.rs index 2335de075037..43408efbb1a9 100644 --- a/lib/llm/src/protocols/openai/chat_completions/aggregator.rs +++ b/lib/llm/src/protocols/openai/chat_completions/aggregator.rs @@ -311,6 +311,7 @@ impl DeltaAggregator { service_tier: aggregator.service_tier, }, nvext: aggregator.nvext, + prompt_token_ids: None, }; Ok(response) diff --git a/lib/llm/src/protocols/openai/chat_completions/delta.rs b/lib/llm/src/protocols/openai/chat_completions/delta.rs index 8bf31756bfa4..1ab5cbf47a44 100644 --- a/lib/llm/src/protocols/openai/chat_completions/delta.rs +++ b/lib/llm/src/protocols/openai/chat_completions/delta.rs @@ -66,6 +66,15 @@ impl NvCreateChatCompletionRequest { }) .unwrap_or(false); + let enable_completion_token_ids = self + .nvext() + .map(|nv| { + nv.extra_fields + .as_ref() + .is_some_and(|fields| fields.iter().any(|f| f == "completion_token_ids")) + }) + .unwrap_or(false); + let options = DeltaGeneratorOptions { enable_usage: self .inner @@ -82,6 +91,7 @@ impl NvCreateChatCompletionRequest { enable_logprobs: self.inner.logprobs.unwrap_or(false) || self.inner.top_logprobs.unwrap_or(0) > 0, enable_tracking, + enable_completion_token_ids, runtime_config: ModelRuntimeConfig::default(), }; @@ -100,6 +110,10 @@ pub struct DeltaGeneratorOptions { pub enable_logprobs: bool, /// Determines whether request tracking (timing, KV hit rate) should be enabled. pub enable_tracking: bool, + /// Determines whether completion token IDs should be accumulated and returned + /// in `nvext.completion_token_ids`. Enabled when client requests + /// `extra_fields: ["completion_token_ids"]`. + pub enable_completion_token_ids: bool, pub runtime_config: ModelRuntimeConfig, } @@ -125,6 +139,10 @@ pub struct DeltaGenerator { options: DeltaGeneratorOptions, /// Optional request tracker for per-request metrics (shared with PreprocessedRequest). tracker: Option>, + /// Accumulated output token IDs across chunks. Only used when + /// `options.enable_completion_token_ids` is true. Emitted in `nvext.completion_token_ids` + /// on the final (finish_reason-bearing) chunk. + accumulated_completion_token_ids: Vec, } impl DeltaGenerator { @@ -172,6 +190,7 @@ impl DeltaGenerator { msg_counter: 0, options, tracker, + accumulated_completion_token_ids: Vec::new(), } } @@ -365,6 +384,12 @@ impl crate::protocols::openai::DeltaGeneratorExt for timing: timing_info, token_ids: token_ids.clone(), routed_experts, + completion_token_ids: None, }; if let Ok(nvext_json) = serde_json::to_value(&nvext_response) { diff --git a/lib/llm/src/protocols/openai/nvext.rs b/lib/llm/src/protocols/openai/nvext.rs index ea1f1e60f230..cdb935aafc01 100644 --- a/lib/llm/src/protocols/openai/nvext.rs +++ b/lib/llm/src/protocols/openai/nvext.rs @@ -114,6 +114,12 @@ pub struct NvExtResponse { /// Routed expert capture payload (SGLang-specific) #[serde(skip_serializing_if = "Option::is_none")] pub routed_experts: Option, + + /// Output token IDs generated by the engine. + /// Populated when client requests `extra_fields: ["completion_token_ids"]`. + /// For RL: len(completion_token_ids) == len(logprobs.content) is a hard invariant. + #[serde(skip_serializing_if = "Option::is_none")] + pub completion_token_ids: Option>, } /// NVIDIA LLM extensions to the OpenAI API diff --git a/lib/llm/src/protocols/openai/responses/mod.rs b/lib/llm/src/protocols/openai/responses/mod.rs index 5750c66ee985..bcee0384c29e 100644 --- a/lib/llm/src/protocols/openai/responses/mod.rs +++ b/lib/llm/src/protocols/openai/responses/mod.rs @@ -533,6 +533,8 @@ impl TryFrom for NvCreateChatCompletionRequest { chat_template_args: None, media_io_kwargs: None, unsupported_fields: Default::default(), + tokens: None, + return_token_ids: None, }) } } diff --git a/lib/llm/src/protocols/openai/tokenization.rs b/lib/llm/src/protocols/openai/tokenization.rs new file mode 100644 index 000000000000..95559684ad89 --- /dev/null +++ b/lib/llm/src/protocols/openai/tokenization.rs @@ -0,0 +1,124 @@ +// SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +use std::collections::HashMap; + +use serde::{Deserialize, Serialize}; + +use crate::preprocessor::media::MediaDecoder; +use crate::types::TokenIdType; + +fn default_true() -> bool { + true +} + +fn default_false() -> bool { + false +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct TokenizeCompletionRequest { + #[serde(default, skip_serializing_if = "Option::is_none")] + pub model: Option, + pub prompt: String, + #[serde(default = "default_true")] + pub add_special_tokens: bool, + #[serde(default = "default_false")] + pub return_token_strs: bool, +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct TokenizeChatRequest { + #[serde(default, skip_serializing_if = "Option::is_none")] + pub model: Option, + pub messages: Vec, + #[serde(default = "default_true")] + pub add_generation_prompt: bool, + #[serde(default = "default_false")] + pub return_token_strs: bool, + #[serde(default = "default_false")] + pub continue_final_message: bool, + #[serde(default = "default_false")] + pub add_special_tokens: bool, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub chat_template: Option, + #[serde( + default, + skip_serializing_if = "Option::is_none", + alias = "chat_template_args" + )] + pub chat_template_kwargs: Option>, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub media_io_kwargs: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub mm_processor_kwargs: Option>, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub tools: Option>, +} + +impl TokenizeChatRequest { + pub fn validate(&self) -> Result<(), String> { + if self.continue_final_message && self.add_generation_prompt { + return Err( + "Cannot set both `continue_final_message` and `add_generation_prompt` to True." + .to_string(), + ); + } + + Ok(()) + } + + pub fn merged_chat_template_kwargs(&self) -> HashMap { + let mut kwargs = self.chat_template_kwargs.clone().unwrap_or_default(); + kwargs.insert( + "add_generation_prompt".to_string(), + serde_json::Value::Bool(self.add_generation_prompt), + ); + kwargs.insert( + "continue_final_message".to_string(), + serde_json::Value::Bool(self.continue_final_message), + ); + kwargs + } +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +#[serde(untagged)] +#[allow(clippy::large_enum_variant)] +pub enum TokenizeRequest { + Completion(TokenizeCompletionRequest), + Chat(TokenizeChatRequest), +} + +impl TokenizeRequest { + pub fn model(&self) -> Option<&str> { + match self { + Self::Completion(request) => request.model.as_deref(), + Self::Chat(request) => request.model.as_deref(), + } + } +} + +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)] +pub struct TokenizeResponse { + pub count: usize, + pub max_model_len: u32, + pub tokens: Vec, + #[serde(skip_serializing_if = "Option::is_none")] + pub token_strs: Option>, +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct DetokenizeRequest { + #[serde(default, skip_serializing_if = "Option::is_none")] + pub model: Option, + pub tokens: Vec, +} + +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)] +pub struct DetokenizeResponse { + pub prompt: String, +} diff --git a/lib/llm/src/protocols/openai/validate.rs b/lib/llm/src/protocols/openai/validate.rs index 237e84bc75be..fff6ea8edd65 100644 --- a/lib/llm/src/protocols/openai/validate.rs +++ b/lib/llm/src/protocols/openai/validate.rs @@ -97,16 +97,24 @@ pub const MAX_REPETITION_PENALTY: f32 = 2.0; // Shared Fields // +/// Fields that Prime-RL / verifiers may send as extra_body hints which Dynamo +/// does not implement but should not reject with a 400. They are silently +/// accepted and ignored so the RL client stack is forward-compatible. +const PASSTHROUGH_EXTRA_FIELDS: &[&str] = &[ + "cache_salt", // KV prefix-cache isolation hint from prime-rl orchestrator +]; + /// Validates that no unsupported fields are present in the request pub fn validate_no_unsupported_fields( unsupported_fields: &std::collections::HashMap, ) -> Result<(), anyhow::Error> { - if !unsupported_fields.is_empty() { - let fields: Vec<_> = unsupported_fields - .keys() - .map(|s| format!("`{}`", s)) - .collect(); - anyhow::bail!("Unsupported parameter(s): {}", fields.join(", ")); + let unknown: Vec<_> = unsupported_fields + .keys() + .filter(|k| !PASSTHROUGH_EXTRA_FIELDS.contains(&k.as_str())) + .map(|s| format!("`{}`", s)) + .collect(); + if !unknown.is_empty() { + anyhow::bail!("Unsupported parameter(s): {}", unknown.join(", ")); } Ok(()) } diff --git a/lib/llm/src/tokenizers.rs b/lib/llm/src/tokenizers.rs index 561642965d09..bd99650bf119 100644 --- a/lib/llm/src/tokenizers.rs +++ b/lib/llm/src/tokenizers.rs @@ -61,6 +61,14 @@ pub mod traits { pub trait Encoder: Send + Sync { fn encode(&self, input: &str) -> Result; fn encode_batch(&self, inputs: &[&str]) -> Result>; + + fn encode_with_special_tokens( + &self, + input: &str, + _add_special_tokens: bool, + ) -> Result { + self.encode(input) + } } /// Result of decoding token IDs to text. @@ -126,8 +134,17 @@ pub mod traits { } pub trait Tokenizer: Encoder + Decoder { - // fn get_vocab_size(&self) -> usize; - // fn make_unique_clone(&self) -> Box; + fn convert_ids_to_tokens(&self, token_ids: &[TokenIdType]) -> Result> { + // Decoder::decode returns DecodeResult (Complete/Partial); the existing + // `impl From for String` unwraps to the inner string. + token_ids + .iter() + .map(|id| { + self.decode(std::slice::from_ref(id), false) + .map(String::from) + }) + .collect() + } } } @@ -148,6 +165,18 @@ impl Tokenizer { Ok(Tokenizer(create_tokenizer_from_file(file_path)?)) } + pub fn encode_with_special_tokens( + &self, + input: &str, + add_special_tokens: bool, + ) -> Result { + self.0.encode_with_special_tokens(input, add_special_tokens) + } + + pub fn convert_ids_to_tokens(&self, token_ids: &[TokenIdType]) -> Result> { + self.0.convert_ids_to_tokens(token_ids) + } + /// Create a stateful sequence object for decoding token_ids into text pub fn decode_stream( &self, diff --git a/lib/llm/src/tokenizers/fastokens.rs b/lib/llm/src/tokenizers/fastokens.rs index 25d8a7d83f11..5bdca6deec0f 100644 --- a/lib/llm/src/tokenizers/fastokens.rs +++ b/lib/llm/src/tokenizers/fastokens.rs @@ -39,16 +39,28 @@ impl FastTokenizer { impl Encoder for FastTokenizer { fn encode(&self, input: &str) -> Result { + self.encode_with_special_tokens(input, false) + } + + fn encode_batch(&self, inputs: &[&str]) -> Result> { + inputs.par_iter().map(|input| self.encode(input)).collect() + } + + fn encode_with_special_tokens( + &self, + input: &str, + add_special_tokens: bool, + ) -> Result { + if add_special_tokens { + return self.hf_decoder.encode_with_special_tokens(input, true); + } + let ids = self .fast_encoder .encode(input) .map_err(|e| Error::msg(format!("Fastokens encode error: {e}")))?; Ok(Encoding::Sp(ids)) } - - fn encode_batch(&self, inputs: &[&str]) -> Result> { - inputs.par_iter().map(|input| self.encode(input)).collect() - } } impl Decoder for FastTokenizer { @@ -57,7 +69,11 @@ impl Decoder for FastTokenizer { } } -impl Tokenizer for FastTokenizer {} +impl Tokenizer for FastTokenizer { + fn convert_ids_to_tokens(&self, token_ids: &[TokenIdType]) -> Result> { + self.hf_decoder.convert_ids_to_tokens(token_ids) + } +} #[cfg(test)] mod tests { diff --git a/lib/llm/src/tokenizers/hf.rs b/lib/llm/src/tokenizers/hf.rs index 080a775719fe..68c720c3ea8c 100644 --- a/lib/llm/src/tokenizers/hf.rs +++ b/lib/llm/src/tokenizers/hf.rs @@ -27,19 +27,18 @@ impl HuggingFaceTokenizer { impl Encoder for HuggingFaceTokenizer { fn encode(&self, input: &str) -> Result { - // This self.tokenizer is the library - let encoding = self - .tokenizer - .encode(input, false) - .map_err(|err| Error::msg(format!("Error tokenizing input: {err}")))?; - - Ok(Encoding::Hf(Box::new(encoding))) + // Use add_special_tokens=true to match TikTokenTokenizer::encode() behaviour. + // Both backends must agree on whether BOS/EOS are included so that callers + // (e.g. /v1/tokenize, rl_tokenize_prompt) get consistent token counts + // regardless of which backend is active. Callers that explicitly need no + // special tokens should call encode_with_special_tokens(input, false) directly. + self.encode_with_special_tokens(input, true) } fn encode_batch(&self, inputs: &[&str]) -> Result> { let hf_encodings = self .tokenizer - .encode_batch(inputs.to_vec(), false) + .encode_batch(inputs.to_vec(), true) // true to match encode() above .map_err(|err| Error::msg(format!("Error batch tokenizing input: {err}")))?; let encodings = hf_encodings @@ -49,6 +48,20 @@ impl Encoder for HuggingFaceTokenizer { Ok(encodings) } + + fn encode_with_special_tokens( + &self, + input: &str, + add_special_tokens: bool, + ) -> Result { + // This self.tokenizer is the library + let encoding = self + .tokenizer + .encode(input, add_special_tokens) + .map_err(|err| Error::msg(format!("Error tokenizing input: {err}")))?; + + Ok(Encoding::Hf(Box::new(encoding))) + } } impl Decoder for HuggingFaceTokenizer { @@ -63,7 +76,14 @@ impl Decoder for HuggingFaceTokenizer { } } -impl Tokenizer for HuggingFaceTokenizer {} +impl Tokenizer for HuggingFaceTokenizer { + fn convert_ids_to_tokens(&self, token_ids: &[TokenIdType]) -> Result> { + Ok(token_ids + .iter() + .map(|&id| self.tokenizer.id_to_token(id).unwrap_or_default()) + .collect()) + } +} impl From for HuggingFaceTokenizer { fn from(tokenizer: HfTokenizer) -> Self { diff --git a/lib/llm/src/tokenizers/tiktoken.rs b/lib/llm/src/tokenizers/tiktoken.rs index 9d311c971e35..20ef09d6c6a5 100644 --- a/lib/llm/src/tokenizers/tiktoken.rs +++ b/lib/llm/src/tokenizers/tiktoken.rs @@ -24,6 +24,8 @@ const KIMI_PATTERN: &str = r#"[\p{Han}]+|[^\r\n\p{L}\p{N}]?[\p{Lu}\p{Lt}\p{Lm}\p pub struct TikTokenTokenizer { bpe: CoreBPE, special_token_ids: HashSet, + decoder_tokens: FxHashMap>, + special_tokens_decoder: FxHashMap>, } impl TikTokenTokenizer { @@ -39,6 +41,14 @@ impl TikTokenTokenizer { special_tokens: FxHashMap, ) -> Result { let encoder = parse_tiktoken_file(path)?; + let decoder_tokens: FxHashMap> = encoder + .iter() + .map(|(bytes, &id)| (id, bytes.clone())) + .collect(); + let special_tokens_decoder: FxHashMap> = special_tokens + .iter() + .map(|(token, &id)| (id, token.as_bytes().to_vec())) + .collect(); let special_token_ids: HashSet = special_tokens.values().copied().collect(); let bpe = CoreBPE::new(encoder, special_tokens, pattern) @@ -47,6 +57,8 @@ impl TikTokenTokenizer { Ok(Self { bpe, special_token_ids, + decoder_tokens, + special_tokens_decoder, }) } @@ -62,9 +74,17 @@ impl TikTokenTokenizer { let pattern = detect_bpe_pattern(directory)?; let encoder = parse_tiktoken_file(path)?; + let decoder_tokens: FxHashMap> = encoder + .iter() + .map(|(bytes, &id)| (id, bytes.clone())) + .collect(); // Use max rank + 1 (not len) to avoid ID collisions with sparse/non-contiguous ranks let num_base_tokens = encoder.values().max().map_or(0, |&m| m + 1) as usize; let special_tokens = load_special_tokens(directory, num_base_tokens)?; + let special_tokens_decoder: FxHashMap> = special_tokens + .iter() + .map(|(token, &id)| (id, token.as_bytes().to_vec())) + .collect(); let special_token_ids: HashSet = special_tokens.values().copied().collect(); let bpe = CoreBPE::new(encoder, special_tokens, pattern) @@ -73,19 +93,33 @@ impl TikTokenTokenizer { Ok(Self { bpe, special_token_ids, + decoder_tokens, + special_tokens_decoder, }) } } impl Encoder for TikTokenTokenizer { fn encode(&self, input: &str) -> Result { - let token_ids: Vec = self.bpe.encode_with_special_tokens(input); - Ok(Encoding::Sp(token_ids)) + self.encode_with_special_tokens(input, true) } fn encode_batch(&self, inputs: &[&str]) -> Result> { inputs.par_iter().map(|input| self.encode(input)).collect() } + + fn encode_with_special_tokens( + &self, + input: &str, + add_special_tokens: bool, + ) -> Result { + let token_ids: Vec = if add_special_tokens { + self.bpe.encode_with_special_tokens(input) + } else { + self.bpe.encode_ordinary(input) + }; + Ok(Encoding::Sp(token_ids)) + } } impl Decoder for TikTokenTokenizer { @@ -119,7 +153,20 @@ impl Decoder for TikTokenTokenizer { } } -impl Tokenizer for TikTokenTokenizer {} +impl Tokenizer for TikTokenTokenizer { + fn convert_ids_to_tokens(&self, token_ids: &[TokenIdType]) -> Result> { + Ok(token_ids + .iter() + .map(|id| { + self.decoder_tokens + .get(id) + .or_else(|| self.special_tokens_decoder.get(id)) + .map(|bytes| String::from_utf8_lossy(bytes).into_owned()) + .unwrap_or_default() + }) + .collect()) + } +} /// Parse a tiktoken model file (base64-encoded token + rank per line). fn parse_tiktoken_file(path: &str) -> Result, u32>> {