diff --git a/Cargo.lock b/Cargo.lock index e1859a858ceb..9de7c2e8485b 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2459,6 +2459,7 @@ dependencies = [ "dynamo-mocker", "dynamo-parsers", "dynamo-protocols", + "dynamo-rl", "dynamo-runtime", "dynamo-tokenizers", "dynamo-tokens", @@ -2637,6 +2638,20 @@ dependencies = [ "uuid", ] +[[package]] +name = "dynamo-rl" +version = "1.2.0" +dependencies = [ + "anyhow", + "axum 0.8.4", + "dynamo-runtime", + "futures", + "serde", + "serde_json", + "tokio", + "tracing", +] + [[package]] name = "dynamo-runtime" version = "1.2.0" diff --git a/Cargo.toml b/Cargo.toml index 03f0aace838a..7c749cd866dc 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -4,6 +4,7 @@ [workspace] members = [ "lib/llm", + "lib/rl", "lib/runtime", "lib/config", "lib/tokenizers", @@ -41,6 +42,7 @@ keywords = ["llm", "genai", "inference", "nvidia", "distributed"] # Local crates dynamo-runtime = { path = "lib/runtime", version = "1.2.0" } dynamo-llm = { path = "lib/llm", version = "1.2.0" } +dynamo-rl = { path = "lib/rl", version = "1.2.0" } dynamo-config = { path = "lib/config", version = "1.2.0" } dynamo-tokenizers = { path = "lib/tokenizers", version = "1.2.0" } dynamo-tokens = { path = "lib/tokens", version = "1.2.0" } diff --git a/components/src/dynamo/vllm/backend_args.py b/components/src/dynamo/vllm/backend_args.py index b44d2adeef79..d40428e147a2 100644 --- a/components/src/dynamo/vllm/backend_args.py +++ b/components/src/dynamo/vllm/backend_args.py @@ -102,6 +102,27 @@ def add_arguments(self, parser) -> None: default=False, help="Enable multimodal processing. If not set, none of the multimodal components can be used.", ) + # Mirror SGLang's `--enable-rl` (sglang/backend_args.py:109) so both + # backends accept the same CLI / env var. The RL request-plane routes + # (pause_generation, update_weights_*, etc.) are unconditionally + # registered on the vLLM backend today, so this flag is currently a + # no-op gate that exists for parity + future opt-in gating. The + # canonical env var is `DYN_ENABLE_RL` (not `DYN_VLLM_ENABLE_RL`) + # because RL endpoints live on the frontend side too and they share + # this env var (see lib/llm/src/http/service/service_v2.rs). + add_negatable_bool_argument( + g, + flag_name="--enable-rl", + env_var="DYN_ENABLE_RL", + default=False, + help=( + "Enable RL training support. Mirrors --enable-rl on the SGLang " + "backend. RL admin routes (/v1/rl/engine pause_generation, " + "update_weights_from_disk, resume_generation, etc.) are " + "registered unconditionally today; this flag is reserved as the " + "future on/off gate." + ), + ) add_argument( g, flag_name="--mm-prompt-template", @@ -254,6 +275,10 @@ class DynamoVllmConfig(ConfigBase): multimodal_worker: bool multimodal_decode_worker: bool enable_multimodal: bool + # RL parity with SGLang. Reserved as the future on/off gate for RL routes; + # vLLM registers the routes unconditionally today, so this flag is a no-op + # signal that the worker is running in an RL deployment. + enable_rl: bool = False mm_prompt_template: str frontend_decoding: bool embedding_transfer_mode: Union[ diff --git a/components/src/dynamo/vllm/handlers.py b/components/src/dynamo/vllm/handlers.py index 04955398aca3..a77c54820fa4 100644 --- a/components/src/dynamo/vllm/handlers.py +++ b/components/src/dynamo/vllm/handlers.py @@ -14,7 +14,7 @@ from abc import ABC, abstractmethod from contextlib import asynccontextmanager from dataclasses import dataclass -from typing import Any, AsyncIterator, Dict, Final, Generic, Optional, TypeVar +from typing import Any, AsyncIterator, Dict, Final, Generic, NoReturn, Optional, TypeVar import torch from vllm.config import ModelConfig, VllmConfig @@ -289,6 +289,135 @@ def _compute_mm_uuids( return {"image": uuids} +# --------------------------------------------------------------------------- +# nvext.engine_data helpers (mirror of SGLang PR #8119) +# --------------------------------------------------------------------------- +# +# Clients (e.g. verifiers' OpenAIChatCompletionsTokenClient when targeting +# Dynamo) opt-in to engine-side metadata by setting +# `nvext.extra_fields: ["engine_data"]` on the request. When set, the +# response.nvext.engine_data payload from this backend includes the exact +# engine-emitted completion_token_ids and per-token logprobs (flat list of +# float indexed by sampled token) -- removing the need for the client to +# re-tokenize message.content to recover token IDs (the root cause of the +# 144x Mismatch-KL gap observed on Dynamo vs native vLLM RL runs). +# +# Symmetric with SGLang's `_nvext_extra_field_requested` helper added in +# `components/src/dynamo/sglang/request_handlers/llm/decode_handler.py`. + + +def _serialize_prompt_logprobs( + raw_prompt_logprobs: list, +) -> list: + """Convert vLLM's ``RequestOutput.prompt_logprobs`` into the dict shape + expected by Dynamo's Rust ``PromptLogprobEntry`` (serde deserialization). + + vLLM shape: ``list[dict[int, Logprob] | None]`` where ``Logprob`` has + ``.logprob``, ``.rank``, ``.decoded_token`` attributes. + + Output shape: ``list[dict[str, {"logprob": float, ...}] | None]`` — + matches ``LLMEngineOutput.prompt_logprobs`` so the Rust postprocessor + can surface it on ``NvExtResponse.prompt_logprobs``. + + NOTE: token_id keys are emitted as **strings** (not ints) so that + pythonize → serde → JSON survives the worker→frontend transport. JSON + object keys are required to be strings; ``HashMap`` on the Rust + side deserializes string keys via ``u32::from_str``. Emitting int keys + here causes the chunk to be silently dropped on pythonize → JSON, which + surfaces as ``"Stream ended before generation completed"`` on the + frontend (worker emits cleanly, frontend never sees ``complete_final``). + """ + result: list = [] + for entry in raw_prompt_logprobs: + if entry is None: + result.append(None) + else: + converted: Dict[str, Dict[str, Any]] = {} + for token_id, logprob_obj in entry.items(): + lp_dict: Dict[str, Any] = { + "logprob": float(logprob_obj.logprob), + } + rank = getattr(logprob_obj, "rank", None) + if rank is not None: + lp_dict["rank"] = int(rank) + decoded = getattr(logprob_obj, "decoded_token", None) + if decoded is not None: + lp_dict["decoded_token"] = decoded + converted[str(int(token_id))] = lp_dict + result.append(converted) + return result + + +def _nvext_extra_field_requested(request: Dict[str, Any], field: str) -> bool: + """Return True iff the request opted into the given nvext extra field. + + Mirrors the SGLang backend helper of the same name (PR #8119) so the + two backends honor `nvext.extra_fields=[...]` identically. + + Looks in two places (in order): + 1. `request["nvext"]["extra_fields"]` — raw OpenAI request shape + (used by SGLang's handler; also covers any direct test injection). + 2. `request["extra_args"]["nvext"]["extra_fields"]` — what the Rust + preprocessor stashes when building PreprocessedRequest from an + OpenAI request (PreprocessedRequest itself has no nvext field). + """ + for source in ( + request.get("nvext"), + (request.get("extra_args") or {}).get("nvext") + if isinstance(request.get("extra_args"), dict) + else None, + ): + if not isinstance(source, dict): + continue + extra_fields = source.get("extra_fields") + if isinstance(extra_fields, list) and field in extra_fields: + return True + return False + + +def _flatten_logprobs( + log_probs: Any, +) -> Optional[list[float]]: + """Coerce a per-token logprob sequence into a flat list[float]. + + The backend's per-chunk `log_probs` field is one float per emitted + token (the logprob the engine sampled). Some upstream paths wrap it + in dicts or nested lists; this helper accepts: + - list[float] -> returned as-is + - list[list[float]] -> flattened + - list[dict{logprob: ...}] -> .logprob extracted + - None -> None + + Any element that isn't coercible to float is dropped silently rather + than poisoning the entire payload. The trainer treats missing + per-token logprobs as off-policy correction = 0 for that token, so a + partial list is preferable to a hard failure. + """ + if log_probs is None: + return None + if not isinstance(log_probs, list): + return None + out: list[float] = [] + for item in log_probs: + if isinstance(item, (int, float)): + out.append(float(item)) + elif isinstance(item, list) and item: + head = item[0] + if isinstance(head, (int, float)): + out.append(float(head)) + elif isinstance(head, dict) and "logprob" in head: + try: + out.append(float(head["logprob"])) + except (TypeError, ValueError): + continue + elif isinstance(item, dict) and "logprob" in item: + try: + out.append(float(item["logprob"])) + except (TypeError, ValueError): + continue + return out or None + + # LoRAManager singleton - initialized lazily when DYN_LORA_ENABLED is set # None = not yet initialized, False = disabled/failed, LoRAManager = initialized _lora_manager = None @@ -327,12 +456,41 @@ def build_sampling_params( Args: request: The PreprocessedRequest dict with 'sampling_options', 'stop_conditions', and 'output_options' - default_sampling_params: Default sampling parameters to initialize with + default_sampling_params: Default sampling parameters from the model's + ``generation_config.json`` (vLLM ``ModelConfig.get_diff_sampling_param``). + Used for non-RL/chat clients that want the model's recommended + sampling defaults applied transparently. Returns: SamplingParams configured from the request + + RL/TITO parity note (rl-sdk-2): when the client opted into pre-tokenized + prompts via ``nvext.token_data``, this function **bypasses the + ``generation_config.json`` defaults** for sampling-distribution params + (top_p, top_k, min_p, temperature, repetition_penalty). RL clients want + vLLM's vanilla defaults so the request is bit-equivalent to a vLLM-native + ``/inference/v1/generate`` call. Stop-token defaults are still honored so + the model still terminates on EOS. Without this gate, Qwen3's + ``top_p=0.95, top_k=20`` ride along for every request and produce a ~16× + Mismatch-KL gap vs vLLM-native at non-greedy temperatures (see + plans/may-13/parity_test.md attribution probes). """ - sampling_params = SamplingParams(**default_sampling_params) + nvext_args = (request.get("extra_args") or {}).get("nvext") or {} + is_tito_request = bool(request.get("nvext") and request["nvext"].get("token_data")) + # The preprocessor strips nvext when it produces a PreprocessedRequest; + # tito mode is detectable from the presence of pre-set ``token_ids`` and + # absence of any prompt text. Use the explicit nvext.extra_fields signal + # from the Rust preprocessor pass-through (PASSTHROUGH_EXTRA_FIELDS path) + # — RL clients always include at least one of ``engine_data``, + # ``completion_token_ids``, ``prompt_logprobs`` in ``extra_fields``. + if not is_tito_request and isinstance(nvext_args.get("extra_fields"), list): + is_tito_request = bool(nvext_args["extra_fields"]) + + if is_tito_request: + # Strict vLLM defaults — no generation_config overlay. + sampling_params = SamplingParams() + else: + sampling_params = SamplingParams(**default_sampling_params) sampling_params.detokenize = False # Handle guided_decoding - convert to StructuredOutputsParams @@ -405,6 +563,10 @@ def build_sampling_params( f"Invalid prompt_logprobs value: {prompt_logprobs_value} (must be integer), ignoring" ) + skip_special_tokens_value = output_options.get("skip_special_tokens") + if skip_special_tokens_value is not None: + sampling_params.skip_special_tokens = bool(skip_special_tokens_value) + # If max_tokens wasn't provided (None or missing), compute a dynamic default provided_max_tokens = request.get("stop_conditions", {}).get("max_tokens", None) token_ids = request.get("token_ids", []) @@ -676,6 +838,24 @@ def __init__( # Store shutdown event for graceful shutdown monitoring self.shutdown_event = shutdown_event + # RL state + self._paused: bool = False + self._weight_version: str = "initial" + # Maps method name → async callable for the request-plane rl_dispatch endpoint. + # Populated by register_engine_routes() in worker_factory.py. + self._rl_routes: dict[str, Any] = {} + + def _shutdown_on_engine_dead(self, e: EngineDeadError) -> NoReturn: + """Handle EngineDeadError from RL admin handlers. + + Logs the error, shuts down the Dynamo runtime, and hard-exits so the + pod is restarted rather than silently failing RL operations. + """ + logger.error(f"[RL] EngineDeadError: {e}") + logger.warning("[RL] Initiating Dynamo Runtime shutdown.") + self.runtime.shutdown() + os._exit(1) + def init_embedding_loader( self, config: Config, encode_worker_client: Optional[Client] = None ) -> Optional[MultiModalEmbeddingLoader]: @@ -933,6 +1113,509 @@ async def stop_profile(self, body: dict) -> dict: logger.error(f"Failed to stop profiling: {e}") return {"status": "error", "message": str(e)} + # ------------------------------------------------------------------------- + # RL admin — request-plane dispatcher (served on dyn://..rl) + # ------------------------------------------------------------------------- + + async def rl_dispatch(self, request=None): + """Request-plane dispatcher for the worker's ``rl`` endpoint. + + The Dynamo frontend discovers live ``rl`` endpoint instances via the + discovery plane and fans out ``{method, kwargs}`` payloads here via + strict request-plane direct routing (NATS / TCP). + + Wire shape (inbound): ``{"method": str, "kwargs": dict}`` + + Special method names: + ``__describe__`` Returns ``{"registered_methods": [...]}`` so the + frontend's ``GET /v1/rl/engine`` can list what each + worker supports. + """ + if request is None: + yield {"status": "error", "message": "rl_dispatch: request required"} + return + + method = request.get("method") + kwargs = request.get("kwargs") or {} + + if not isinstance(method, str) or not method: + yield {"status": "error", "message": "rl_dispatch: missing 'method' (str)"} + return + + # Describe — list registered methods (no handler lookup needed) + if method == "__describe__": + yield { + "status": "ok", + "registered_methods": sorted(self._rl_routes), + } + return + + handler_fn = self._rl_routes.get(method) + if handler_fn is None: + yield { + "status": "error", + "message": ( + f"rl_dispatch: unknown method {method!r}; " + f"registered: {sorted(self._rl_routes)}" + ), + } + return + + try: + result = await handler_fn(kwargs) + yield result if result is not None else {"status": "ok"} + except Exception as e: + logger.exception(f"[RL] rl_dispatch method={method!r} failed: {e}") + yield {"status": "error", "method": method, "message": str(e)} + + # ------------------------------------------------------------------------- + # RL admin routes + # Registered at startup via register_engine_routes(); called by the frontend + # via POST /v1/rl/engine {"method": "", "kwargs": {...}}. + # ------------------------------------------------------------------------- + + async def liveness_probe(self, body: dict) -> dict: + """Engine event-loop liveness probe. + + Confirms the engine is responsive by performing a lightweight + check_health() round-trip to the EngineCore. A hung event loop or + wedged engine will time out at the frontend rather than returning a + stale ok. + """ + body = body or {} + try: + if hasattr(self.engine_client, "check_health"): + await self.engine_client.check_health() + else: + # Fallback: the collective_rpc round-trip itself is the liveness + # signal for older engines that don't expose check_health(). + await self.engine_client.collective_rpc( + "liveness_probe", kwargs={} + ) + return {"status": "ok", "alive": True} + except EngineDeadError as e: + self._shutdown_on_engine_dead(e) + except Exception as e: + logger.warning(f"[RL] liveness_probe failed: {e}") + return {"status": "error", "alive": False, "message": str(e)} + + async def pause_generation(self, body: dict) -> dict: + """Pause the engine before a weight update. + + Does NOT release GPU memory (no sleep) and does NOT unregister from + discovery. New requests queue; in-flight requests are handled according + to ``mode``. + + Body (all optional): + mode "keep" | "wait" | "abort" (default "keep") + clear_cache bool (default False) + """ + body = body or {} + mode = body.get("mode", "keep") + clear_cache = bool(body.get("clear_cache", False)) + if mode not in ("keep", "wait", "abort"): + return { + "status": "error", + "message": f"Invalid mode '{mode}'; expected keep|wait|abort", + } + try: + await self.engine_client.pause_generation( + mode=mode, clear_cache=clear_cache + ) + self._paused = True + logger.info(f"[RL] Engine paused (mode={mode}, clear_cache={clear_cache})") + return { + "status": "ok", + "message": "Engine paused", + "mode": mode, + "clear_cache": clear_cache, + } + except EngineDeadError as e: + self._shutdown_on_engine_dead(e) + 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() + self._paused = False + logger.info("[RL] Engine resumed") + return {"status": "ok", "message": "Engine resumed"} + except EngineDeadError as e: + self._shutdown_on_engine_dead(e) + 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 the prefix / KV cache.""" + body = body or {} + try: + await self.engine_client.reset_prefix_cache() + logger.debug("[RL] Prefix cache flushed") + return {"status": "ok", "message": "Cache flushed"} + except EngineDeadError as e: + self._shutdown_on_engine_dead(e) + except Exception as e: + logger.error(f"[RL] Failed to flush cache: {e}") + return {"status": "error", "message": str(e)} + + async def abort_request(self, body: dict) -> dict: + """Abort a single in-flight request by its request_id. + + Body: {"request_id": str} + """ + body = body or {} + request_id = body.get("request_id") + if not request_id: + return {"status": "error", "message": "Missing 'request_id' in body"} + try: + await self.engine_client.abort(request_id) + logger.debug(f"[RL] Aborted request {request_id}") + return {"status": "ok", "request_id": request_id} + except EngineDeadError as e: + self._shutdown_on_engine_dead(e) + except Exception as e: + logger.error(f"[RL] Failed to abort request {request_id}: {e}") + return {"status": "error", "message": str(e)} + + async def get_weight_version(self, body: dict) -> dict: + """Return the current weight version tag.""" + return {"status": "ok", "version": getattr(self, "_weight_version", "initial")} + + async def update_weights_from_disk(self, body: dict) -> dict: + """Load weights from a shared filesystem checkpoint. + + Body: + model_path str path to the safetensors checkpoint directory + weight_version str version tag to record (default "unknown") + engine_rpc str collective_rpc target (default "reload_weights") + + When engine_rpc is "reload_weights" (vLLM built-in), the kwargs key is + "weights_path". For the FileSystemWeightUpdateWorker extension the target + is "update_weights_from_path" and the kwargs key is "weight_path". + """ + body = body or {} + if not getattr(self, "_paused", False): + return { + "status": "error", + "message": ( + "Worker must be paused via pause_generation() before updating " + "weights. Call pause_generation() first, then update, then " + "resume_generation()." + ), + } + path = body.get("model_path") + if not path: + return {"status": "error", "message": "Missing 'model_path' in body"} + version = body.get("weight_version", "unknown") + rpc = body.get("engine_rpc", "reload_weights") + kwargs = {"weights_path": path} if rpc == "reload_weights" else {"weight_path": path} + try: + await self.engine_client.collective_rpc(rpc, kwargs=kwargs) + self._weight_version = version + logger.info(f"[RL] Weights loaded from {path} (version={version}, rpc={rpc})") + return {"status": "ok", "version": version} + except EngineDeadError as e: + self._shutdown_on_engine_dead(e) + except Exception as e: + logger.error(f"[RL] update_weights_from_disk failed: {e}") + return {"status": "error", "message": str(e)} + + async def update_weights_from_distributed(self, body: dict) -> dict: + """Receive weights via a distributed transport (e.g. NCCL). + + Requires init_weights_update_group to have been called first. + + Body: + weight_version str version tag to record (default "unknown") + engine_rpc str collective_rpc target (default "update_weights_from_path") + forwarded to the rpc as kwargs (e.g. weight_dir for + NCCLWeightUpdateWorker.update_weights_from_path) + """ + body = body or {} + if not getattr(self, "_paused", False): + return { + "status": "error", + "message": ( + "Worker must be paused via pause_generation() before updating " + "weights. Call pause_generation() first, then update, then " + "resume_generation()." + ), + } + version = body.get("weight_version", "unknown") + rpc = body.get("engine_rpc", "update_weights_from_path") + rpc_kwargs = {k: v for k, v in body.items() if k not in ("engine_rpc", "weight_version")} + try: + await self.engine_client.collective_rpc(rpc, kwargs=rpc_kwargs) + self._weight_version = version + logger.info(f"[RL] Weights received via distributed (version={version}, rpc={rpc})") + return {"status": "ok", "version": version} + except EngineDeadError as e: + self._shutdown_on_engine_dead(e) + except Exception as e: + logger.error(f"[RL] update_weights_from_distributed failed: {e}") + return {"status": "error", "message": str(e)} + + async def update_weights_from_tensor(self, body: dict) -> dict: + """Not implemented — in-process tensor transfer is not yet supported.""" + return { + "status": "error", + "message": "update_weights_from_tensor is not implemented", + } + + async def init_weights_update_group(self, body: dict) -> dict: + """Initialize the distributed weight-update communication group. + + Body: + master_address str + master_port int + rank_offset int + world_size int + timeout int seconds (default 180) + engine_rpc str collective_rpc target (default "init_broadcaster") + """ + body = body or {} + rpc = body.get("engine_rpc", "init_broadcaster") + kwargs = {k: v for k, v in body.items() if k != "engine_rpc"} + try: + await self.engine_client.collective_rpc(rpc, kwargs=kwargs) + logger.info(f"[RL] Weight update group initialized (rpc={rpc})") + return {"status": "ok", "message": "Weight update group initialized"} + except EngineDeadError as e: + self._shutdown_on_engine_dead(e) + except Exception as e: + logger.error(f"[RL] init_weights_update_group failed: {e}") + return {"status": "error", "message": str(e)} + + async def destroy_weights_update_group(self, body: dict) -> dict: + """Tear down the distributed weight-update communication group. + + Body: + engine_rpc str collective_rpc target (default "destroy_broadcaster") + """ + body = body or {} + rpc = body.get("engine_rpc", "destroy_broadcaster") + kwargs = {k: v for k, v in body.items() if k != "engine_rpc"} + try: + await self.engine_client.collective_rpc(rpc, kwargs=kwargs) + logger.info(f"[RL] Weight update group destroyed (rpc={rpc})") + return {"status": "ok", "message": "Weight update group destroyed"} + except EngineDeadError as e: + self._shutdown_on_engine_dead(e) + except Exception as e: + logger.error(f"[RL] destroy_weights_update_group failed: {e}") + return {"status": "error", "message": str(e)} + + async def load_lora_adapter(self, body: dict) -> dict: + """Load (or hot-swap) a LoRA adapter from a filesystem path. + + Body: {"lora_name": str, "lora_path": str} + + Hot-swaps if the adapter is already loaded: removes the old id first, + reloads, then resets the prefix cache so stale KV entries don't + contaminate new rollouts. Publishes a ModelDeploymentCard on first + load so the frontend can route requests with model=. + """ + 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 + + # Save old info but do NOT pop yet + old_info = self.loaded_loras.get(lora_name) if is_hot_swap else None + if is_hot_swap and old_info is not None: + old_id = old_info.id + try: + await self.engine_client.remove_lora(old_id) + except Exception as e: + logger.error( + f"[RL] remove_lora({lora_name}, id={old_id}) failed during hot-swap: {e}" + ) + return { + "status": "error", + "message": f"Failed to remove existing LoRA '{lora_name}' before hot-swap: {e}", + "lora_name": lora_name, + } + + try: + await self.engine_client.add_lora( + LoRARequest( + lora_name=lora_name, + lora_int_id=lora_id, + lora_path=lora_path, + ) + ) + except Exception as e: + # add_lora failed — rollback by re-adding old adapter + if is_hot_swap and old_info is not None: + try: + await self.engine_client.add_lora( + LoRARequest( + lora_name=lora_name, + lora_int_id=old_info.id, + lora_path=old_info.path, + ) + ) + logger.warning( + f"[RL] add_lora({lora_name}) failed; rolled back to old adapter (id={old_info.id}): {e}" + ) + except Exception as rollback_err: + self.loaded_loras.pop(lora_name, None) + logger.error( + f"[RL] add_lora({lora_name}) failed AND rollback failed: {rollback_err}" + ) + return { + "status": "error", + "message": f"Failed to add LoRA '{lora_name}': {e}", + "lora_name": lora_name, + "lora_id": lora_id, + } + + # Only now update loaded_loras (add_lora succeeded) + self.loaded_loras[lora_name] = LoRAInfo(id=lora_id, path=lora_path) + + if is_hot_swap: + try: + await self.engine_client.reset_prefix_cache() + except Exception as e: + logger.error(f"[RL] reset_prefix_cache after LoRA swap failed: {e}") + return { + "status": "error", + "message": ( + f"LoRA '{lora_name}' loaded but prefix cache reset failed; " + "worker is not safe to serve until next successful swap." + ), + "lora_name": lora_name, + "lora_id": lora_id, + } + + 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, + ) + except Exception as e: + logger.exception(f"[RL] Failed to publish LoRA '{lora_name}' MDC: {e}; rolling back") + try: + await self.engine_client.remove_lora(lora_id) + except Exception as rollback_err: + logger.error(f"[RL] Rollback remove_lora failed — adapter leaked: {rollback_err}") + self.loaded_loras.pop(lora_name, None) + return { + "status": "error", + "message": f"Failed to register LoRA '{lora_name}' in discovery: {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", + "lora_name": lora_name, + "lora_id": lora_id, + "hot_swap": is_hot_swap, + } + except EngineDeadError as e: + self._shutdown_on_engine_dead(e) + except Exception as e: + logger.exception(f"[RL] Failed to load LoRA adapter '{lora_name}': {e}") + return {"status": "error", "message": str(e)} + + async def unload_lora_adapter(self, body: dict) -> dict: + """Unload a LoRA adapter previously loaded via load_lora_adapter. + + Body: {"lora_name": str} + + Idempotent: unloading an 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 + lora_path = lora.path # save for potential rollback + await self.engine_client.remove_lora(lora_id) + del self.loaded_loras[lora_name] + + if self.generate_endpoint is not None: + try: + await unregister_model( + endpoint=self.generate_endpoint, + lora_name=lora_name, + ) + except Exception as e: + logger.error( + f"[RL] Failed to unregister LoRA '{lora_name}' MDC after engine removal: {e} — rolling back" + ) + # Rollback: re-add adapter to engine and restore loaded_loras + try: + 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) + except Exception as rollback_err: + logger.error( + f"[RL] Rollback add_lora({lora_name}) also failed — adapter gone from engine but MDC may still route here: {rollback_err}" + ) + return { + "status": "error", + "message": f"LoRA '{lora_name}' removal rolled back: discovery unregister failed ({e}). Retry unload_lora_adapter.", + "lora_name": lora_name, + "lora_id": lora_id, + } + + logger.info(f"[RL] LoRA adapter unloaded: name={lora_name} id={lora_id}") + return { + "status": "ok", + "lora_name": lora_name, + "lora_id": lora_id, + } + except EngineDeadError as e: + self._shutdown_on_engine_dead(e) + 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 @@ -2160,6 +2843,10 @@ async def generate_tokens( request_output=res, embedding_sequence_length=embedding_sequence_length, ) + if res.prompt_logprobs is not None: + out["prompt_logprobs"] = _serialize_prompt_logprobs( + res.prompt_logprobs + ) # Log completion with LoRA info (debug level to avoid log spam) self._log_with_lora_context( "Completed token generation for request {request_id}{lora_info}: " @@ -2356,6 +3043,12 @@ async def _generate_token_mode(self, request, context, request_id): yield error return + extra_args = request.get("extra_args") or {} + nvext_args = extra_args.get("nvext") or {} + cache_salt = nvext_args.get("cache_salt") + if cache_salt is not None and isinstance(prompt, dict): + prompt["cache_salt"] = cache_salt + # Build sampling params from request sampling_params = build_sampling_params( request, self.default_sampling_params, self.model_max_len @@ -2402,6 +3095,28 @@ async def _generate_token_mode(self, request, context, request_id): async with self._abort_monitor( context, request_id, abort_guard=abort_guard ): + # nvext.engine_data opt-in: if the client requested + # `nvext.extra_fields=["engine_data"]`, we accumulate + # per-chunk token_ids and logprobs and attach them to the + # FINAL chunk (the one carrying `finish_reason`). The Rust + # frontend's response builder (delta.rs) gates emission via + # `NvExtResponseFieldSelection.engine_data` so this payload + # only reaches clients that asked for it. + want_engine_data = _nvext_extra_field_requested( + request, "engine_data" + ) + # Prompt token IDs the engine actually saw. Either the + # pre-tokenized `nvext.token_data` (TITO) or whatever the + # preprocessor produced from messages (MITO). We echo them + # back in engine_data so the client doesn't have to re-derive + # them from a request it might no longer hold. + request_prompt_token_ids = ( + list(request.get("token_ids") or []) + if want_engine_data + else None + ) + accumulated_token_ids: list[int] = [] + accumulated_log_probs: list[float] = [] try: async for tok in self.generate_tokens( prompt, @@ -2421,6 +3136,35 @@ async def _generate_token_mode(self, request, context, request_id): tok["completion_usage"][ "prompt_tokens_details" ] = prefill_prompt_tokens_details + + if want_engine_data: + new_token_ids = tok.get("token_ids") + if isinstance(new_token_ids, list): + accumulated_token_ids.extend( + int(t) for t in new_token_ids + ) + flat_lp = _flatten_logprobs(tok.get("log_probs")) + if flat_lp: + accumulated_log_probs.extend(flat_lp) + if tok.get("finish_reason") is not None: + # Final chunk -- attach the cumulative + # engine_data payload. Schema mirrors PR #8119 + # (SGLang) so clients see a single shape. + engine_data: Dict[str, Any] = { + "completion_token_ids": list( + accumulated_token_ids + ), + "finished": True, + } + if accumulated_log_probs: + engine_data["completion_logprobs"] = list( + accumulated_log_probs + ) + if request_prompt_token_ids: + engine_data["prompt_token_ids"] = list( + request_prompt_token_ids + ) + tok["engine_data"] = engine_data yield tok except EngineDeadError as e: logger.error(f"vLLM EngineDeadError: {e}") diff --git a/components/src/dynamo/vllm/main.py b/components/src/dynamo/vllm/main.py index e089973bdb6e..0b3bad44a356 100644 --- a/components/src/dynamo/vllm/main.py +++ b/components/src/dynamo/vllm/main.py @@ -110,6 +110,23 @@ async def worker() -> None: if not config.served_model_name: config.served_model_name = config.engine_args.served_model_name = config.model + # rl-sdk-2 TITO parity: when running as an RL backend, default + # ``logprobs_mode`` to ``processed_logprobs`` so per-token logprobs are + # reported from the temperature-applied (post-softmax-with-T) distribution. + # vLLM defaults to ``raw_logprobs`` (pre-temperature), which produces + # systematically different values vs prime-rl's ``inference @`` entrypoint + # (which sets ``processed_logprobs`` at ``inference.py:566``) — leading to + # a ~17× Mismatch KL gap on T>0 RL workloads even though the engines + # otherwise agree byte-for-byte. Honor an explicit ``--logprobs-mode`` + # override if the caller set one; only inject the RL-friendly default + # when the field is still at the vLLM default ``raw_logprobs``. + if config.enable_rl and config.engine_args.logprobs_mode == "raw_logprobs": + config.engine_args.logprobs_mode = "processed_logprobs" + logger.info( + "Defaulting logprobs_mode=processed_logprobs (--enable-rl active); " + "override with --logprobs-mode=raw_logprobs to restore vLLM's default." + ) + # Download the model if necessary using modelexpress. # We want it on disk before we start vllm to avoid downloading from HuggingFace. # diff --git a/components/src/dynamo/vllm/tests/test_vllm_tito_parity.py b/components/src/dynamo/vllm/tests/test_vllm_tito_parity.py new file mode 100644 index 000000000000..eb61fc9c73ad --- /dev/null +++ b/components/src/dynamo/vllm/tests/test_vllm_tito_parity.py @@ -0,0 +1,188 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Unit tests for vLLM TITO (token-in-token-out) parity fixes: + + D1 – _serialize_prompt_logprobs converts vLLM's prompt_logprobs into + the dict shape expected by the Rust PromptLogprobEntry. + D2 – cache_salt is forwarded from extra_args["nvext"]["cache_salt"] + to the prompt dict so vLLM's input_processor picks it up. + D3 – skip_special_tokens from output_options is applied to + SamplingParams via build_sampling_params. +""" + +from __future__ import annotations + +import importlib.util +from types import SimpleNamespace +from unittest.mock import MagicMock + +import pytest + +pytestmark = [ + pytest.mark.pre_merge, + pytest.mark.vllm, + pytest.mark.gpu_0, + pytest.mark.skipif( + importlib.util.find_spec("vllm") is None, + reason="vllm not installed in this container", + ), +] + + +# --------------------------------------------------------------------------- +# D1 – _serialize_prompt_logprobs +# --------------------------------------------------------------------------- + + +class TestSerializePromptLogprobs: + """Validate _serialize_prompt_logprobs against various vLLM outputs.""" + + @staticmethod + def _import(): + from dynamo.vllm.handlers import _serialize_prompt_logprobs + + return _serialize_prompt_logprobs + + def test_none_entries_preserved(self): + fn = self._import() + raw = [None, None] + assert fn(raw) == [None, None] + + def test_single_token_entry(self): + fn = self._import() + logprob = SimpleNamespace(logprob=-1.5, rank=1, decoded_token="hello") + raw = [{42: logprob}] + result = fn(raw) + assert len(result) == 1 + entry = result[0] + assert 42 in entry + assert entry[42]["logprob"] == pytest.approx(-1.5) + assert entry[42]["rank"] == 1 + assert entry[42]["decoded_token"] == "hello" + + def test_mixed_none_and_entries(self): + fn = self._import() + lp1 = SimpleNamespace(logprob=-0.1, rank=1, decoded_token="a") + lp2 = SimpleNamespace(logprob=-2.3, rank=5, decoded_token="b") + raw = [None, {10: lp1, 20: lp2}, None] + result = fn(raw) + assert result[0] is None + assert result[2] is None + assert set(result[1].keys()) == {10, 20} + + def test_missing_optional_attributes(self): + """Logprob objects without rank/decoded_token should omit those keys.""" + fn = self._import() + logprob = SimpleNamespace(logprob=-3.0) + raw = [{7: logprob}] + result = fn(raw) + assert result[0][7]["logprob"] == pytest.approx(-3.0) + assert "rank" not in result[0][7] + assert "decoded_token" not in result[0][7] + + def test_empty_list(self): + fn = self._import() + assert fn([]) == [] + + def test_multiple_tokens_per_position(self): + fn = self._import() + lp_a = SimpleNamespace(logprob=-0.5, rank=1, decoded_token="x") + lp_b = SimpleNamespace(logprob=-1.2, rank=2, decoded_token="y") + lp_c = SimpleNamespace(logprob=-3.0, rank=3, decoded_token="z") + raw = [{100: lp_a, 200: lp_b, 300: lp_c}] + result = fn(raw) + assert len(result[0]) == 3 + + +# --------------------------------------------------------------------------- +# D2 – cache_salt forwarding +# --------------------------------------------------------------------------- + + +class TestCacheSaltWiring: + """Verify cache_salt is extracted from extra_args and placed on the prompt.""" + + @staticmethod + def _build_token_mode_request(cache_salt=None, token_ids=None): + """Build a minimal TITO request dict mirroring the Rust preprocessor.""" + req = { + "token_ids": token_ids or [1, 2, 3], + "sampling_options": {}, + "stop_conditions": {}, + "output_options": {}, + } + if cache_salt is not None: + req["extra_args"] = {"nvext": {"cache_salt": cache_salt}} + return req + + def test_cache_salt_attached_to_prompt(self): + """When extra_args.nvext.cache_salt is set, the prompt dict gets it.""" + from vllm.inputs import TokensPrompt + + req = self._build_token_mode_request(cache_salt="step_42") + extra_args = req.get("extra_args") or {} + nvext_args = extra_args.get("nvext") or {} + salt = nvext_args.get("cache_salt") + prompt = TokensPrompt(prompt_token_ids=req["token_ids"]) + if salt is not None and isinstance(prompt, dict): + prompt["cache_salt"] = salt + + assert prompt.get("cache_salt") == "step_42" + + def test_no_cache_salt_when_absent(self): + """When extra_args has no cache_salt, prompt should not gain the key.""" + from vllm.inputs import TokensPrompt + + req = self._build_token_mode_request() + extra_args = req.get("extra_args") or {} + nvext_args = extra_args.get("nvext") or {} + salt = nvext_args.get("cache_salt") + prompt = TokensPrompt(prompt_token_ids=req["token_ids"]) + if salt is not None and isinstance(prompt, dict): + prompt["cache_salt"] = salt + + assert "cache_salt" not in prompt + + +# --------------------------------------------------------------------------- +# D3 – skip_special_tokens in build_sampling_params +# --------------------------------------------------------------------------- + + +class TestSkipSpecialTokens: + """Verify skip_special_tokens from output_options flows to SamplingParams.""" + + @staticmethod + def _build(output_options=None): + from dynamo.vllm.handlers import build_sampling_params + + req = { + "token_ids": [1, 2, 3], + "sampling_options": {}, + "stop_conditions": {}, + "output_options": output_options or {}, + } + return build_sampling_params(req, {}) + + def test_skip_special_tokens_true(self): + sp = self._build(output_options={"skip_special_tokens": True}) + assert sp.skip_special_tokens is True + + def test_skip_special_tokens_false(self): + sp = self._build(output_options={"skip_special_tokens": False}) + assert sp.skip_special_tokens is False + + def test_skip_special_tokens_absent(self): + """When not provided, build_sampling_params hardcodes detokenize=False + and SamplingParams default for skip_special_tokens should be unchanged.""" + sp = self._build(output_options={}) + assert sp.detokenize is False + + def test_prompt_logprobs_still_works(self): + """Regression: prompt_logprobs should still be wired alongside skip_special_tokens.""" + sp = self._build( + output_options={"prompt_logprobs": 5, "skip_special_tokens": True} + ) + assert sp.prompt_logprobs == 5 + assert sp.skip_special_tokens is True diff --git a/components/src/dynamo/vllm/worker_factory.py b/components/src/dynamo/vllm/worker_factory.py index 5ff59c718e65..a6e77b76a3ab 100644 --- a/components/src/dynamo/vllm/worker_factory.py +++ b/components/src/dynamo/vllm/worker_factory.py @@ -241,10 +241,17 @@ async def _create_decode_worker( clear_endpoint = runtime.endpoint( f"{config.namespace}.{config.component}.clear_kv_blocks" ) + # RL admin endpoint: discoverable as ..rl on the request + # plane (NATS / TCP). The frontend fans out POST /v1/rl/engine payloads + # here via strict request-plane direct routing. + rl_endpoint = runtime.endpoint( + f"{config.namespace}.{config.component}.rl" + ) shutdown_endpoints[:] = [ generate_endpoint, clear_endpoint, + rl_endpoint, ] lora_enabled = config.engine_args.enable_lora @@ -369,7 +376,7 @@ async def _create_decode_worker( ) # Register engine routes - self.register_engine_routes(runtime, handler) + self.register_engine_routes(runtime, handler, lora_enabled=lora_enabled) # Parse endpoint types from --endpoint-types flag model_type = parse_endpoint_types(config.endpoint_types) @@ -444,6 +451,11 @@ async def _create_decode_worker( handler.get_perf_metrics, metrics_labels=model_metrics_labels, ), + # RL admin: receive {method, kwargs} fan-out payloads from the frontend + rl_endpoint.serve_endpoint( + handler.rl_dispatch, + metrics_labels=model_metrics_labels, + ), ] if lora_enabled: @@ -491,6 +503,9 @@ async def _create_prefill_worker( clear_endpoint = runtime.endpoint( f"{config.namespace}.{config.component}.clear_kv_blocks" ) + rl_endpoint = runtime.endpoint( + f"{config.namespace}.{config.component}.rl" + ) # Use pre-created engine if provided (checkpoint mode), otherwise create new fpm_worker_id = str(generate_endpoint.connection_id()) @@ -580,7 +595,8 @@ async def _create_prefill_worker( ) # Register engine routes - self.register_engine_routes(runtime, handler) + lora_enabled = config.engine_args.enable_lora + self.register_engine_routes(runtime, handler, lora_enabled=lora_enabled) await self._maybe_wait_for_failover_lock(handler, runtime, config) @@ -594,7 +610,7 @@ async def _create_prefill_worker( perf_endpoint = runtime.endpoint( f"{config.namespace}.{config.component}.get_perf_metrics" ) - shutdown_endpoints[:] = [generate_endpoint, clear_endpoint, perf_endpoint] + shutdown_endpoints[:] = [generate_endpoint, clear_endpoint, rl_endpoint, perf_endpoint] # Register prefill model with ModelType.Prefill model_input = ( @@ -641,6 +657,10 @@ async def _create_prefill_worker( handler.get_perf_metrics, metrics_labels=prefill_metrics_labels, ), + rl_endpoint.serve_endpoint( + handler.rl_dispatch, + metrics_labels=prefill_metrics_labels, + ), ) logger.debug("serve_endpoint completed for prefill worker") except Exception as e: @@ -666,19 +686,57 @@ async def _maybe_get_encode_worker_client( return None def register_engine_routes( - self, runtime: DistributedRuntime, handler: BaseWorkerHandler + self, runtime: DistributedRuntime, handler: BaseWorkerHandler, lora_enabled: bool = False ) -> None: """Register all engine routes for this handler. + Two registration paths: + 1. System status server (/engine/): via runtime.register_engine_route. + Used for direct management calls (profiling, sleep/wake). + 2. Request-plane rl endpoint (dyn://..rl): handler._rl_routes dict. + Used by the frontend fan-out via POST /v1/rl/engine. + Args: runtime: The DistributedRuntime instance to register routes on. """ + # System / operational routes (system status server only) runtime.register_engine_route("start_profile", handler.start_profile) runtime.register_engine_route("stop_profile", handler.stop_profile) 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 admin routes — registered in both paths. + # handler._rl_routes is used by rl_dispatch (request plane). + # runtime.register_engine_route also registers on the system status server + # for direct management / testing convenience. + rl_routes: dict = { + # Control plane + "liveness_probe": handler.liveness_probe, + "pause_generation": handler.pause_generation, + "resume_generation": handler.resume_generation, + "flush_cache": handler.flush_cache, + "abort_request": handler.abort_request, + # Weight management + "update_weights_from_disk": handler.update_weights_from_disk, + "update_weights_from_distributed": handler.update_weights_from_distributed, + "update_weights_from_tensor": handler.update_weights_from_tensor, + "init_weights_update_group": handler.init_weights_update_group, + "destroy_weights_update_group": handler.destroy_weights_update_group, + "get_weight_version": handler.get_weight_version, + } + + if lora_enabled: + rl_routes["load_lora_adapter"] = handler.load_lora_adapter + rl_routes["unload_lora_adapter"] = handler.unload_lora_adapter + + for name, fn in rl_routes.items(): + handler._rl_routes[name] = fn + runtime.register_engine_route(name, fn) + logger.info( - "Registered engine routes: /engine/sleep, /engine/wake_up, /engine/scale_elastic_ep, /engine/start_profile, /engine/stop_profile" + "Registered engine routes: sleep, wake_up, scale_elastic_ep, " + "start_profile, stop_profile, and RL admin routes: %s%s", + ", ".join(sorted(rl_routes)), + " (LoRA routes: load_lora_adapter, unload_lora_adapter)" if lora_enabled else "", ) diff --git a/lib/bindings/python/Cargo.lock b/lib/bindings/python/Cargo.lock index 79bc7c2ea797..bef76adccc17 100644 --- a/lib/bindings/python/Cargo.lock +++ b/lib/bindings/python/Cargo.lock @@ -2112,6 +2112,7 @@ dependencies = [ "dynamo-mocker", "dynamo-parsers", "dynamo-protocols", + "dynamo-rl", "dynamo-runtime", "dynamo-tokenizers", "dynamo-tokens", @@ -2287,6 +2288,20 @@ dependencies = [ "uuid", ] +[[package]] +name = "dynamo-rl" +version = "1.2.0" +dependencies = [ + "anyhow", + "axum", + "dynamo-runtime", + "futures", + "serde", + "serde_json", + "tokio", + "tracing", +] + [[package]] name = "dynamo-runtime" version = "1.2.0" diff --git a/lib/llm/Cargo.toml b/lib/llm/Cargo.toml index bd6878704620..fdbd9bc1ed3b 100644 --- a/lib/llm/Cargo.toml +++ b/lib/llm/Cargo.toml @@ -57,6 +57,7 @@ dynamo-config = { workspace = true } dynamo-kv-router = { workspace = true, features = ["metrics", "runtime-protocols"] } dynamo-memory = { workspace = true } dynamo-mocker = { workspace = true } +dynamo-rl = { workspace = true } dynamo-runtime = { workspace = true } dynamo-tokenizers = { workspace = true } dynamo-tokens = { workspace = true } diff --git a/lib/llm/src/backend.rs b/lib/llm/src/backend.rs index 2df3cdf374c9..8914743a9799 100644 --- a/lib/llm/src/backend.rs +++ b/lib/llm/src/backend.rs @@ -389,6 +389,7 @@ impl completion_usage: data.completion_usage, disaggregated_params: data.disaggregated_params, engine_data: data.engine_data, + prompt_logprobs: data.prompt_logprobs, }) }) }); diff --git a/lib/llm/src/entrypoint/input/http.rs b/lib/llm/src/entrypoint/input/http.rs index 09b5d25bd1e1..a8d9ee2517e3 100644 --- a/lib/llm/src/entrypoint/input/http.rs +++ b/lib/llm/src/entrypoint/input/http.rs @@ -64,6 +64,11 @@ pub async fn run( http_service_builder = http_service_builder.drt_discovery(Some(distributed_runtime.discovery())); + // Wire the DistributedRuntime so that RL admin routes can use discovery + + // request-plane fan-out when DYN_ENABLE_RL_ENDPOINTS=true. + http_service_builder = + http_service_builder.runtime(Some(Arc::new(distributed_runtime.clone()))); + let http_service = match engine_config { EngineConfig::Dynamic { ref model, diff --git a/lib/llm/src/http/service/openai.rs b/lib/llm/src/http/service/openai.rs index f0a76b5f5ade..11248bd7d6a2 100644 --- a/lib/llm/src/http/service/openai.rs +++ b/lib/llm/src/http/service/openai.rs @@ -2634,6 +2634,35 @@ pub fn audios_router( (vec![doc], router) } +// --------------------------------------------------------------------------- +// RL admin router +// --------------------------------------------------------------------------- + +/// Build the RL admin axum router and mount it on the dedicated RL listener. +/// +/// Enabled when `DYN_ENABLE_RL_ENDPOINTS=true` or when the caller sets +/// `HttpServiceConfig.enable_rl = true`. The router is served on a separate +/// port (`DYN_RL_PORT`, default 8002). +/// +/// The namespace is read from `DYN_NAMESPACE` (default "dynamo"). Callers can +/// restrict fan-out to a specific component via `DYN_RL_COMPONENT`. +pub(super) fn rl_router( + drt: std::sync::Arc, +) -> anyhow::Result { + let namespace = std::env::var("DYN_NAMESPACE").unwrap_or_else(|_| "dynamo".into()); + + let mut config = dynamo_rl::RlClientConfig::new(drt, namespace); + if let Ok(timeout) = std::env::var("DYN_RL_DEFAULT_TIMEOUT_SECS") + && let Ok(secs) = timeout.parse::() + { + config.default_request_timeout = std::time::Duration::from_secs_f64(secs); + } + + let client = dynamo_rl::RlClient::new(config)?; + let state = dynamo_rl::RlState::new(client); + Ok(dynamo_rl::rl_router(state)) +} + #[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 fc4913b60df2..71cdbd32684c 100644 --- a/lib/llm/src/http/service/service_v2.rs +++ b/lib/llm/src/http/service/service_v2.rs @@ -11,6 +11,7 @@ use std::time::Duration; use axum::body::Body; use axum::http::Response; +use dynamo_runtime::DistributedRuntime; use super::Metrics; use super::RouteDoc; @@ -209,6 +210,9 @@ pub struct HttpService { tls_cert_path: Option, tls_key_path: Option, route_docs: Vec, + /// RL admin router, served on a dedicated port when `DYN_ENABLE_RL_ENDPOINTS=true`. + rl_router: Option, + rl_port: u16, } #[derive(Clone, Builder)] @@ -264,6 +268,27 @@ pub struct HttpServiceConfig { /// are registered using discovery.instance_id() and exposed on /metrics. #[builder(default = "None")] drt_discovery: Option>, + + /// When true (or `DYN_ENABLE_RL_ENDPOINTS=true`), mount the RL admin routes + /// on a dedicated listener at `rl_port`. Requires `runtime` to be set. + #[builder(default = "false")] + enable_rl: bool, + + /// Port for the RL admin listener (default 8002). Ignored when `enable_rl` is false. + #[builder(default = "default_rl_port()")] + rl_port: u16, + + /// The DistributedRuntime used by the RL client for discovery and fan-out. + /// Required when `enable_rl` is true or `DYN_ENABLE_RL_ENDPOINTS=true`. + #[builder(default = "None")] + runtime: Option>, +} + +fn default_rl_port() -> u16 { + std::env::var("DYN_RL_PORT") + .ok() + .and_then(|s| s.parse().ok()) + .unwrap_or(8002) } impl HttpService { @@ -327,6 +352,9 @@ impl HttpService { .handle(handle.clone()) .serve(router.into_make_service()); + // TLS bind succeeded — now safe to start the RL listener. + self.spawn_rl_listener_if_configured(&cancel_token); + // Spawn canary after all fallible startup so it won't leak on early errors tokio::spawn(tokio_metrics_and_canary_loop(cancel_token.clone())); @@ -367,6 +395,9 @@ impl HttpService { } })?; + // Main port is bound — now safe to start the RL admin listener. + self.spawn_rl_listener_if_configured(&cancel_token); + // Spawn canary after all fallible startup so it won't leak on early errors tokio::spawn(tokio_metrics_and_canary_loop(cancel_token.clone())); @@ -388,6 +419,44 @@ impl HttpService { Ok(()) } + /// Spawn the RL admin listener as a background task, if configured. + /// + /// Must be called only after the main HTTP/HTTPS bind has succeeded, so the + /// RL listener is never live while the main port is still in-flight or has + /// already failed to bind. + fn spawn_rl_listener_if_configured(&self, cancel_token: &CancellationToken) { + let Some(rl_router) = self.rl_router.clone() else { + return; + }; + let rl_addr = format!("{}:{}", self.host, self.rl_port); + let rl_cancel = cancel_token.child_token(); + tokio::spawn(async move { + match tokio::net::TcpListener::bind(&rl_addr).await { + Ok(listener) => { + tracing::info!( + address = %rl_addr, + "RL admin listener started on dedicated port" + ); + if let Err(e) = axum::serve(listener, rl_router) + .with_graceful_shutdown(async move { + rl_cancel.cancelled_owned().await; + }) + .await + { + tracing::error!("RL admin listener error: {e}"); + } + } + Err(e) => { + tracing::error!( + address = %rl_addr, + error = %e, + "Failed to bind RL admin listener" + ); + } + } + }); + } + /// Documentation of exposed HTTP endpoints pub fn route_docs(&self) -> &[RouteDoc] { &self.route_docs @@ -579,6 +648,52 @@ impl HttpServiceConfigBuilder { // Echo x-request-id from request to response headers for client correlation let router = router.layer(axum::middleware::from_fn(echo_request_id_header)); + // RL admin router: served on a dedicated listener at `rl_port` AND merged + // onto the main port so that admin clients can use the same base URL as the + // OpenAI-compat endpoint (no separate port required in single-node setups). + // Enabled when `enable_rl || DYN_ENABLE_RL_ENDPOINTS=true` and `runtime` is set. + let (router, rl_router) = if config.enable_rl || env_is_truthy("DYN_ENABLE_RL_ENDPOINTS") { + match config.runtime.as_ref() { + Some(drt) => { + // Build the plain RL router (shared routes, no extra layers). + match super::openai::rl_router(drt.clone()) { + Ok(plain) => { + tracing::info!( + rl_port = config.rl_port, + "RL admin routes enabled at /v1/rl/engine \ + (main port + dedicated listener)" + ); + // Merge plain routes onto main router. + let router = router.merge(plain.clone()); + // Dedicated-port copy gets full tracing. + let dedicated = plain + .layer( + TraceLayer::new_for_http() + .make_span_with(make_system_request_span) + .on_response(on_response), + ) + .layer(axum::middleware::from_fn(echo_request_id_header)); + (router, Some(dedicated)) + } + Err(e) => { + tracing::error!("Failed to build RL router: {e}"); + (router, None) + } + } + } + None => { + return Err(anyhow::anyhow!( + "RL admin routes were requested (DYN_ENABLE_RL_ENDPOINTS=true \ + or enable_rl) but HttpServiceConfig.runtime is not set. \ + The caller must supply a DistributedRuntime via \ + HttpServiceConfigBuilder::runtime()." + )); + } + } + } else { + (router, None) + }; + Ok(HttpService { state, router, @@ -588,6 +703,8 @@ impl HttpServiceConfigBuilder { tls_cert_path: config.tls_cert_path, tls_key_path: config.tls_key_path, route_docs: all_docs, + rl_router, + rl_port: config.rl_port, }) } diff --git a/lib/llm/src/migration.rs b/lib/llm/src/migration.rs index 09143bbb44ee..dee974379292 100644 --- a/lib/llm/src/migration.rs +++ b/lib/llm/src/migration.rs @@ -333,6 +333,7 @@ mod tests { disaggregated_params: None, completion_usage: None, engine_data: None, + prompt_logprobs: None, }) } diff --git a/lib/llm/src/preprocessor.rs b/lib/llm/src/preprocessor.rs index b4d0d3773e41..ace1c16c27e6 100644 --- a/lib/llm/src/preprocessor.rs +++ b/lib/llm/src/preprocessor.rs @@ -464,7 +464,21 @@ impl OpenAIPreprocessor { .with_label_values(&[STAGE_PREPROCESS]) .observe(preprocess_start.elapsed().as_secs_f64()); - Ok((builder.build()?, annotations, prompt_injected_reasoning)) + let mut preprocessed = builder.build()?; + + // rl-sdk-2 TITO parity (Step A5): when the client omits `max_tokens` + // we default to (context_length - prompt_len), mirroring prime-rl's + // vLLM `serving_tokens.py` behavior. Done after token_ids is known + // and clamps to 0 (saturating_sub) for the pathological case where + // the prompt alone fills the context window (validate_token_count + // rejects that earlier, but we belt-and-suspenders here too). + if preprocessed.stop_conditions.max_tokens.is_none() && self.context_length > 0 { + let prompt_len = preprocessed.token_ids.len() as u32; + preprocessed.stop_conditions.max_tokens = + Some(self.context_length.saturating_sub(prompt_len)); + } + + Ok((preprocessed, annotations, prompt_injected_reasoning)) } pub fn builder< @@ -546,6 +560,28 @@ impl OpenAIPreprocessor { session_control: nvext.session_control.clone(), }; builder.routing(Some(routing)); + + // Forward `nvext.extra_fields` (and `nvext.cache_salt`) to the + // backend handler via `extra_args["nvext"]` so it can decide + // what to emit on the response (`engine_data`, `stop_reason`, + // etc.). Mirrors PR #8119's SGLang `_nvext_extra_field_requested` + // helper — both backends now read `request["nvext"]["extra_fields"]` + // identically. The actual extra_args insertion happens at the + // end of build_preprocessed_request (after the multimodal block + // also computes extra_args), so both can coexist on a single + // PreprocessedRequest. + if nvext.extra_fields.is_some() || nvext.cache_salt.is_some() { + let mut nvext_passthrough = serde_json::Map::new(); + if let Some(ref fields) = nvext.extra_fields { + nvext_passthrough.insert("extra_fields".to_string(), serde_json::json!(fields)); + } + if let Some(ref salt) = nvext.cache_salt { + nvext_passthrough.insert("cache_salt".to_string(), serde_json::json!(salt)); + } + builder.extra_args(Some( + serde_json::json!({ "nvext": serde_json::Value::Object(nvext_passthrough) }), + )); + } } else if lora_name.is_some() { // Ensure routing hints exist when we have LoRA, // even when nvext is absent. diff --git a/lib/llm/src/protocols/common.rs b/lib/llm/src/protocols/common.rs index 257bcc170bd2..6f83268a07e9 100644 --- a/lib/llm/src/protocols/common.rs +++ b/lib/llm/src/protocols/common.rs @@ -343,6 +343,12 @@ pub struct SamplingOptions { /// Guided Decoding Options pub guided_decoding: Option, + + /// rl-sdk-2 TITO parity: when `Some(false)`, the engine should skip text + /// decoding of generated token IDs. Mirrors vLLM `SamplingParams.detokenize`. + /// Backends MAY ignore this hint; the only downside is wasted detok work. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub detokenize: Option, } /// Guided Decoding Options diff --git a/lib/llm/src/protocols/common/llm_backend.rs b/lib/llm/src/protocols/common/llm_backend.rs index b6b5f3fad001..8f1ec4b8ed68 100644 --- a/lib/llm/src/protocols/common/llm_backend.rs +++ b/lib/llm/src/protocols/common/llm_backend.rs @@ -14,6 +14,24 @@ use dynamo_runtime::protocols::maybe_error::MaybeError; pub type TokenType = Option; pub type LogProbs = Vec; +/// Per-position prompt logprob entry (rl-sdk-2 TITO parity). +/// +/// Mirrors vLLM's `Logprob` shape so prime-rl's parser at +/// `orchestrator/utils.py:compute_teacher_logprobs` reads it unchanged. +/// `rank` and `decoded_token` are optional; only `logprob` is guaranteed. +#[derive(Serialize, Deserialize, utoipa::ToSchema, Debug, Clone, PartialEq)] +pub struct PromptLogprobEntry { + pub logprob: f32, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub rank: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub decoded_token: Option, +} + +/// Per-token map of `token_id -> PromptLogprobEntry`. The first position +/// is `None` (no logprob exists for BOS / the very first prompt token). +pub type PromptLogprobs = Vec>>; + /// Output type discriminator for different modalities #[derive(Serialize, Deserialize, Debug, Clone, Copy, PartialEq, Eq, Default)] #[serde(rename_all = "lowercase")] @@ -111,6 +129,14 @@ pub struct BackendOutput { /// Dynamo does not inspect this field; it is serialized as-is into `nvext.engine_data`. #[serde(default, skip_serializing_if = "Option::is_none")] pub engine_data: Option, + + /// Per-prompt-token top-k logprobs (rl-sdk-2 TITO parity). + /// + /// Carried through from `LLMEngineOutput.prompt_logprobs`; surfaced to + /// clients on the final response chunk via `nvext.prompt_logprobs` when + /// they opted in with `nvext.extra_fields = ["prompt_logprobs"]`. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub prompt_logprobs: Option, } /// The LLM engine and backnd with manage it's own state, specifically translating how a @@ -177,6 +203,16 @@ pub struct LLMEngineOutput { /// Dynamo does not inspect this field; it is serialized as-is into `nvext.engine_data`. #[serde(default, skip_serializing_if = "Option::is_none")] pub engine_data: Option, + + /// Per-prompt-token top-k logprobs (rl-sdk-2 TITO parity). + /// + /// Populated by engine adapters when the request set + /// `common_ext.prompt_logprobs = Some(k)`. Surfaced to clients via + /// `nvext.prompt_logprobs` when they also include + /// `"prompt_logprobs"` in `nvext.extra_fields`. Backends that don't + /// support prompt logprobs leave this as `None`. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub prompt_logprobs: Option, } impl LLMEngineOutput { @@ -197,6 +233,7 @@ impl LLMEngineOutput { extra_args: None, completion_usage: None, engine_data: None, + prompt_logprobs: None, } } @@ -217,6 +254,7 @@ impl LLMEngineOutput { extra_args: None, completion_usage: None, engine_data: None, + prompt_logprobs: None, } } @@ -237,6 +275,7 @@ impl LLMEngineOutput { extra_args: None, completion_usage: None, engine_data: None, + prompt_logprobs: None, } } @@ -257,6 +296,7 @@ impl LLMEngineOutput { extra_args: None, completion_usage: None, engine_data: None, + prompt_logprobs: None, } } } diff --git a/lib/llm/src/protocols/openai.rs b/lib/llm/src/protocols/openai.rs index 0234ccc6f59e..971b90e9627e 100644 --- a/lib/llm/src/protocols/openai.rs +++ b/lib/llm/src/protocols/openai.rs @@ -160,6 +160,10 @@ impl SamplingOptionsProvid } }; + // rl-sdk-2 TITO parity: surface `detokenize` from CommonExt so the + // engine adapter can disable text decoding when the client asked for it. + let detokenize = self.common_ext().and_then(|c| c.detokenize); + Ok(common::SamplingOptions { n, best_of, @@ -175,6 +179,7 @@ impl SamplingOptionsProvid length_penalty: None, guided_decoding, include_stop_str_in_output, + detokenize, }) } } diff --git a/lib/llm/src/protocols/openai/chat_completions.rs b/lib/llm/src/protocols/openai/chat_completions.rs index c88cfed8d9b4..a541e3ce87a4 100644 --- a/lib/llm/src/protocols/openai/chat_completions.rs +++ b/lib/llm/src/protocols/openai/chat_completions.rs @@ -264,6 +264,14 @@ impl CommonExtProvider for NvCreateChatCompletionRequest { fn get_skip_special_tokens(&self) -> Option { self.common.skip_special_tokens } + + fn get_prompt_logprobs_count(&self) -> Option { + self.common.prompt_logprobs + } + + fn get_detokenize(&self) -> Option { + self.common.detokenize + } } /// Implements `OpenAIStopConditionsProvider` for `NvCreateChatCompletionRequest`, @@ -294,7 +302,18 @@ impl OpenAIStopConditionsProvider for NvCreateChatCompletionRequest { } fn get_stop_token_ids(&self) -> Option> { - self.inner.stop.as_ref().and_then(|stop| stop.token_ids()) + // PR #8119 path: integer IDs in the standard OpenAI `stop` array. + if let Some(ids) = self.inner.stop.as_ref().and_then(|stop| stop.token_ids()) { + return Some(ids); + } + // RL-client compat path: `extra_body.stop_token_ids` (whitelisted via + // `PASSTHROUGH_EXTRA_FIELDS`). Verifiers' `OpenAIChatCompletionsTokenClient` + // and prime-rl's orchestrator both put `stop_token_ids` here when + // controlling generation by token ID. The whitelist accepts it at + // validation; this reader plumbs it into `common::StopConditions`. + self.unsupported_fields + .get("stop_token_ids") + .and_then(|v| serde_json::from_value::>(v.clone()).ok()) } /// Returns a reference to the optional `NvExt` extension, if available. @@ -326,7 +345,11 @@ impl OpenAIOutputOptionsProvider for NvCreateChatCompletionRequest { } fn get_prompt_logprobs(&self) -> Option { - None + // rl-sdk-2 TITO parity: clients (e.g. prime-rl teacher-logprobs path) + // request prompt logprobs by setting top-level `prompt_logprobs: N` on + // the chat-completions request body. Surfaced via CommonExt and then + // copied into common::OutputOptions.prompt_logprobs for the engine. + self.common.prompt_logprobs } fn get_skip_special_tokens(&self) -> Option { @@ -510,14 +533,21 @@ mod tests { serde_json::from_value(scalar_token_id_stop); assert!(result.is_err()); - let unsupported_stop_token_ids = json!({ + // rl-sdk-2: `stop_token_ids` is in PASSTHROUGH_EXTRA_FIELDS (the + // RL-client compat allowlist). Validation must now ACCEPT it (the + // provider trait reads it from `unsupported_fields["stop_token_ids"]` + // and plumbs it into `common::StopConditions.stop_token_ids`). + let whitelisted_stop_token_ids = json!({ "model": "test-model", "messages": [{"role": "user", "content": "Hello"}], "stop_token_ids": [576] }); let request: NvCreateChatCompletionRequest = - serde_json::from_value(unsupported_stop_token_ids) + serde_json::from_value(whitelisted_stop_token_ids) .expect("Failed to deserialize request"); - assert!(ValidateRequest::validate(&request).is_err()); + assert!( + ValidateRequest::validate(&request).is_ok(), + "stop_token_ids must be accepted via PASSTHROUGH_EXTRA_FIELDS" + ); } } diff --git a/lib/llm/src/protocols/openai/chat_completions/delta.rs b/lib/llm/src/protocols/openai/chat_completions/delta.rs index 6f42b2f54556..ca03612d1501 100644 --- a/lib/llm/src/protocols/openai/chat_completions/delta.rs +++ b/lib/llm/src/protocols/openai/chat_completions/delta.rs @@ -415,12 +415,21 @@ impl crate::protocols::openai::DeltaGeneratorExt, + + /// Number of log probabilities to return per prompt token (rl-sdk-2 TITO parity). + /// + /// When set, the engine emits the top-`prompt_logprobs` per-position + /// logprobs over the prompt sequence. Mirrors vLLM's `prompt_logprobs` + /// `SamplingParams` field. Surfaced on the response via + /// `nvext.prompt_logprobs` when the client also opts in with + /// `nvext.extra_fields = ["prompt_logprobs"]`. Required by prime-rl's + /// teacher-logprobs path (`orchestrator/utils.py:compute_teacher_logprobs`). + #[serde(default, skip_serializing_if = "Option::is_none")] + #[builder(default, setter(strip_option))] + pub prompt_logprobs: Option, + + /// Skip text decoding when false (rl-sdk-2 TITO parity). + /// + /// When `Some(false)`, the engine should not run the detokenizer over + /// the generated token IDs. `choices[0].message.content` will be `None` + /// (or empty) and clients are expected to read `nvext.completion_token_ids` + /// instead. Mirrors vLLM's `SamplingParams.detokenize=False`. Defaults + /// to `None` (engine default = true). Backends that do not honor this + /// hint MAY ignore it; the only downside is unnecessary detok work. + #[serde(default, skip_serializing_if = "Option::is_none")] + #[builder(default, setter(strip_option))] + pub detokenize: Option, } impl CommonExt { @@ -111,6 +135,12 @@ pub trait CommonExtProvider { /// Output Options fn get_skip_special_tokens(&self) -> Option; + + /// rl-sdk-2 TITO parity: top-`k` prompt logprobs request (verifier path). + fn get_prompt_logprobs_count(&self) -> Option; + + /// rl-sdk-2 TITO parity: skip text decode when `Some(false)`. + fn get_detokenize(&self) -> Option; } #[cfg(test)] @@ -206,6 +236,8 @@ mod tests { guided_decoding_backend: None, guided_whitespace_pattern: None, skip_special_tokens: None, + prompt_logprobs: None, + detokenize: None, }; assert!(common_ext.validate().is_ok()); } diff --git a/lib/llm/src/protocols/openai/completions.rs b/lib/llm/src/protocols/openai/completions.rs index e60890214bb9..aa4ab7bbb5f8 100644 --- a/lib/llm/src/protocols/openai/completions.rs +++ b/lib/llm/src/protocols/openai/completions.rs @@ -236,6 +236,14 @@ impl CommonExtProvider for NvCreateCompletionRequest { fn get_skip_special_tokens(&self) -> Option { self.common.skip_special_tokens } + + fn get_prompt_logprobs_count(&self) -> Option { + self.common.prompt_logprobs + } + + fn get_detokenize(&self) -> Option { + self.common.detokenize + } } impl OpenAIStopConditionsProvider for NvCreateCompletionRequest { @@ -733,13 +741,18 @@ mod tests { assert_eq!(request.get_stop(), Some(vec!["token_id:576".to_string()])); assert_eq!(request.get_stop_token_ids(), None); - let unsupported_stop_token_ids = json!({ + // rl-sdk-2: `stop_token_ids` is in PASSTHROUGH_EXTRA_FIELDS — must now + // validate OK (it's accepted as a passthrough hint, not a 400). + let whitelisted_stop_token_ids = json!({ "model": "test-model", "prompt": [1, 2, 3], "stop_token_ids": [576] }); - let request: NvCreateCompletionRequest = serde_json::from_value(unsupported_stop_token_ids) + let request: NvCreateCompletionRequest = serde_json::from_value(whitelisted_stop_token_ids) .expect("Failed to deserialize request"); - assert!(ValidateRequest::validate(&request).is_err()); + assert!( + ValidateRequest::validate(&request).is_ok(), + "stop_token_ids must be accepted via PASSTHROUGH_EXTRA_FIELDS" + ); } } diff --git a/lib/llm/src/protocols/openai/completions/delta.rs b/lib/llm/src/protocols/openai/completions/delta.rs index 38cc8d66b642..5f603d51d4fb 100644 --- a/lib/llm/src/protocols/openai/completions/delta.rs +++ b/lib/llm/src/protocols/openai/completions/delta.rs @@ -299,6 +299,10 @@ impl crate::protocols::openai::DeltaGeneratorExt for } } + // rl-sdk-2 TITO parity: clone token_ids before moving them into + // create_logprobs, so build_response_nvext can also see them when + // the client asked for `completion_token_ids`. + let completion_token_ids_for_nvext: Vec = delta.token_ids.clone(); let logprobs = self.create_logprobs( delta.tokens, delta.token_ids, @@ -325,12 +329,17 @@ impl crate::protocols::openai::DeltaGeneratorExt for // `NvExtResponseFieldSelection` (see `nvext.rs`). Both chat and // completions delta generators go through the same helper so the gating // rules stay in one place. + // rl-sdk-2 TITO parity: token_ids was already moved into create_logprobs + // above (see clone made into `completion_token_ids_for_nvext`). + let prompt_logprobs_payload = delta.prompt_logprobs; if let Some(nvext_response) = self.options.response_fields.build_response_nvext( self.tracker.as_ref(), delta.disaggregated_params.as_ref(), finish_reason.is_some(), delta.engine_data, stop_reason, + Some(&completion_token_ids_for_nvext), + prompt_logprobs_payload, ) && let Ok(nvext_json) = serde_json::to_value(&nvext_response) { response.nvext = Some(nvext_json); @@ -347,6 +356,12 @@ impl crate::protocols::openai::DeltaGeneratorExt for tokens.len() ); } + if let Some(ref tokens) = nvext_response.completion_token_ids { + tracing::debug!( + "Injected completion_token_ids into completions nvext: {} tokens", + tokens.len() + ); + } } Ok(response) @@ -426,6 +441,7 @@ mod tests { "routed_experts": {"layer_0": [1, 3]} })), engine_data: None, + prompt_logprobs: None, } } @@ -469,6 +485,7 @@ mod tests { "disaggregated_kv_transfer_time_ms": 8.1, "prefill_compute_time_ms": 45.6 })), + prompt_logprobs: None, } } @@ -735,6 +752,7 @@ mod tests { completion_usage: None, disaggregated_params: None, engine_data: None, // engine didn't provide any data + prompt_logprobs: None, }; let response = generator diff --git a/lib/llm/src/protocols/openai/nvext.rs b/lib/llm/src/protocols/openai/nvext.rs index 7c561dd4a59d..4ee83ea23d1b 100644 --- a/lib/llm/src/protocols/openai/nvext.rs +++ b/lib/llm/src/protocols/openai/nvext.rs @@ -9,11 +9,17 @@ use utoipa::ToSchema; use validator::{Validate, ValidationError}; pub use crate::agents::context::AgentContext; +use crate::protocols::TokenIdType; +pub use crate::protocols::common::llm_backend::PromptLogprobs; pub use crate::protocols::common::timing::TimingInfo; pub const HEADER_WORKER_INSTANCE_ID: &str = "x-worker-instance-id"; pub const HEADER_PREFILL_INSTANCE_ID: &str = "x-prefill-instance-id"; pub const HEADER_DP_RANK: &str = "x-dp-rank"; +/// rl-sdk-2 TITO parity: prime-rl's HTTP client sends data-parallel rank as +/// `X-data-parallel-rank` (verifiers ClientConfig). Accepted as an alias of +/// `x-dp-rank` so existing prime-rl clients work unchanged. +pub const HEADER_DP_RANK_ALIAS: &str = "x-data-parallel-rank"; pub const HEADER_PREFILL_DP_RANK: &str = "x-prefill-dp-rank"; const UNSET_DP_RANK_SENTINEL: u32 = u32::MAX; @@ -38,8 +44,13 @@ pub fn apply_header_routing_overrides(nvext: Option, headers: &HeaderMap) .and_then(|v| v.to_str().ok()) .and_then(|s| s.parse::().ok()); + // rl-sdk-2 TITO parity (Step A6): accept `X-data-parallel-rank` as an + // alias for `x-dp-rank`. Verifiers' ClientConfig sets the alias header + // when prime-rl targets a specific DP rank; without this fallback the + // routing hint silently drops on rl-sdk-2. let dp_rank = headers .get(HEADER_DP_RANK) + .or_else(|| headers.get(HEADER_DP_RANK_ALIAS)) .and_then(|v| v.to_str().ok()) .and_then(|s| s.parse::().ok()); @@ -129,6 +140,28 @@ pub struct NvExtResponse { /// If `n > 1` is supported here, this needs an indexed/per-choice shape. #[serde(skip_serializing_if = "Option::is_none")] pub stop_reason: Option, + + /// rl-sdk-2 TITO parity: engine-emitted output token IDs (top-level shape). + /// + /// Populated when the request set `nvext.extra_fields = ["completion_token_ids"]`. + /// For streaming, each chunk carries the **delta** token IDs for that chunk + /// only (not cumulative) — invariant documented on + /// `NvExtResponseFieldSelection::completion_token_ids`. For non-streaming, + /// this is the concatenation of all chunk deltas. + #[serde(skip_serializing_if = "Option::is_none")] + pub completion_token_ids: Option>, + + /// rl-sdk-2 TITO parity: per-prompt-token top-k logprobs. + /// + /// Populated on the **final** chunk only (chunks with + /// `finish_reason.is_some()`) when the request set + /// `nvext.extra_fields = ["prompt_logprobs"]` AND the engine adapter + /// populated `LLMEngineOutput.prompt_logprobs`. The shape mirrors vLLM's + /// `RequestOutput.prompt_logprobs`: `[None, {tok_id: {logprob, ...}}, ...]`. + /// Prime-rl's parser at `orchestrator/utils.py:compute_teacher_logprobs` + /// reads `.logprob` of each entry. + #[serde(skip_serializing_if = "Option::is_none")] + pub prompt_logprobs: Option, } pub(crate) fn merge_response_nvext( @@ -141,7 +174,35 @@ pub(crate) fn merge_response_nvext( match (target.as_mut(), incoming) { (Some(serde_json::Value::Object(target_obj)), serde_json::Value::Object(incoming_obj)) => { - target_obj.extend(incoming_obj); + // rl-sdk-2 TITO parity (Step A9 — aggregator side): per-chunk + // streaming response carries DELTA token IDs; the SSE-to-single + // response aggregator must CONCATENATE rather than overwrite for + // `completion_token_ids`. Same for `prompt_logprobs` (though that + // one is emitted on the final chunk only, so concatenation is a + // no-op in the steady state — defensive anyway). All other keys + // keep "later wins" semantics. + for (key, value) in incoming_obj { + match key.as_str() { + "completion_token_ids" => { + let entry = target_obj + .entry(&key) + .or_insert_with(|| serde_json::Value::Array(Vec::new())); + if let (serde_json::Value::Array(acc), serde_json::Value::Array(new)) = + (entry, value) + { + acc.extend(new); + } + } + "prompt_logprobs" => { + // Final-chunk-only on the producer side, so on collision + // we keep the latest (which is the one with finish_reason). + target_obj.insert(key, value); + } + _ => { + target_obj.insert(key, value); + } + } + } } (_, incoming) => { *target = Some(incoming); @@ -165,6 +226,15 @@ pub struct NvExtResponseFieldSelection { pub routed_experts: bool, pub engine_data: bool, pub stop_reason: bool, + /// rl-sdk-2 TITO parity: emit per-chunk delta `completion_token_ids` and + /// the final concatenated list on the aggregated response. Wire-level + /// invariant for streaming: each chunk carries the **delta** tokens for + /// that chunk only (matches vLLM `RequestOutputKind::DELTA`). + pub completion_token_ids: bool, + /// rl-sdk-2 TITO parity: emit `nvext.prompt_logprobs` on the final chunk + /// only. Inert unless the engine adapter actually populated + /// `LLMEngineOutput.prompt_logprobs`. + pub prompt_logprobs: bool, } impl NvExtResponseFieldSelection { @@ -182,6 +252,8 @@ impl NvExtResponseFieldSelection { "routed_experts" => selection.routed_experts = true, "engine_data" => selection.engine_data = true, "stop_reason" => selection.stop_reason = true, + "completion_token_ids" => selection.completion_token_ids = true, + "prompt_logprobs" => selection.prompt_logprobs = true, _ => {} } } @@ -222,6 +294,9 @@ impl NvExtResponseFieldSelection { finish_reason_present: bool, engine_data_from_backend: Option, stop_reason_from_backend: Option, + // rl-sdk-2 TITO parity additions: + completion_token_ids_from_backend: Option<&[TokenIdType]>, + prompt_logprobs_from_backend: Option, ) -> Option { let worker_id = if self.worker_id { tracker.and_then(|t| t.get_worker_info()) @@ -263,12 +338,34 @@ impl NvExtResponseFieldSelection { None }; + // rl-sdk-2 TITO parity: emit `completion_token_ids` on every chunk + // when the client asked for it. Streaming: per-chunk deltas only — + // the caller passes the chunk's `BackendOutput.token_ids` slice. + // Aggregator: passes the concatenated list. Either way, this helper + // is a pure pass-through (no accumulation done here). + let completion_token_ids = if self.completion_token_ids { + completion_token_ids_from_backend.map(<[u32]>::to_vec) + } else { + None + }; + + // rl-sdk-2 TITO parity: emit `prompt_logprobs` on the **final** chunk + // only. Prompt logprobs are produced once for the whole prompt by the + // engine and we don't want to repeat them on every streaming chunk. + let prompt_logprobs = if self.prompt_logprobs && finish_reason_present { + prompt_logprobs_from_backend + } else { + None + }; + if worker_id.is_none() && token_ids.is_none() && routed_experts.is_none() && timing.is_none() && engine_data.is_none() && stop_reason.is_none() + && completion_token_ids.is_none() + && prompt_logprobs.is_none() { return None; } @@ -280,6 +377,8 @@ impl NvExtResponseFieldSelection { routed_experts, engine_data, stop_reason, + completion_token_ids, + prompt_logprobs, }) } } @@ -328,6 +427,21 @@ pub struct NvExt { #[builder(default, setter(strip_option))] pub max_thinking_tokens: Option, + /// KV prefix-cache isolation hint from RL orchestrators. + /// + /// Prime-RL's orchestrator tags every rollout request with a `cache_salt` + /// derived from the current checkpoint step (e.g. `"step_7"`). When the + /// salt changes across requests, the inference engine treats their prompt + /// prefixes as distinct cache keys even if the token sequences are + /// byte-identical — ensuring KV cache hits from the pre-weight-update + /// policy do not leak into post-update generations. + /// + /// Dynamo passes this through to the backend as a sampling-params hint. + /// Backends that do not support cache salting may ignore it. + #[serde(default, skip_serializing_if = "Option::is_none")] + #[builder(default, setter(strip_option))] + pub cache_salt: Option, + /// Extra fields to be included in the response's nvext /// This is a list of field names that should be populated in the response /// Supported fields include "worker_id", "timing", "routed_experts", "engine_data", @@ -507,6 +621,7 @@ mod tests { assert_eq!(nv_ext.backend_instance_id, None); assert_eq!(nv_ext.token_data, None); assert_eq!(nv_ext.max_thinking_tokens, None); + assert_eq!(nv_ext.cache_salt, None); assert_eq!(nv_ext.extra_fields, None); assert_eq!(nv_ext.prefill_worker_id, None); assert_eq!(nv_ext.decode_worker_id, None); @@ -791,12 +906,12 @@ mod tests { fn test_build_response_nvext_all_false_returns_none() { let sel = sel_all_false(); assert!( - sel.build_response_nvext(None, None, false, None, None) + sel.build_response_nvext(None, None, false, None, None, None, None) .is_none(), "no fields selected → None" ); assert!( - sel.build_response_nvext(None, None, true, None, None) + sel.build_response_nvext(None, None, true, None, None, None, None) .is_none(), "finish_reason alone does not force emission" ); @@ -812,7 +927,7 @@ mod tests { // finish_reason=false: worker_id still emitted (only timing is finish-gated). let out = sel - .build_response_nvext(Some(&tracker), None, false, None, None) + .build_response_nvext(Some(&tracker), None, false, None, None, None, None) .expect("worker_id should emit regardless of finish_reason"); assert!(out.worker_id.is_some()); @@ -831,7 +946,7 @@ mod tests { // timing alone + finish_reason=false → nothing to emit, returns None. assert!( - sel.build_response_nvext(Some(&tracker), None, false, None, None) + sel.build_response_nvext(Some(&tracker), None, false, None, None, None, None) .is_none(), "timing is gated on finish_reason_present" ); @@ -846,7 +961,7 @@ mod tests { let tracker = tracker_with_prefill_worker(); let out = sel - .build_response_nvext(Some(&tracker), None, true, None, None) + .build_response_nvext(Some(&tracker), None, true, None, None, None, None) .expect("timing should emit on finish"); assert!(out.timing.is_some()); @@ -863,7 +978,7 @@ mod tests { }; // finish=true but no tracker → timing not populated → None. assert!( - sel.build_response_nvext(None, None, true, None, None) + sel.build_response_nvext(None, None, true, None, None, None, None) .is_none() ); } @@ -877,7 +992,7 @@ mod tests { let params = disagg_params_full(); let out = sel - .build_response_nvext(None, Some(¶ms), false, None, None) + .build_response_nvext(None, Some(¶ms), false, None, None, None, None) .expect("token_ids should emit when present"); assert_eq!(out.token_ids, Some(vec![11u32, 22, 33])); @@ -896,7 +1011,7 @@ mod tests { let params = serde_json::json!({ "token_ids": "not-an-array" }); assert!( - sel.build_response_nvext(None, Some(¶ms), false, None, None) + sel.build_response_nvext(None, Some(¶ms), false, None, None, None, None) .is_none(), "malformed token_ids silently suppressed; nothing else selected → None" ); @@ -911,7 +1026,7 @@ mod tests { let params = disagg_params_full(); let out = sel - .build_response_nvext(None, Some(¶ms), false, None, None) + .build_response_nvext(None, Some(¶ms), false, None, None, None, None) .expect("routed_experts should emit when present"); assert_eq!( @@ -934,6 +1049,8 @@ mod tests { true, None, Some(StopReason::String("END".to_string())), + None, + None, ) .expect("stop_reason should emit when requested and present"); @@ -952,7 +1069,7 @@ mod tests { }; assert!( - sel.build_response_nvext(None, None, true, None, None) + sel.build_response_nvext(None, None, true, None, None, None, None) .is_none() ); } @@ -966,12 +1083,14 @@ mod tests { routed_experts: true, engine_data: false, stop_reason: false, + completion_token_ids: false, + prompt_logprobs: false, }; let tracker = tracker_with_prefill_worker(); let params = disagg_params_full(); let out = sel - .build_response_nvext(Some(&tracker), Some(¶ms), true, None, None) + .build_response_nvext(Some(&tracker), Some(¶ms), true, None, None, None, None) .expect("all fields selected and available → Some"); assert!(out.worker_id.is_some()); @@ -1003,7 +1122,106 @@ mod tests { routed_experts: true, engine_data: false, stop_reason: false, + completion_token_ids: false, + prompt_logprobs: false, } ); } + + // ---------- rl-sdk-2 TITO parity coverage (Steps A3/A4/A9) ---------- + + #[test] + fn tito_parity_completion_token_ids_pass_through() { + // Step A4: per-chunk delta tokens land verbatim on + // `nvext.completion_token_ids` when the client opted in. + let sel = NvExtResponseFieldSelection { + completion_token_ids: true, + ..Default::default() + }; + let chunk_tokens: &[u32] = &[101, 102, 103]; + let out = sel + .build_response_nvext(None, None, false, None, None, Some(chunk_tokens), None) + .expect("completion_token_ids must be present when requested + provided"); + assert_eq!(out.completion_token_ids, Some(vec![101u32, 102, 103])); + // Strict-delta invariant doc: the helper is a pure pass-through. + // No accumulation here; that's the aggregator's job (verified by + // merge_response_nvext concat semantics). + assert!(out.prompt_logprobs.is_none()); + assert!(out.engine_data.is_none()); + } + + #[test] + fn tito_parity_prompt_logprobs_final_chunk_only() { + // Step A3: prompt_logprobs payload must NOT leak onto intermediate + // chunks; only emitted when finish_reason.is_some(). This guarantees + // prime-rl's teacher-logprobs parser sees one logprobs blob per + // request, not duplicates across chunks. + let sel = NvExtResponseFieldSelection { + prompt_logprobs: true, + ..Default::default() + }; + let mut entry = std::collections::HashMap::new(); + entry.insert( + 42u32, + crate::protocols::common::llm_backend::PromptLogprobEntry { + logprob: -1.234, + rank: Some(1), + decoded_token: None, + }, + ); + let payload: crate::protocols::common::llm_backend::PromptLogprobs = + vec![None, Some(entry)]; + + // Intermediate chunk (no finish): suppressed. + assert!( + sel.build_response_nvext(None, None, false, None, None, None, Some(payload.clone())) + .is_none(), + "prompt_logprobs must be suppressed on intermediate chunks" + ); + + // Final chunk: surfaced. + let out = sel + .build_response_nvext(None, None, true, None, None, None, Some(payload.clone())) + .expect("prompt_logprobs must emit on the final chunk"); + let got = out.prompt_logprobs.expect("prompt_logprobs payload"); + assert_eq!(got.len(), 2); + assert!(got[0].is_none()); + assert_eq!( + got[1].as_ref().unwrap().get(&42u32).unwrap().logprob, + -1.234 + ); + } + + #[test] + fn tito_parity_aggregator_concatenates_completion_token_ids() { + // Step A9: SSE aggregator must concatenate per-chunk completion_token_ids, + // not overwrite (which would lose all but the final chunk's tokens). + let mut target: Option = None; + // Chunk 1: tokens [10, 11, 12] + merge_response_nvext( + &mut target, + Some(serde_json::json!({ "completion_token_ids": [10, 11, 12] })), + ); + // Chunk 2: tokens [13, 14] — must append, not replace. + merge_response_nvext( + &mut target, + Some(serde_json::json!({ "completion_token_ids": [13, 14] })), + ); + // Chunk 3 (final): one more token + non-token field that should overwrite. + merge_response_nvext( + &mut target, + Some(serde_json::json!({ + "completion_token_ids": [15], + "worker_id": { "decode_worker_id": 7 } + })), + ); + + let aggregated = target.expect("aggregator state"); + assert_eq!( + aggregated["completion_token_ids"], + serde_json::json!([10, 11, 12, 13, 14, 15]), + "completion_token_ids must concatenate across chunks" + ); + assert_eq!(aggregated["worker_id"]["decode_worker_id"], 7); + } } diff --git a/lib/llm/src/protocols/openai/validate.rs b/lib/llm/src/protocols/openai/validate.rs index 559dd109ac1b..6e99361dde8b 100644 --- a/lib/llm/src/protocols/openai/validate.rs +++ b/lib/llm/src/protocols/openai/validate.rs @@ -97,16 +97,45 @@ pub const MAX_REPETITION_PENALTY: f32 = 2.0; // Shared Fields // -/// Validates that no unsupported fields are present in the request +/// Fields that RL clients (prime-rl, verifiers) may send as `extra_body` +/// hints which Dynamo does not implement as first-class request fields +/// but should not reject with a 400. They are accepted at the validator +/// layer; downstream consumers (provider trait, preprocessor, backend) +/// MAY choose to read them for their own purposes. +/// +/// Canonical locations going forward: +/// - `cache_salt` → `nvext.cache_salt` (this allowlist entry is for +/// backward compat with clients that still send it +/// under `extra_body.cache_salt`). +/// - `stop_token_ids` → `nvext.stop_token_ids` (planned). For now the +/// provider trait reads from this allowlist via +/// `unsupported_fields["stop_token_ids"]` and +/// plumbs into `common::StopConditions.stop_token_ids`. +/// - `bad_words_token_ids`, `allowed_token_ids`, `truncate_prompt_tokens` +/// → vLLM `SamplingParams` parity passthroughs. +pub const PASSTHROUGH_EXTRA_FIELDS: &[&str] = &[ + "cache_salt", + "stop_token_ids", + "bad_words_token_ids", + "allowed_token_ids", + "truncate_prompt_tokens", +]; + +/// Validates that no unsupported fields are present in the request. +/// +/// Fields in `PASSTHROUGH_EXTRA_FIELDS` are silently accepted (they may +/// be read elsewhere — e.g. by `get_stop_token_ids()` on the provider +/// trait — but their absence from this validator does not cause a 400). 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/protocols/unified.rs b/lib/llm/src/protocols/unified.rs index c748678126e5..b4e2eeac73f3 100644 --- a/lib/llm/src/protocols/unified.rs +++ b/lib/llm/src/protocols/unified.rs @@ -388,6 +388,14 @@ impl CommonExtProvider for UnifiedRequest { fn get_skip_special_tokens(&self) -> Option { self.inner.common.skip_special_tokens } + + fn get_prompt_logprobs_count(&self) -> Option { + self.inner.common.prompt_logprobs + } + + fn get_detokenize(&self) -> Option { + self.inner.common.detokenize + } } impl OpenAIStopConditionsProvider for UnifiedRequest { diff --git a/lib/llm/tests/test_streaming_usage.rs b/lib/llm/tests/test_streaming_usage.rs index eb91c305da68..e36a07019058 100644 --- a/lib/llm/tests/test_streaming_usage.rs +++ b/lib/llm/tests/test_streaming_usage.rs @@ -111,6 +111,7 @@ fn build_backend_outputs_with_cached_tokens(cached_tokens: Option) -> Vec) -> Vec) -> Vec BackendOutput { completion_usage: None, disaggregated_params: None, engine_data: None, + prompt_logprobs: None, } } @@ -305,6 +306,7 @@ async fn test_streaming_named_tool_buffers_until_finish() { completion_usage: None, disaggregated_params: None, engine_data: None, + prompt_logprobs: None, }; let response = generator @@ -373,6 +375,7 @@ async fn test_streaming_required_tool_parallel() { completion_usage: None, disaggregated_params: None, engine_data: None, + prompt_logprobs: None, }; let response = generator @@ -443,6 +446,7 @@ fn test_no_tool_choice_outputs_normal_text() { completion_usage: None, disaggregated_params: None, engine_data: None, + prompt_logprobs: None, }; let response = generator diff --git a/lib/llm/tests/tool_choice_finish_reasons.rs b/lib/llm/tests/tool_choice_finish_reasons.rs index 9f66338d24b7..f797edf9b3fb 100644 --- a/lib/llm/tests/tool_choice_finish_reasons.rs +++ b/lib/llm/tests/tool_choice_finish_reasons.rs @@ -52,6 +52,7 @@ fn build_backend_output_with_finish(text: &str, finish: common::FinishReason) -> completion_usage: None, disaggregated_params: None, engine_data: None, + prompt_logprobs: None, } } diff --git a/lib/rl/Cargo.toml b/lib/rl/Cargo.toml new file mode 100644 index 000000000000..2b97709185c3 --- /dev/null +++ b/lib/rl/Cargo.toml @@ -0,0 +1,26 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +[package] +name = "dynamo-rl" +description = "RL admin control plane — generic fan-out facade for /v1/rl/engine" +version.workspace = true +edition.workspace = true +authors.workspace = true +license.workspace = true +homepage.workspace = true +repository.workspace = true + +# Dependency direction: dynamo-llm -> dynamo-rl -> dynamo-runtime. +# This crate must NOT depend on dynamo-llm. + +[dependencies] +dynamo-runtime = { workspace = true } + +axum = { workspace = true } +serde = { workspace = true } +serde_json = { workspace = true } +tokio = { workspace = true } +tracing = { workspace = true } +anyhow = { workspace = true } +futures = { workspace = true } diff --git a/lib/rl/src/lib.rs b/lib/rl/src/lib.rs new file mode 100644 index 000000000000..65524e135ccd --- /dev/null +++ b/lib/rl/src/lib.rs @@ -0,0 +1,688 @@ +// SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! Dynamo RL admin control plane — generic fan-out facade for `/v1/rl/engine`. +//! +//! ## Architecture +//! +//! Workers register an `rl` endpoint on the request plane at startup: +//! ```text +//! dyn://..rl +//! ``` +//! +//! The frontend discovers live `rl` instances via the discovery plane and fans +//! out `{method, kwargs}` payloads via strict request-plane direct routing +//! (NATS / shared TCP). The Python `rl_dispatch` handler on each worker +//! receives the payload and dispatches to the appropriate registered handler. +//! +//! ## HTTP surface (mounted on the frontend) +//! +//! ```text +//! POST /v1/rl/engine fan-out or direct call, selected by body +//! GET /v1/rl/engine list registered methods per live worker +//! ``` +//! +//! Enabled when `DYN_ENABLE_RL_ENDPOINTS=true`, on port `DYN_RL_PORT` +//! (default: 8002 if unset, or shares the main port). + +use std::{ + collections::{HashMap, hash_map::DefaultHasher}, + hash::{Hash, Hasher}, + sync::Arc, + time::Duration, +}; + +use axum::{Json, Router, extract::State, http::StatusCode, response::IntoResponse, routing::post}; +use dynamo_runtime::{ + DistributedRuntime, + component::Client, + discovery::{DiscoveryInstance, DiscoveryQuery}, + pipeline::{ + SingleIn, + network::egress::push_router::{PushRouter, RouterMode}, + }, + protocols::annotated::Annotated, +}; +use futures::{FutureExt, StreamExt}; + +// --------------------------------------------------------------------------- +// Public constants +// --------------------------------------------------------------------------- + +pub const DEFAULT_RL_ENDPOINT: &str = "rl"; +pub const DEFAULT_REQUEST_TIMEOUT_SECS: u64 = 30; + +// --------------------------------------------------------------------------- +// Error types +// --------------------------------------------------------------------------- + +#[derive(Debug, Clone)] +pub enum RlError { + NoWorkers { + namespace: String, + rl_endpoint: String, + }, + MembershipChanged { + before_epoch: u64, + after_epoch: u64, + }, + UnknownInstance { + instance_id: u64, + }, +} + +impl std::fmt::Display for RlError { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + RlError::NoWorkers { + namespace, + rl_endpoint, + } => write!( + f, + "no live RL workers found in namespace '{namespace}' for endpoint '{rl_endpoint}'" + ), + RlError::MembershipChanged { + before_epoch, + after_epoch, + } => write!( + f, + "RL worker membership changed during fan-out (before={before_epoch}, after={after_epoch})" + ), + RlError::UnknownInstance { instance_id } => { + write!(f, "instance_id {instance_id} not found in live worker set") + } + } + } +} + +impl std::error::Error for RlError {} + +// --------------------------------------------------------------------------- +// Client configuration +// --------------------------------------------------------------------------- + +pub struct RlClientConfig { + pub runtime: Arc, + pub namespace: String, + /// Worker endpoint name (default: "rl") + pub rl_endpoint: String, + /// Per-worker call timeout default; overridden per call via `timeout_secs` + pub default_request_timeout: Duration, +} + +impl RlClientConfig { + pub fn new(runtime: Arc, namespace: impl Into) -> Self { + Self { + runtime, + namespace: namespace.into(), + rl_endpoint: DEFAULT_RL_ENDPOINT.to_string(), + default_request_timeout: Duration::from_secs(DEFAULT_REQUEST_TIMEOUT_SECS), + } + } +} + +// --------------------------------------------------------------------------- +// Wire types +// --------------------------------------------------------------------------- + +/// Payload sent from the frontend to each worker's `rl` endpoint. +#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)] +pub struct RlRequest { + pub method: String, + #[serde(default)] + pub kwargs: serde_json::Value, +} + +/// Per-worker result collected by the frontend after fan-out. +#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)] +pub struct WorkerResult { + pub instance_id: u64, + pub component: String, + pub status: WorkerStatus, + #[serde(skip_serializing_if = "Option::is_none")] + pub response: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub error: Option, +} + +#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)] +#[serde(rename_all = "lowercase")] +pub enum WorkerStatus { + Ok, + Error, +} + +impl WorkerResult { + fn ok(target: &WorkerTarget, response: serde_json::Value) -> Self { + Self { + instance_id: target.instance_id, + component: target.component.clone(), + status: WorkerStatus::Ok, + response: Some(response), + error: None, + } + } + + fn error(target: &WorkerTarget, error: impl Into) -> Self { + Self { + instance_id: target.instance_id, + component: target.component.clone(), + status: WorkerStatus::Error, + response: None, + error: Some(error.into()), + } + } +} + +// --------------------------------------------------------------------------- +// Discovery / membership +// --------------------------------------------------------------------------- + +#[derive( + Debug, Clone, serde::Serialize, serde::Deserialize, PartialEq, Eq, Hash, PartialOrd, Ord, +)] +pub struct WorkerTarget { + pub namespace: String, + pub component: String, + pub endpoint: String, + pub instance_id: u64, +} + +#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)] +pub struct MembershipSnapshot { + pub epoch: u64, + pub targets: Vec, +} + +impl MembershipSnapshot { + fn new(mut targets: Vec) -> Self { + targets.sort(); + targets.dedup(); + + let mut hasher = DefaultHasher::new(); + targets.hash(&mut hasher); + let epoch = hasher.finish(); + + Self { epoch, targets } + } + + pub fn is_empty(&self) -> bool { + self.targets.is_empty() + } +} + +/// Aggregated result from a fan-out call. +#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)] +pub struct FanoutReport { + pub epoch: u64, + pub workers: Vec, +} + +// --------------------------------------------------------------------------- +// Per-call options +// --------------------------------------------------------------------------- + +#[derive(Debug, Clone)] +pub struct CallOptions { + pub timeout: Duration, + /// Component names to restrict fan-out. `None` = all components. + pub components: Option>, +} + +impl CallOptions { + pub fn new(timeout: Duration) -> Self { + Self { + timeout, + components: None, + } + } + + pub fn with_components(mut self, components: Vec) -> Self { + self.components = Some(components); + self + } +} + +// --------------------------------------------------------------------------- +// RlClient +// --------------------------------------------------------------------------- + +#[derive(Clone)] +pub struct RlClient { + runtime: Arc, + namespace: String, + rl_endpoint: String, + default_request_timeout: Duration, +} + +impl RlClient { + pub fn new(config: RlClientConfig) -> anyhow::Result { + if config.namespace.trim().is_empty() { + anyhow::bail!("RlClientConfig.namespace must not be empty"); + } + if config.rl_endpoint.trim().is_empty() { + anyhow::bail!("RlClientConfig.rl_endpoint must not be empty"); + } + Ok(Self { + runtime: config.runtime, + namespace: config.namespace, + rl_endpoint: config.rl_endpoint, + default_request_timeout: config.default_request_timeout, + }) + } + + /// Snapshot the current live worker set. + pub async fn snapshot(&self, opts: &CallOptions) -> anyhow::Result { + let instances = self + .runtime + .discovery() + .list(DiscoveryQuery::NamespacedEndpoints { + namespace: self.namespace.clone(), + }) + .await?; + + let targets = instances + .into_iter() + .filter_map(|instance| match instance { + DiscoveryInstance::Endpoint(ep) if ep.endpoint == self.rl_endpoint => Some(ep), + _ => None, + }) + .filter(|ep| { + opts.components + .as_ref() + .map(|cs| cs.iter().any(|c| c == &ep.component)) + .unwrap_or(true) + }) + .map(|ep| WorkerTarget { + namespace: ep.namespace, + component: ep.component, + endpoint: ep.endpoint, + instance_id: ep.instance_id, + }) + .collect(); + + Ok(MembershipSnapshot::new(targets)) + } + + /// Fan out `method` + `kwargs` to all live workers (or those matched by + /// `opts.components`). Returns a `FanoutReport` with per-worker results. + /// Returns 503 if zero workers are discovered. + /// Returns 409 if the worker set changes between pre- and post-fan-out + /// snapshots (abort-on-membership-change). + pub async fn engine_call( + &self, + method: &str, + kwargs: serde_json::Value, + opts: CallOptions, + ) -> anyhow::Result { + let snapshot = self.snapshot(&opts).await?; + if snapshot.is_empty() { + return Err(RlError::NoWorkers { + namespace: self.namespace.clone(), + rl_endpoint: self.rl_endpoint.clone(), + } + .into()); + } + self.fanout_snapshot(&snapshot, method, kwargs, opts.timeout) + .await + } + + /// Call exactly one worker (strict-direct). Returns 409 if the instance + /// has vanished from the live set before the call is dispatched. + pub async fn engine_call_one( + &self, + method: &str, + kwargs: serde_json::Value, + instance_id: u64, + opts: CallOptions, + ) -> anyhow::Result { + let snapshot = self.snapshot(&opts).await?; + let target = snapshot + .targets + .iter() + .find(|t| t.instance_id == instance_id) + .cloned() + .ok_or(RlError::UnknownInstance { instance_id })?; + + let result = call_worker_target(&self.runtime, &target, method, kwargs, opts.timeout).await; + Ok(FanoutReport { + epoch: snapshot.epoch, + workers: vec![result], + }) + } + + async fn fanout_snapshot( + &self, + snapshot: &MembershipSnapshot, + method: &str, + kwargs: serde_json::Value, + timeout: Duration, + ) -> anyhow::Result { + // Group targets by (namespace, component, endpoint) — one PushRouter per group. + let mut grouped: HashMap<(String, String, String), Vec> = HashMap::new(); + for target in &snapshot.targets { + grouped + .entry(( + target.namespace.clone(), + target.component.clone(), + target.endpoint.clone(), + )) + .or_default() + .push(target.clone()); + } + + let mut calls: Vec> = Vec::new(); + + for ((namespace, component, endpoint_name), targets) in grouped { + let endpoint = match self + .runtime + .namespace(&namespace) + .and_then(|ns| ns.component(&component)) + { + Ok(comp) => comp.endpoint(endpoint_name), + Err(err) => { + for target in targets { + let err_str = format!("endpoint build failed: {err}"); + calls.push( + futures::future::ready(WorkerResult::error(&target, err_str)).boxed(), + ); + } + continue; + } + }; + + let client = match endpoint.client().await { + Ok(c) => c, + Err(err) => { + for target in targets { + let err_str = format!("client create failed: {err}"); + calls.push( + futures::future::ready(WorkerResult::error(&target, err_str)).boxed(), + ); + } + continue; + } + }; + + let target_ids: Vec = targets.iter().map(|t| t.instance_id).collect(); + wait_for_client_targets(&client, &target_ids, Duration::from_secs(5)).await; + + let router = + match PushRouter::>::from_client( + client, + RouterMode::Direct, + ) + .await + { + Ok(r) => r, + Err(err) => { + for target in targets { + let err_str = format!("PushRouter build failed: {err}"); + calls.push( + futures::future::ready(WorkerResult::error(&target, err_str)) + .boxed(), + ); + } + continue; + } + }; + + for target in targets { + calls.push( + call_worker( + router.clone(), + target, + method.to_string(), + kwargs.clone(), + timeout, + ) + .boxed(), + ); + } + } + + let workers = futures::future::join_all(calls).await; + + // Abort-on-membership-change: check epoch after fan-out completes. + let after_opts = CallOptions::new(timeout); + let after = self.snapshot(&after_opts).await?; + if after.epoch != snapshot.epoch { + return Err(RlError::MembershipChanged { + before_epoch: snapshot.epoch, + after_epoch: after.epoch, + } + .into()); + } + + Ok(FanoutReport { + epoch: snapshot.epoch, + workers, + }) + } + + pub fn default_timeout(&self) -> Duration { + self.default_request_timeout + } +} + +// --------------------------------------------------------------------------- +// Internal helpers +// --------------------------------------------------------------------------- + +async fn wait_for_client_targets(client: &Client, target_ids: &[u64], timeout: Duration) { + let wait = async { + loop { + let ids = client.instance_ids(); + if target_ids.iter().all(|id| ids.contains(id)) { + break; + } + tokio::time::sleep(Duration::from_millis(25)).await; + } + }; + let _ = tokio::time::timeout(timeout, wait).await; +} + +async fn call_worker_target( + runtime: &Arc, + target: &WorkerTarget, + method: &str, + kwargs: serde_json::Value, + timeout: Duration, +) -> WorkerResult { + let endpoint = match runtime + .namespace(&target.namespace) + .and_then(|ns| ns.component(&target.component)) + { + Ok(comp) => comp.endpoint(target.endpoint.clone()), + Err(err) => return WorkerResult::error(target, format!("endpoint build failed: {err}")), + }; + + let client = match endpoint.client().await { + Ok(c) => c, + Err(err) => return WorkerResult::error(target, format!("client create failed: {err}")), + }; + + wait_for_client_targets(&client, &[target.instance_id], Duration::from_secs(5)).await; + + let router = match PushRouter::>::from_client( + client, + RouterMode::Direct, + ) + .await + { + Ok(r) => r, + Err(err) => return WorkerResult::error(target, format!("PushRouter build failed: {err}")), + }; + + call_worker(router, target.clone(), method.to_string(), kwargs, timeout).await +} + +async fn call_worker( + router: PushRouter>, + target: WorkerTarget, + method: String, + kwargs: serde_json::Value, + timeout: Duration, +) -> WorkerResult { + let request_value = serde_json::json!({ + "method": method, + "kwargs": kwargs, + }); + + let instance_id = target.instance_id; + + let dispatch = async { + let req = SingleIn::new(request_value); + let mut stream = router.direct(req, instance_id).await?; + + while let Some(chunk) = stream.next().await { + if let Some(data) = chunk.data { + return anyhow::Ok(data); + } + if let Some(err) = chunk.error { + anyhow::bail!(err.to_string()); + } + } + + anyhow::bail!("empty response stream from worker"); + }; + + match tokio::time::timeout(timeout, dispatch).await { + Ok(Ok(response)) => WorkerResult::ok(&target, response), + Ok(Err(err)) => WorkerResult::error(&target, format!("dispatch failed: {err}")), + Err(_) => WorkerResult::error( + &target, + format!("dispatch timed out after {}s", timeout.as_secs()), + ), + } +} + +// --------------------------------------------------------------------------- +// HTTP facade +// --------------------------------------------------------------------------- + +/// HTTP request body for `POST /v1/rl/engine`. +#[derive(Debug, serde::Deserialize)] +pub struct RlEngineRequest { + pub method: String, + #[serde(default)] + pub kwargs: serde_json::Value, + /// If present, call only this instance (strict-direct). Absent = fan-out. + pub instance_id: Option, + /// Per-call timeout override (seconds). Falls back to `RlClientConfig.default_request_timeout`. + pub timeout_secs: Option, + /// Restrict fan-out to these component names. `None` = all. + pub components: Option>, +} + +impl RlEngineRequest { + fn call_options(&self, default_timeout: Duration) -> CallOptions { + let timeout = self + .timeout_secs + .map(Duration::from_secs_f64) + .unwrap_or(default_timeout); + let mut opts = CallOptions::new(timeout); + if let Some(ref cs) = self.components { + opts = opts.with_components(cs.clone()); + } + opts + } +} + +/// Shared state for the RL HTTP facade. +#[derive(Clone)] +pub struct RlState { + pub client: Arc, +} + +impl RlState { + pub fn new(client: RlClient) -> Self { + Self { + client: Arc::new(client), + } + } +} + +/// Build the RL axum router. +/// +/// ```text +/// POST /v1/rl/engine — fan-out or direct call +/// GET /v1/rl/engine — describe: list registered methods per worker +/// ``` +pub fn rl_router(state: RlState) -> Router { + Router::new() + .route("/v1/rl/engine", post(engine_handler).get(describe_handler)) + .with_state(state) +} + +fn rl_error_response(err: anyhow::Error) -> (StatusCode, Json) { + let (status, error_type) = match err.downcast_ref::() { + Some(RlError::NoWorkers { .. }) => (StatusCode::SERVICE_UNAVAILABLE, "no_workers"), + Some(RlError::MembershipChanged { .. }) => (StatusCode::CONFLICT, "membership_changed"), + Some(RlError::UnknownInstance { .. }) => (StatusCode::CONFLICT, "unknown_instance"), + None => (StatusCode::INTERNAL_SERVER_ERROR, "fanout_failed"), + }; + + ( + status, + Json(serde_json::json!({ + "error": error_type, + "message": err.to_string(), + })), + ) +} + +fn fanout_report_to_response(report: FanoutReport) -> Json { + Json(serde_json::json!({ + "epoch": report.epoch, + "workers": report.workers, + })) +} + +/// `POST /v1/rl/engine` — fan-out or direct call. +async fn engine_handler( + State(state): State, + Json(req): Json, +) -> impl IntoResponse { + let opts = req.call_options(state.client.default_timeout()); + + let result = if let Some(id) = req.instance_id { + state + .client + .engine_call_one(&req.method, req.kwargs, id, opts) + .await + } else { + state + .client + .engine_call(&req.method, req.kwargs, opts) + .await + }; + + match result { + Ok(report) => fanout_report_to_response(report).into_response(), + Err(err) => { + let (status, body) = rl_error_response(err); + (status, body).into_response() + } + } +} + +/// `GET /v1/rl/engine` — return live worker set with registered method lists. +/// +/// Sends `{"method": "__describe__"}` to each worker's `rl_dispatch` and +/// aggregates `registered_methods` lists. +async fn describe_handler(State(state): State) -> impl IntoResponse { + let opts = CallOptions::new(state.client.default_timeout()); + let result = state + .client + .engine_call("__describe__", serde_json::json!({}), opts) + .await; + + match result { + Ok(report) => fanout_report_to_response(report).into_response(), + Err(err) => { + let (status, body) = rl_error_response(err); + (status, body).into_response() + } + } +} diff --git a/pyproject.toml b/pyproject.toml index d766e151d4f4..43ef9ad03111 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -51,7 +51,7 @@ trtllm =[ vllm = [ "uvloop", "nixl[cu12]<=0.10.1", - "vllm[flashinfer,runai,otel]==0.20.1", + "vllm[flashinfer,runai,otel]==0.20.2", # vllm-omni is installed separately in container builds (see # container/deps/vllm/install_vllm.sh). Do not add it to ai-dynamo[vllm]: # pip/uv dependency resolution for omni can override the vLLM torch stack. diff --git a/tests/rl/make_lora.py b/tests/rl/make_lora.py new file mode 100755 index 000000000000..d4a82feb76b0 --- /dev/null +++ b/tests/rl/make_lora.py @@ -0,0 +1,49 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Generate a tiny untrained LoRA adapter for Qwen3-0.6B. + +Used by the RL smoke test to exercise load_lora_adapter / unload_lora_adapter +without needing to download a pretrained adapter from HuggingFace. + +Usage: + python make_lora.py +""" + +import os +import sys + +# Force offline so we don't hit HF Hub (cached weights only). +os.environ.setdefault("HF_HUB_OFFLINE", "1") +os.environ.setdefault("TRANSFORMERS_OFFLINE", "1") + +from peft import LoraConfig, get_peft_model +from transformers import AutoModelForCausalLM + + +def main() -> None: + if len(sys.argv) != 2: + print("usage: make_lora.py ", file=sys.stderr) + sys.exit(2) + + output_dir = sys.argv[1] + model_name = "Qwen/Qwen3-0.6B" + + print(f"[make_lora] Loading base model {model_name}") + base = AutoModelForCausalLM.from_pretrained(model_name) + + # Small rank, untrained — only meant to verify the load/unload code path. + config = LoraConfig( + r=8, + lora_alpha=16, + target_modules=["q_proj", "v_proj"], + lora_dropout=0.0, + bias="none", + task_type="CAUSAL_LM", + ) + peft_model = get_peft_model(base, config) + peft_model.save_pretrained(output_dir) + print(f"[make_lora] Saved adapter to {output_dir}") + + +if __name__ == "__main__": + main() diff --git a/tests/rl/nccl_broadcaster.py b/tests/rl/nccl_broadcaster.py new file mode 100755 index 000000000000..2190f958381a --- /dev/null +++ b/tests/rl/nccl_broadcaster.py @@ -0,0 +1,148 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Minimal NCCL broadcast SENDER (rank 0) for the RL admin smoke test. + +Mirrors the wire protocol of prime_rl's `NCCLWeightBroadcastSender` so that +prime_rl's `NCCLWeightUpdateWorker` (rank >= 1) can receive the broadcast and +load it via `model.load_weights(state_iter)`. + +Coordination: + 1. Process starts and loads the model. + 2. Calls StatelessProcessGroup.create(rank=0, world_size=2) — BLOCKS until + the inference worker also calls create (via POST init_weights_update_group). + 3. Once the NCCL communicator is ready, reads a line from stdin. Sending "GO" + triggers the broadcast; anything else aborts. + +Usage: + python nccl_broadcaster.py --host 127.0.0.1 --port 29501 \\ + --world-size 2 --model Qwen/Qwen3-0.6B +""" + +from __future__ import annotations + +import argparse +import os +import pickle +import sys +from typing import Generator + +# Force offline so transformers/HF don't try to reach the hub. +os.environ.setdefault("HF_HUB_OFFLINE", "1") +os.environ.setdefault("TRANSFORMERS_OFFLINE", "1") + +import torch +from torch import Tensor +from transformers import AutoModelForCausalLM +from vllm.distributed.device_communicators.pynccl import PyNcclCommunicator +from vllm.distributed.utils import StatelessProcessGroup + +LAYER_PREFIX = "model.layers." + + +def broadcast_integer(integer: int, communicator: PyNcclCommunicator) -> None: + t = torch.tensor([integer], dtype=torch.long).cuda() + communicator.broadcast(t, src=0) + + +def broadcast_state_dict(state_dict: dict[str, Tensor], communicator: PyNcclCommunicator) -> None: + """Wire protocol matches prime_rl.trainer.rl.broadcast.nccl.broadcast_state_dict. + + Sends, in order: + 1. integer = byte-length of pickled metadata + 2. raw pickle bytes (state_tensor of uint8) + 3. for each dtype group: one concatenated tensor (flatten + cat) + """ + dtype_groups: dict[torch.dtype, list[tuple[str, Tensor]]] = {} + for key, value in state_dict.items(): + dtype_groups.setdefault(value.dtype, []).append((key, value)) + + metadata = {dt: [(k, v.shape, v.numel()) for k, v in items] for dt, items in dtype_groups.items()} + state = pickle.dumps(metadata) + size_tensor = torch.tensor([len(state)], dtype=torch.long).cuda() + communicator.broadcast(size_tensor, src=0) + state_tensor = torch.ByteTensor(list(state)).cuda() + communicator.broadcast(state_tensor, src=0) + + for dtype, items in dtype_groups.items(): + flat = [v.flatten() for _, v in items] + concatenated = torch.cat(flat) + communicator.broadcast(concatenated, src=0) + del concatenated + + +def filter_state_dict_by_layers( + state_dict: dict[str, Tensor], num_layers: int +) -> Generator[tuple[int, dict[str, Tensor]], None, None]: + """Yield (-1, non-layer weights), then (i, layer_i weights) for each i.""" + yield -1, {k: v for k, v in state_dict.items() if not k.startswith(LAYER_PREFIX)} + for i in range(num_layers): + yield i, {k: v for k, v in state_dict.items() if k.startswith(f"{LAYER_PREFIX}{i}.")} + + +def get_max_layer_num(state_dict: dict[str, Tensor]) -> int: + max_layer = -1 + for key in state_dict: + if key.startswith(LAYER_PREFIX): + tail = key[len(LAYER_PREFIX) :].split(".", 1)[0] + if tail.isdigit(): + max_layer = max(max_layer, int(tail)) + return max_layer + 1 + + +def main() -> None: + ap = argparse.ArgumentParser() + ap.add_argument("--host", default="127.0.0.1") + ap.add_argument("--port", type=int, required=True, + help="TCP port for StatelessProcessGroup rendezvous. " + "Must be dynamically allocated by the caller to avoid " + "collisions between concurrent smoke-test runs.") + ap.add_argument("--world-size", type=int, default=2, + help="Total ranks (trainer + inference workers). Default 2 = 1 trainer + 1 worker.") + ap.add_argument("--model", default="Qwen/Qwen3-0.6B") + ap.add_argument("--dtype", default="bfloat16") + ap.add_argument("--timeout", type=int, default=120, help="StatelessProcessGroup store timeout (s)") + args = ap.parse_args() + + dtype = getattr(torch, args.dtype) + + print(f"[broadcaster] Loading {args.model} (dtype={args.dtype}) on cuda ...", flush=True) + model = AutoModelForCausalLM.from_pretrained(args.model, dtype=dtype).cuda().eval() + sd = {k: v.detach() for k, v in model.state_dict().items()} + num_layers = get_max_layer_num(sd) + print(f"[broadcaster] Model loaded: {len(sd)} tensors, {num_layers} layers", flush=True) + + print( + f"[broadcaster] Creating StatelessProcessGroup({args.host}:{args.port}, " + f"rank=0, world_size={args.world_size}) — blocks until peers join ...", + flush=True, + ) + pg = StatelessProcessGroup.create( + host=args.host, + port=args.port, + rank=0, + world_size=args.world_size, + store_timeout=args.timeout, + ) + print("[broadcaster] StatelessProcessGroup created — all peers joined", flush=True) + communicator = PyNcclCommunicator(pg, device=torch.cuda.current_device()) + print("[broadcaster] PyNcclCommunicator ready", flush=True) + + print("[broadcaster] Waiting for 'GO' on stdin ...", flush=True) + line = sys.stdin.readline().strip() + if line != "GO": + print(f"[broadcaster] aborting: expected 'GO', got {line!r}", flush=True) + sys.exit(2) + + num_state_dict_to_send = num_layers + 1 # non-layer chunk + one per layer + print(f"[broadcaster] Broadcasting {num_state_dict_to_send} state-dict chunks ...", flush=True) + broadcast_integer(num_state_dict_to_send, communicator) + for layer_idx, layer_sd in filter_state_dict_by_layers(sd, num_layers): + layer_sd = {k: v.to(dtype) for k, v in layer_sd.items()} + broadcast_state_dict(layer_sd, communicator) + print(f"[broadcaster] sent chunk {layer_idx} ({len(layer_sd)} tensors)", flush=True) + + print("[broadcaster] Broadcast complete", flush=True) + + +if __name__ == "__main__": + main() diff --git a/tests/rl/probe_engine_data.sh b/tests/rl/probe_engine_data.sh new file mode 100755 index 000000000000..a9cd9a6e60ca --- /dev/null +++ b/tests/rl/probe_engine_data.sh @@ -0,0 +1,174 @@ +#!/bin/bash +# Server-side TITO probe: validates nvext.engine_data emission from the +# vLLM backend after the rl-sdk-2 changes. +# +# What it tests (single roundtrip, no orchestrator/trainer): +# 1. PASSTHROUGH_EXTRA_FIELDS accepts cache_salt + stop_token_ids +# without 400 Bad Request. +# 2. Request `nvext.token_data=[...]` skips server-side tokenization +# (validated indirectly: usage.prompt_tokens == len(token_data)). +# 3. Request `nvext.extra_fields=["engine_data"]` causes the response +# to include `nvext.engine_data.completion_token_ids` as the +# exact engine-emitted IDs. +# 4. `nvext.engine_data.completion_logprobs` is a flat list[float] +# indexed by sampled token (one float per completion_token_id). +# +# Pass criteria (asserted by the python block): +# * status == 200 +# * len(engine_data.completion_token_ids) == usage.completion_tokens +# * (optional) len(completion_logprobs) >= 1 if logprobs were requested + +set -euo pipefail + +DYNAMO_VENV=/home/biswaranjanp/dev/rl/dynamo/.venv +WORKDIR=/tmp/probe-engine-data +mkdir -p "$WORKDIR" + +# Clean up any leftover dynamo processes +pkill -9 -f "dynamo.vllm" 2>/dev/null || true +pkill -9 -f "dynamo.frontend" 2>/dev/null || true +pkill -9 -f "VLLM::EngineCore" 2>/dev/null || true +sleep 3 + +source "$DYNAMO_VENV/bin/activate" + +# Frontend on :8000 (no DYN_ENABLE_RL needed for engine_data path — gating +# is purely from nvext.extra_fields=["engine_data"] on the request) +CUDA_VISIBLE_DEVICES="" \ + nohup python3 -m dynamo.frontend --http-port 8000 \ + > "$WORKDIR/frontend.log" 2>&1 & +FRONTEND_PID=$! +echo "[probe] frontend PID=$FRONTEND_PID" +sleep 5 + +# vLLM worker - Qwen3-0.6B, eager, single GPU, small mem +CUDA_VISIBLE_DEVICES=0 \ + nohup python3 -m dynamo.vllm \ + --model Qwen/Qwen3-0.6B \ + --served-model-name Qwen/Qwen3-0.6B \ + --enforce-eager \ + --max-model-len 2048 \ + --max-num-seqs 4 \ + --gpu-memory-utilization 0.30 \ + > "$WORKDIR/vllm_worker.log" 2>&1 & +WORKER_PID=$! +echo "[probe] worker PID=$WORKER_PID" + +cleanup() { + echo "[probe] cleanup" + kill "$FRONTEND_PID" "$WORKER_PID" 2>/dev/null || true + pkill -9 -f "VLLM::EngineCore" 2>/dev/null || true +} +trap cleanup EXIT INT TERM + +echo "[probe] waiting for /v1/models to register Qwen/Qwen3-0.6B..." +for i in $(seq 1 60); do + if curl -fs http://localhost:8000/v1/models 2>/dev/null | grep -q "Qwen/Qwen3-0.6B"; then + echo "[probe] worker ready after ${i}*5s" + break + fi + sleep 5 +done + +# Pre-tokenize a known prompt ourselves so we can compare prompt_tokens later. +python3 <<'PYEOF' +import json, sys, urllib.request + +# Hardcoded short prompt tokenization for Qwen3 chat template. +# We'll send a placeholder "(token-in mode)" string in messages so the +# preprocessor still has SOMETHING to apply chat-template selection on, +# but the actual prompt tokens override via nvext.token_data. +# These IDs are arbitrary valid Qwen3 vocab tokens; the server should +# treat them as the prompt and emit usage.prompt_tokens == 8. +prompt_ids = [151644, 8948, 198, 9707, 11, 1879, 0, 151645] # 8 tokens + +body = { + "model": "Qwen/Qwen3-0.6B", + "messages": [{"role": "user", "content": "(token-in mode)"}], + "max_completion_tokens": 12, + "temperature": 0.7, + "logprobs": True, + # Test PASSTHROUGH for both cache_salt + stop_token_ids: + "cache_salt": "probe_step_1", + "stop_token_ids": [151643], # Qwen3 <|endoftext|> + "nvext": { + "token_data": prompt_ids, + "extra_fields": ["engine_data"], + # Also test nvext.cache_salt as the canonical location: + "cache_salt": "probe_step_1_nvext", + }, +} + +req = urllib.request.Request( + "http://localhost:8000/v1/chat/completions", + data=json.dumps(body).encode(), + headers={"Content-Type": "application/json"}, +) +try: + with urllib.request.urlopen(req, timeout=120) as resp: + status = resp.status + payload = json.loads(resp.read().decode()) +except urllib.error.HTTPError as e: + print(f"[probe] HTTP {e.code}: {e.read().decode()}", file=sys.stderr) + sys.exit(1) + +print(f"[probe] HTTP {status}") +print(json.dumps(payload, indent=2)) + +# Assertions +assert status == 200, f"expected 200, got {status}" + +nvext_resp = payload.get("nvext") +assert isinstance(nvext_resp, dict), f"missing response.nvext (got {type(nvext_resp)})" + +engine_data = nvext_resp.get("engine_data") +assert isinstance(engine_data, dict), ( + f"missing response.nvext.engine_data; got nvext={nvext_resp}" +) + +ctids = engine_data.get("completion_token_ids") +assert isinstance(ctids, list) and ctids, ( + f"missing/empty completion_token_ids: {ctids}" +) + +usage = payload["usage"] +assert len(ctids) == usage["completion_tokens"], ( + f"completion_token_ids length ({len(ctids)}) != " + f"usage.completion_tokens ({usage['completion_tokens']})" +) + +# usage.prompt_tokens should equal what we sent via nvext.token_data +assert usage["prompt_tokens"] == len(prompt_ids), ( + f"prompt_tokens={usage['prompt_tokens']} != sent token_data len={len(prompt_ids)}; " + f"server did NOT consume nvext.token_data" +) + +clp = engine_data.get("completion_logprobs") +if clp is not None: + assert isinstance(clp, list), f"completion_logprobs not a list: {type(clp)}" + assert all(isinstance(x, (int, float)) for x in clp), ( + f"completion_logprobs not flat list[float]: {[type(x).__name__ for x in clp][:5]}" + ) + assert len(clp) == len(ctids), ( + f"len(completion_logprobs)={len(clp)} != len(completion_token_ids)={len(ctids)}" + ) + +print() +print("[probe] PASS — all engine_data assertions satisfied") +print(f" completion_token_ids: {ctids}") +print(f" completion_logprobs (first 5): {clp[:5] if clp else 'none'}") +print(f" usage: {usage}") +PYEOF + +PROBE_EXIT=$? +echo +if [ "$PROBE_EXIT" -eq 0 ]; then + echo "[probe] === PASS ===" +else + echo "[probe] === FAIL (exit=$PROBE_EXIT) ===" + echo "--- frontend.log tail ---" + tail -30 "$WORKDIR/frontend.log" + echo "--- vllm_worker.log tail ---" + tail -30 "$WORKDIR/vllm_worker.log" +fi +exit $PROBE_EXIT diff --git a/tests/rl/smoke_test.sh b/tests/rl/smoke_test.sh new file mode 100755 index 000000000000..aa67bc036f1d --- /dev/null +++ b/tests/rl/smoke_test.sh @@ -0,0 +1,280 @@ +#!/bin/bash +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Smoke test for the RL admin control plane. +# +# Starts: +# 1. NATS server (discovery plane) +# 2. Dynamo frontend (with DYN_ENABLE_RL_ENDPOINTS=true) +# 3. Dynamo vLLM worker (Qwen3-0.6B + FileSystemWeightUpdateWorker) +# +# Then exercises: GET /v1/rl/engine, POST /v1/rl/engine pause_generation, +# POST /v1/rl/engine update_weights_from_disk, POST /v1/rl/engine resume_generation +# +# Usage: +# cd +# source dynamo/bin/activate +# export PRIME_RL_SRC=/path/to/prime-rl/src +# bash tests/rl/smoke_test.sh [] + +set -euo pipefail + +BGPIDS=() +# Track background PIDs and kill only them on exit (not the shell itself). +cleanup() { + trap - EXIT INT TERM + echo "[smoke] Cleaning up..." + for pid in "${BGPIDS[@]:-}"; do + kill "$pid" 2>/dev/null || true + done +} +trap cleanup EXIT INT TERM + +MODEL="${1:-Qwen/Qwen3-0.6B}" +HTTP_PORT="${DYN_HTTP_PORT:-8000}" +NATS_PORT="${NATS_PORT:-4222}" +# prime_rl source must be on PYTHONPATH so the spawned vLLM worker subprocess +# can import prime_rl.inference.vllm.worker.filesystem.FileSystemWeightUpdateWorker. +# Set PRIME_RL_SRC to the prime-rl src directory before running this script. +: "${PRIME_RL_SRC:?Set PRIME_RL_SRC to the prime-rl src directory (e.g. export PRIME_RL_SRC=/path/to/prime-rl/src)}" +PYTHON="${SMOKE_PYTHON:-python}" +LOG_DIR="${TMPDIR:-/tmp}/dynamo-rl-smoke-$$" +mkdir -p "$LOG_DIR" + +WEIGHT_DIR="${LOG_DIR}/weights" +mkdir -p "$WEIGHT_DIR" + +echo "[smoke] Log dir: $LOG_DIR" +echo "[smoke] Model: $MODEL" +echo "[smoke] Frontend port: $HTTP_PORT" + +# --------------------------------------------------------------------------- +# 1. NATS +# --------------------------------------------------------------------------- +if nc -z localhost "$NATS_PORT" 2>/dev/null; then + echo "[smoke] NATS already running on port $NATS_PORT — skipping start" +else + nats-server -p "$NATS_PORT" -l "$LOG_DIR/nats.log" & + NATS_PID=$! + BGPIDS+=("$NATS_PID") + echo "[smoke] NATS started (pid=$NATS_PID)" + sleep 1 +fi + +# --------------------------------------------------------------------------- +# 2. Dynamo frontend (RL enabled on same port as OpenAI-compat) +# --------------------------------------------------------------------------- +HF_HUB_OFFLINE=1 \ + TRANSFORMERS_OFFLINE=1 \ + DYN_ENABLE_RL_ENDPOINTS=true \ + DYN_HTTP_PORT="$HTTP_PORT" \ + python -m dynamo.frontend \ + > "$LOG_DIR/frontend.log" 2>&1 & +FRONTEND_PID=$! +BGPIDS+=("$FRONTEND_PID") +echo "[smoke] Frontend started (pid=$FRONTEND_PID)" + +# --------------------------------------------------------------------------- +# 3. Dynamo vLLM worker with FileSystemWeightUpdateWorker extension +# --------------------------------------------------------------------------- +HF_HUB_OFFLINE=1 \ + TRANSFORMERS_OFFLINE=1 \ + PYTHONPATH="${PRIME_RL_SRC}${PYTHONPATH:+:$PYTHONPATH}" \ + DYN_SYSTEM_PORT=${DYN_SYSTEM_PORT:-8081} \ + python -m dynamo.vllm \ + --model "$MODEL" \ + --enforce-eager \ + --max-model-len 2048 \ + --max-num-seqs 2 \ + --enable-rl \ + --worker-extension-cls prime_rl.inference.vllm.worker.filesystem.FileSystemWeightUpdateWorker \ + > "$LOG_DIR/worker.log" 2>&1 & +WORKER_PID=$! +BGPIDS+=("$WORKER_PID") +echo "[smoke] Worker started (pid=$WORKER_PID, log=$LOG_DIR/worker.log)" + +# --------------------------------------------------------------------------- +# 4. Wait for the RL endpoint to be live +# --------------------------------------------------------------------------- +echo "[smoke] Waiting for /v1/rl/engine to become live..." +DEADLINE=$(( $(date +%s) + 180 )) +while true; do + if curl -sf -X GET "http://localhost:${HTTP_PORT}/v1/rl/engine" -o /dev/null 2>&1; then + echo "[smoke] RL endpoint is live" + break + fi + if [ "$(date +%s)" -ge "$DEADLINE" ]; then + echo "[smoke] TIMEOUT: RL endpoint not live after 180s" + echo "=== frontend.log ===" + tail -30 "$LOG_DIR/frontend.log" + echo "=== worker.log ===" + tail -30 "$LOG_DIR/worker.log" + exit 1 + fi + sleep 3 +done + +# --------------------------------------------------------------------------- +# 5. GET /v1/rl/engine — describe registered methods +# --------------------------------------------------------------------------- +echo "" +echo "[smoke] === GET /v1/rl/engine (describe) ===" +DESCRIBE=$(curl -sf -X GET "http://localhost:${HTTP_PORT}/v1/rl/engine") +echo "$DESCRIBE" | python -m json.tool +echo "" + +# --------------------------------------------------------------------------- +# 6. pause_generation +# --------------------------------------------------------------------------- +echo "[smoke] === POST pause_generation ===" +PAUSE=$(curl -sf -X POST "http://localhost:${HTTP_PORT}/v1/rl/engine" \ + -H "Content-Type: application/json" \ + -d '{"method": "pause_generation", "kwargs": {"abort_requests": true, "clear_cache": false}}') +echo "$PAUSE" | python -m json.tool +echo "" + +# Verify all workers paused (status ok) +if echo "$PAUSE" | python -c " +import sys, json +data = json.load(sys.stdin) +workers = data.get('workers', []) +failed = [w for w in workers if w.get('status') != 'ok'] +if failed: + print('FAIL: some workers not ok:', failed) + sys.exit(1) +elif not workers: + print('FAIL: no workers reported') + sys.exit(1) +else: + print(f'PASS: {len(workers)} worker(s) paused ok') +"; then + : +else + echo "[smoke] FAIL: pause_generation" + exit 1 +fi + +# --------------------------------------------------------------------------- +# 7. update_weights_from_disk (same path — no-op but tests round-trip) +# --------------------------------------------------------------------------- +echo "[smoke] === POST update_weights_from_disk ===" +# Use the original model path from HF cache as weight_path so load is valid +MODEL_CACHE=$(python -c " +import huggingface_hub, os +try: + p = huggingface_hub.snapshot_download('$MODEL', local_files_only=True) + print(p) +except Exception as e: + print(os.path.expanduser('~/.cache/huggingface/hub')) +" 2>/dev/null) +echo "[smoke] weight path: $MODEL_CACHE" + +UPDATE=$(curl -sf -X POST "http://localhost:${HTTP_PORT}/v1/rl/engine" \ + -H "Content-Type: application/json" \ + --max-time 300 \ + -d "{\"method\": \"update_weights_from_disk\", + \"kwargs\": {\"model_path\": \"${MODEL_CACHE}\", + \"weight_version\": \"smoke_v1\", + \"engine_rpc\": \"update_weights_from_path\"}, + \"timeout_secs\": 240}") +echo "$UPDATE" | python -m json.tool +echo "" + +if echo "$UPDATE" | python -c " +import sys, json +data = json.load(sys.stdin) +workers = data.get('workers', []) +failed = [w for w in workers if w.get('status') != 'ok'] +if failed: + print('FAIL: workers not ok:', failed) + sys.exit(1) +elif not workers: + print('FAIL: no workers reported') + sys.exit(1) +else: + print(f'PASS: {len(workers)} worker(s) updated ok') +"; then + : +else + echo "[smoke] FAIL: update_weights_from_disk" + exit 1 +fi + +# --------------------------------------------------------------------------- +# 8. get_weight_version — verify version updated +# --------------------------------------------------------------------------- +echo "[smoke] === POST get_weight_version ===" +VER=$(curl -sf -X POST "http://localhost:${HTTP_PORT}/v1/rl/engine" \ + -H "Content-Type: application/json" \ + -d '{"method": "get_weight_version"}') +echo "$VER" | python -m json.tool + +if echo "$VER" | python -c " +import sys, json +data = json.load(sys.stdin) +for w in data.get('workers', []): + resp = w.get('response', {}) + version = resp.get('version', resp.get('weight_version', '')) + if version != 'smoke_v1': + print(f'FAIL: expected smoke_v1 got {version!r}') + sys.exit(1) +print('PASS: version=smoke_v1') +"; then + : +else + echo "[smoke] WARN: weight version check failed (non-fatal)" +fi +echo "" + +# --------------------------------------------------------------------------- +# 9. resume_generation +# --------------------------------------------------------------------------- +echo "[smoke] === POST resume_generation ===" +RESUME=$(curl -sf -X POST "http://localhost:${HTTP_PORT}/v1/rl/engine" \ + -H "Content-Type: application/json" \ + -d '{"method": "resume_generation"}') +echo "$RESUME" | python -m json.tool +echo "" + +if echo "$RESUME" | python -c " +import sys, json +data = json.load(sys.stdin) +workers = data.get('workers', []) +failed = [w for w in workers if w.get('status') != 'ok'] +if failed: + print('FAIL: workers not ok:', failed) + sys.exit(1) +elif not workers: + print('FAIL: no workers reported') + sys.exit(1) +else: + print(f'PASS: {len(workers)} worker(s) resumed ok') +"; then + : +else + echo "[smoke] FAIL: resume_generation" + exit 1 +fi + +# --------------------------------------------------------------------------- +# 10. Quick inference check — verify model still serves requests +# --------------------------------------------------------------------------- +echo "[smoke] === Quick inference check ===" +INF=$(curl -sf -X POST "http://localhost:${HTTP_PORT}/v1/chat/completions" \ + -H "Content-Type: application/json" \ + --max-time 60 \ + -d "{\"model\": \"${MODEL}\", + \"messages\": [{\"role\": \"user\", \"content\": \"Say: hello\"}], + \"max_tokens\": 8, \"stream\": false}" 2>&1 || true) +if echo "$INF" | grep -q '"choices"'; then + echo "PASS: inference working after weight update" +else + echo "WARN: inference check inconclusive (model may not be ready yet)" + echo "$INF" | head -5 +fi + +echo "" +echo "========================================" +echo "[smoke] ALL TESTS PASSED" +echo "========================================" diff --git a/tests/rl/smoke_test_lora.sh b/tests/rl/smoke_test_lora.sh new file mode 100755 index 000000000000..61ac81c3ecd9 --- /dev/null +++ b/tests/rl/smoke_test_lora.sh @@ -0,0 +1,268 @@ +#!/bin/bash +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# LoRA smoke test for the RL admin control plane. +# +# Verifies load_lora_adapter / unload_lora_adapter via POST /v1/rl/engine. +# +# Starts: +# 1. NATS (skipped if already running) +# 2. Dynamo frontend (DYN_ENABLE_RL_ENDPOINTS=true) +# 3. Dynamo vLLM worker with --enable-lora and FileSystemWeightUpdateWorker +# +# Then exercises: +# GET /v1/rl/engine → describe (incl. load/unload_lora_adapter) +# POST /v1/rl/engine load_lora_adapter → load tiny untrained LoRA +# POST /v1/chat/completions model= → inference uses LoRA +# POST /v1/rl/engine unload_lora_adapter → unload +# +# Usage: +# cd /home/biswaranjanp/dev/rl/dynamo +# source dynamo/bin/activate +# bash tests/rl/smoke_test_lora.sh [] + +set -euo pipefail + +BGPIDS=() +cleanup() { + trap - EXIT INT TERM + echo "[smoke-lora] Cleaning up..." + for pid in "${BGPIDS[@]:-}"; do + kill "$pid" 2>/dev/null || true + done +} +trap cleanup EXIT INT TERM + +MODEL="${1:-Qwen/Qwen3-0.6B}" +HTTP_PORT="${DYN_HTTP_PORT:-8000}" +NATS_PORT="${NATS_PORT:-4222}" +# Set PRIME_RL_SRC to the prime-rl src directory before running this script. +: "${PRIME_RL_SRC:?Set PRIME_RL_SRC to the prime-rl src directory (e.g. export PRIME_RL_SRC=/path/to/prime-rl/src)}" + +LOG_DIR="${TMPDIR:-/tmp}/dynamo-rl-smoke-lora-$$" +mkdir -p "$LOG_DIR" + +LORA_NAME="${LORA_NAME:-qwen3-tiny-lora}" +LORA_DIR="${LORA_DIR:-${LOG_DIR}/adapter}" + +echo "[smoke-lora] Log dir: $LOG_DIR" +echo "[smoke-lora] Model: $MODEL" +echo "[smoke-lora] LoRA name: $LORA_NAME" +echo "[smoke-lora] LoRA dir: $LORA_DIR" + +# --------------------------------------------------------------------------- +# 0. Build the tiny LoRA adapter (peft must be installed in the venv) +# --------------------------------------------------------------------------- +echo "[smoke-lora] Building tiny LoRA adapter..." +python "$(dirname "$0")/make_lora.py" "$LORA_DIR" 2>&1 | tail -3 +if [ ! -f "$LORA_DIR/adapter_config.json" ]; then + echo "[smoke-lora] FAIL: LoRA adapter not created at $LORA_DIR" + exit 1 +fi +echo "[smoke-lora] LoRA adapter ready" + +# --------------------------------------------------------------------------- +# 1. NATS +# --------------------------------------------------------------------------- +if nc -z localhost "$NATS_PORT" 2>/dev/null; then + echo "[smoke-lora] NATS already running on port $NATS_PORT — skipping start" +else + nats-server -p "$NATS_PORT" -l "$LOG_DIR/nats.log" & + NATS_PID=$! + BGPIDS+=("$NATS_PID") + echo "[smoke-lora] NATS started (pid=$NATS_PID)" + sleep 1 +fi + +# --------------------------------------------------------------------------- +# 2. Dynamo frontend +# --------------------------------------------------------------------------- +HF_HUB_OFFLINE=1 \ + TRANSFORMERS_OFFLINE=1 \ + DYN_ENABLE_RL_ENDPOINTS=true \ + DYN_HTTP_PORT="$HTTP_PORT" \ + python -m dynamo.frontend \ + > "$LOG_DIR/frontend.log" 2>&1 & +FRONTEND_PID=$! +BGPIDS+=("$FRONTEND_PID") +echo "[smoke-lora] Frontend started (pid=$FRONTEND_PID)" + +# --------------------------------------------------------------------------- +# 3. Dynamo vLLM worker with --enable-lora +# --------------------------------------------------------------------------- +HF_HUB_OFFLINE=1 \ + TRANSFORMERS_OFFLINE=1 \ + PYTHONPATH="${PRIME_RL_SRC}${PYTHONPATH:+:$PYTHONPATH}" \ + DYN_SYSTEM_PORT=${DYN_SYSTEM_PORT:-8081} \ + python -m dynamo.vllm \ + --model "$MODEL" \ + --enforce-eager \ + --max-model-len 2048 \ + --max-num-seqs 2 \ + --enable-rl \ + --enable-lora \ + --max-lora-rank 32 \ + --max-loras 4 \ + --worker-extension-cls prime_rl.inference.vllm.worker.filesystem.FileSystemWeightUpdateWorker \ + > "$LOG_DIR/worker.log" 2>&1 & +WORKER_PID=$! +BGPIDS+=("$WORKER_PID") +echo "[smoke-lora] Worker started (pid=$WORKER_PID, log=$LOG_DIR/worker.log)" + +# --------------------------------------------------------------------------- +# 4. Wait for the RL endpoint to be live +# --------------------------------------------------------------------------- +echo "[smoke-lora] Waiting for /v1/rl/engine to become live..." +DEADLINE=$(( $(date +%s) + 180 )) +while true; do + if curl -sf -X GET "http://localhost:${HTTP_PORT}/v1/rl/engine" -o /dev/null 2>&1; then + echo "[smoke-lora] RL endpoint is live" + break + fi + if [ "$(date +%s)" -ge "$DEADLINE" ]; then + echo "[smoke-lora] TIMEOUT: RL endpoint not live after 180s" + echo "=== frontend.log ===" + tail -30 "$LOG_DIR/frontend.log" + echo "=== worker.log ===" + tail -30 "$LOG_DIR/worker.log" + exit 1 + fi + sleep 3 +done + +# --------------------------------------------------------------------------- +# 5. GET /v1/rl/engine — verify load_lora_adapter and unload_lora_adapter are registered +# --------------------------------------------------------------------------- +echo "" +echo "[smoke-lora] === GET /v1/rl/engine (describe) ===" +DESCRIBE=$(curl -sf -X GET "http://localhost:${HTTP_PORT}/v1/rl/engine") +echo "$DESCRIBE" | python -m json.tool + +if echo "$DESCRIBE" | python -c " +import sys, json +data = json.load(sys.stdin) +workers = data.get('workers', []) +if not workers: + print('FAIL: no workers reported'); sys.exit(1) +methods = set(workers[0].get('response', {}).get('registered_methods', [])) +required = {'load_lora_adapter', 'unload_lora_adapter'} +missing = required - methods +if missing: + print(f'FAIL: missing methods: {missing}'); sys.exit(1) +print(f'PASS: lora methods registered on {len(workers)} worker(s)') +"; then + : +else + echo "[smoke-lora] FAIL: describe missing LoRA methods" + exit 1 +fi + +# --------------------------------------------------------------------------- +# 6. POST load_lora_adapter +# --------------------------------------------------------------------------- +echo "" +echo "[smoke-lora] === POST load_lora_adapter ===" +LOAD=$(curl -sf -X POST "http://localhost:${HTTP_PORT}/v1/rl/engine" \ + -H "Content-Type: application/json" \ + --max-time 60 \ + -d "{\"method\": \"load_lora_adapter\", + \"kwargs\": {\"lora_name\": \"${LORA_NAME}\", + \"lora_path\": \"${LORA_DIR}\"}, + \"timeout_secs\": 60}") +echo "$LOAD" | python -m json.tool + +if echo "$LOAD" | python -c " +import sys, json +data = json.load(sys.stdin) +workers = data.get('workers', []) +failed = [w for w in workers if w.get('status') != 'ok'] +if failed: + print('FAIL: workers not ok:', failed); sys.exit(1) +elif not workers: + print('FAIL: no workers reported'); sys.exit(1) +print(f'PASS: {len(workers)} worker(s) loaded LoRA ok') +"; then + : +else + echo "[smoke-lora] FAIL: load_lora_adapter" + exit 1 +fi + +# Give discovery a moment to propagate the new LoRA model registration. +sleep 2 + +# --------------------------------------------------------------------------- +# 7. Inference with model= +# --------------------------------------------------------------------------- +echo "" +echo "[smoke-lora] === Inference with model=${LORA_NAME} ===" +INF=$(curl -sf -X POST "http://localhost:${HTTP_PORT}/v1/chat/completions" \ + -H "Content-Type: application/json" \ + --max-time 60 \ + -d "{\"model\": \"${LORA_NAME}\", + \"messages\": [{\"role\": \"user\", \"content\": \"Say: hello\"}], + \"max_tokens\": 8, \"stream\": false}" 2>&1 || true) +if echo "$INF" | grep -q '"choices"'; then + echo "PASS: inference with LoRA model name returned choices" +else + echo "WARN: LoRA inference inconclusive — first 5 lines:" + echo "$INF" | head -5 +fi + +# --------------------------------------------------------------------------- +# 8. POST unload_lora_adapter +# --------------------------------------------------------------------------- +echo "" +echo "[smoke-lora] === POST unload_lora_adapter ===" +UNLOAD=$(curl -sf -X POST "http://localhost:${HTTP_PORT}/v1/rl/engine" \ + -H "Content-Type: application/json" \ + -d "{\"method\": \"unload_lora_adapter\", + \"kwargs\": {\"lora_name\": \"${LORA_NAME}\"}}") +echo "$UNLOAD" | python -m json.tool + +if echo "$UNLOAD" | python -c " +import sys, json +data = json.load(sys.stdin) +workers = data.get('workers', []) +failed = [w for w in workers if w.get('status') != 'ok'] +if failed: + print('FAIL: workers not ok:', failed); sys.exit(1) +elif not workers: + print('FAIL: no workers reported'); sys.exit(1) +print(f'PASS: {len(workers)} worker(s) unloaded LoRA ok') +"; then + : +else + echo "[smoke-lora] FAIL: unload_lora_adapter" + exit 1 +fi + +# --------------------------------------------------------------------------- +# 9. Idempotency: second unload should be a no-op success +# --------------------------------------------------------------------------- +echo "" +echo "[smoke-lora] === POST unload_lora_adapter (idempotency check) ===" +UNLOAD2=$(curl -sf -X POST "http://localhost:${HTTP_PORT}/v1/rl/engine" \ + -H "Content-Type: application/json" \ + -d "{\"method\": \"unload_lora_adapter\", + \"kwargs\": {\"lora_name\": \"${LORA_NAME}\"}}") +if echo "$UNLOAD2" | python -c " +import sys, json +data = json.load(sys.stdin) +workers = data.get('workers', []) +failed = [w for w in workers if w.get('status') != 'ok'] +if failed: + print('FAIL: idempotent unload should be ok:', failed); sys.exit(1) +print('PASS: idempotent unload returns ok') +"; then + : +else + echo "[smoke-lora] FAIL: idempotent unload_lora_adapter" + exit 1 +fi + +echo "" +echo "========================================" +echo "[smoke-lora] ALL TESTS PASSED" +echo "========================================" diff --git a/tests/rl/smoke_test_nccl.sh b/tests/rl/smoke_test_nccl.sh new file mode 100755 index 000000000000..d64f28fc09d2 --- /dev/null +++ b/tests/rl/smoke_test_nccl.sh @@ -0,0 +1,384 @@ +#!/bin/bash +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# E2E smoke test for NCCL-based weight updates via the RL admin control plane. +# +# Architecture: +# NCCL group (world_size = 2) +# ┌───────────────────────────┐ ┌────────────────────────────────┐ +# │ Dynamo frontend (port 8000)│ POST /v1/rl/ │ NCCLWeightUpdateWorker │ +# │ /v1/rl/engine fan-out ├──── engine ──▶│ rank = 1, GPU 0 │ +# └───────────────────────────┘ └────────────────────────────────┘ +# ▲ ▲ +# │ HTTP │ NCCL broadcast (src=0) +# │ │ +# ┌───────┴───────────────────┐ ┌──────────┴────────────────────┐ +# │ This test script │ start + │ nccl_broadcaster.py │ +# │ (orchestrator) │── stdin "GO" ─▶│ rank = 0, GPU 0 │ +# └───────────────────────────┘ └───────────────────────────────┘ +# +# Timing protocol: +# 1. Start frontend + worker (--worker-extension-cls NCCLWeightUpdateWorker) +# 2. Start broadcaster subprocess → it loads model, calls +# StatelessProcessGroup.create(rank=0, world_size=2) and BLOCKS. +# 3. POST init_weights_update_group on worker → worker calls +# StatelessProcessGroup.create(rank=1, world_size=2). Both ranks rendezvous +# and the NCCL communicator comes up on each side. +# 4. POST update_weights_from_distributed (in background — it blocks waiting +# for the broadcast). +# 5. Send "GO" to the broadcaster's stdin. It broadcasts the state dict; +# the worker receives and loads. Both calls return. +# 6. Verify worker status=ok, then run an inference to confirm the model +# still serves after the weight swap. +# +# Usage: +# cd +# source dynamo/bin/activate +# export PRIME_RL_SRC=/path/to/prime-rl/src +# bash tests/rl/smoke_test_nccl.sh [] + +set -euo pipefail + +BGPIDS=() +cleanup() { + trap - EXIT INT TERM + echo "[smoke-nccl] Cleaning up..." + for pid in "${BGPIDS[@]:-}"; do + kill "$pid" 2>/dev/null || true + done +} +trap cleanup EXIT INT TERM + +MODEL="${1:-Qwen/Qwen3-0.6B}" +HTTP_PORT="${DYN_HTTP_PORT:-8000}" +NATS_PORT="${NATS_PORT:-4222}" +NCCL_HOST="${NCCL_HOST:-127.0.0.1}" +NCCL_PORT="${NCCL_PORT:-29501}" +# Set PRIME_RL_SRC to the prime-rl src directory before running this script. +: "${PRIME_RL_SRC:?Set PRIME_RL_SRC to the prime-rl src directory (e.g. export PRIME_RL_SRC=/path/to/prime-rl/src)}" + +LOG_DIR="${TMPDIR:-/tmp}/dynamo-rl-smoke-nccl-$$" +mkdir -p "$LOG_DIR" + +# NCCL fundamentally requires one GPU per rank. With a single GPU, the +# trainer (rank 0) and inference worker (rank 1) collide at PyNcclCommunicator +# init. Detect that up front and choose between full E2E vs wire-path-only. +GPU_COUNT=$(nvidia-smi --query-gpu=name --format=csv,noheader 2>/dev/null | wc -l) +if [ "$GPU_COUNT" -ge 2 ]; then + FULL_E2E=1 + BROADCASTER_GPU="${BROADCASTER_GPU:-1}" + echo "[smoke-nccl] Detected $GPU_COUNT GPUs — running FULL E2E (broadcaster on GPU $BROADCASTER_GPU)" +else + FULL_E2E=0 + BROADCASTER_GPU="" + echo "[smoke-nccl] Only 1 GPU detected — running WIRE-PATH test (NCCL E2E requires >=2 GPUs)" +fi + +echo "[smoke-nccl] Log dir: $LOG_DIR" +echo "[smoke-nccl] Model: $MODEL" +echo "[smoke-nccl] NCCL bind: $NCCL_HOST:$NCCL_PORT world_size=2 (rank0=trainer, rank1=worker)" + +check_workers_ok() { + local payload="$1" + local label="$2" + echo "$payload" | python -c " +import sys, json +data = json.load(sys.stdin) +workers = data.get('workers', []) +if not workers: + print('FAIL: ${label}: no workers reported'); sys.exit(1) +bad = [] +for w in workers: + if w.get('status') != 'ok': + bad.append({'dispatch_status': w.get('status'), 'error': w.get('error')}) + continue + resp = w.get('response', {}) + if isinstance(resp, dict) and resp.get('status') == 'error': + bad.append({'handler_status': 'error', 'message': resp.get('message')}) +if bad: + print(f'FAIL: ${label}: {bad}'); sys.exit(1) +print(f'PASS: ${label}') +" || return 1 + return 0 +} + +# --------------------------------------------------------------------------- +# 1. NATS +# --------------------------------------------------------------------------- +if nc -z localhost "$NATS_PORT" 2>/dev/null; then + echo "[smoke-nccl] NATS already running on port $NATS_PORT — skipping start" +else + nats-server -p "$NATS_PORT" -l "$LOG_DIR/nats.log" & + NATS_PID=$! + BGPIDS+=("$NATS_PID") + echo "[smoke-nccl] NATS started (pid=$NATS_PID)" + sleep 1 +fi + +# --------------------------------------------------------------------------- +# 2. Dynamo frontend +# --------------------------------------------------------------------------- +HF_HUB_OFFLINE=1 \ + TRANSFORMERS_OFFLINE=1 \ + DYN_ENABLE_RL_ENDPOINTS=true \ + DYN_HTTP_PORT="$HTTP_PORT" \ + python -m dynamo.frontend \ + > "$LOG_DIR/frontend.log" 2>&1 & +FRONTEND_PID=$! +BGPIDS+=("$FRONTEND_PID") +echo "[smoke-nccl] Frontend started (pid=$FRONTEND_PID)" + +# --------------------------------------------------------------------------- +# 3. vLLM worker with NCCLWeightUpdateWorker extension +# --------------------------------------------------------------------------- +HF_HUB_OFFLINE=1 \ + TRANSFORMERS_OFFLINE=1 \ + PYTHONPATH="${PRIME_RL_SRC}${PYTHONPATH:+:$PYTHONPATH}" \ + DYN_SYSTEM_PORT=${DYN_SYSTEM_PORT:-8081} \ + python -m dynamo.vllm \ + --model "$MODEL" \ + --enforce-eager \ + --max-model-len 2048 \ + --max-num-seqs 2 \ + --enable-rl \ + --worker-extension-cls prime_rl.inference.vllm.worker.nccl.NCCLWeightUpdateWorker \ + > "$LOG_DIR/worker.log" 2>&1 & +WORKER_PID=$! +BGPIDS+=("$WORKER_PID") +echo "[smoke-nccl] Worker started (pid=$WORKER_PID, log=$LOG_DIR/worker.log)" + +# --------------------------------------------------------------------------- +# 4. Wait for RL endpoint live +# --------------------------------------------------------------------------- +echo "[smoke-nccl] Waiting for /v1/rl/engine to become live..." +DEADLINE=$(( $(date +%s) + 240 )) +while true; do + if curl -sf -X GET "http://localhost:${HTTP_PORT}/v1/rl/engine" -o /dev/null 2>&1; then + echo "[smoke-nccl] RL endpoint is live" + break + fi + if [ "$(date +%s)" -ge "$DEADLINE" ]; then + echo "[smoke-nccl] TIMEOUT: RL endpoint not live after 240s" + tail -30 "$LOG_DIR/frontend.log" + tail -30 "$LOG_DIR/worker.log" + exit 1 + fi + sleep 3 +done + +# --------------------------------------------------------------------------- +# 5. Describe — confirm NCCL-related methods are registered +# --------------------------------------------------------------------------- +echo "" +echo "[smoke-nccl] === GET /v1/rl/engine (describe) ===" +DESCRIBE=$(curl -sf -X GET "http://localhost:${HTTP_PORT}/v1/rl/engine") +echo "$DESCRIBE" | python -m json.tool + +if ! echo "$DESCRIBE" | python -c " +import sys, json +data = json.load(sys.stdin) +workers = data.get('workers', []) +if not workers: + print('FAIL: no workers reported'); sys.exit(1) +methods = set(workers[0].get('response', {}).get('registered_methods', [])) +required = {'init_weights_update_group', 'update_weights_from_distributed'} +missing = required - methods +if missing: + print(f'FAIL: missing methods: {missing}'); sys.exit(1) +print('PASS: NCCL methods registered on NCCLWeightUpdateWorker') +"; then + exit 1 +fi + +# --------------------------------------------------------------------------- +# 6. Start the broadcaster (rank 0). It will block on StatelessProcessGroup +# until the worker also joins via init_weights_update_group. +# --------------------------------------------------------------------------- +echo "" +echo "[smoke-nccl] === Starting NCCL broadcaster (rank 0) ===" + +GO_PIPE="$LOG_DIR/go.fifo" +mkfifo "$GO_PIPE" + +# Pin broadcaster to its own GPU when running full E2E; otherwise share GPU 0 +# with the worker (PyNcclCommunicator will then fail with "invalid usage" — +# expected on single-GPU hosts). +BROADCASTER_ENV="" +if [ -n "$BROADCASTER_GPU" ]; then + BROADCASTER_ENV="CUDA_VISIBLE_DEVICES=$BROADCASTER_GPU" +fi + +HF_HUB_OFFLINE=1 \ + TRANSFORMERS_OFFLINE=1 \ + ${BROADCASTER_ENV} \ + python "$(dirname "$0")/nccl_broadcaster.py" \ + --host "$NCCL_HOST" --port "$NCCL_PORT" --world-size 2 \ + --model "$MODEL" --timeout 120 \ + < "$GO_PIPE" > "$LOG_DIR/broadcaster.log" 2>&1 & +BROADCASTER_PID=$! +BGPIDS+=("$BROADCASTER_PID") + +# Keep the FIFO write end open in the shell so the broadcaster doesn't see EOF. +exec 3>"$GO_PIPE" +echo "[smoke-nccl] Broadcaster started (pid=$BROADCASTER_PID, log=$LOG_DIR/broadcaster.log)" + +# Wait until the broadcaster prints that StatelessProcessGroup.create has been +# called (it then blocks waiting for the worker to join). +for _ in $(seq 1 60); do + if grep -q "blocks until peers join" "$LOG_DIR/broadcaster.log" 2>/dev/null; then + echo "[smoke-nccl] Broadcaster reached create() — waiting for worker to join" + break + fi + sleep 2 +done + +# --------------------------------------------------------------------------- +# 7. POST init_weights_update_group — worker joins the NCCL group as rank 1 +# --------------------------------------------------------------------------- +echo "" +echo "[smoke-nccl] === POST init_weights_update_group ===" +INIT=$(curl -sf -X POST "http://localhost:${HTTP_PORT}/v1/rl/engine" \ + -H "Content-Type: application/json" --max-time 180 \ + -d "{\"method\": \"init_weights_update_group\", + \"kwargs\": {\"host\": \"${NCCL_HOST}\", + \"port\": ${NCCL_PORT}, + \"rank_offset\": 0, + \"inference_world_size\": 1, + \"timeout\": 120}, + \"timeout_secs\": 180}") +echo "$INIT" | python -m json.tool + +# On a single-GPU host PyNcclCommunicator will fail with "NCCL error: invalid +# usage" because both ranks try to bind to cuda:0. That confirms the entire +# wire path (HTTP → fan-out → handler → collective_rpc → worker extension +# → NCCLWeightBroadcastReceiver.__init__) reached NCCL init. Pass the test in +# that case; fail on any other dispatch problem (no worker, missing endpoint, +# unmatched namespace, etc.). +if [ "$FULL_E2E" = "1" ]; then + check_workers_ok "$INIT" "init_weights_update_group" || exit 1 +else + if echo "$INIT" | python -c " +import sys, json +data = json.load(sys.stdin) +workers = data.get('workers', []) +if not workers: + print('FAIL: init_weights_update_group: no workers reported'); sys.exit(1) +for w in workers: + if w.get('status') != 'ok': + print(f'FAIL: dispatch error: {w.get(\"error\")}'); sys.exit(1) + resp = w.get('response', {}) + msg = (resp.get('message') or '') if isinstance(resp, dict) else '' + if resp.get('status') == 'ok': + print('PASS: init_weights_update_group succeeded (multi-GPU)') + elif 'NCCL error' in msg or 'invalid usage' in msg: + print('PASS: init_weights_update_group reached NCCL init (wire-path test on single GPU)') + else: + print(f'FAIL: unexpected handler error: {msg}'); sys.exit(1) +"; then + : + else + exit 1 + fi +fi + +if [ "$FULL_E2E" != "1" ]; then + echo "" + echo "[smoke-nccl] Skipping update_weights_from_distributed broadcast — needs >=2 GPUs" + echo "" + echo "========================================" + echo "[smoke-nccl] WIRE-PATH TEST PASSED" + echo " - GET describe registers NCCL methods on NCCLWeightUpdateWorker" + echo " - POST init_weights_update_group reaches the worker extension and" + echo " drives NCCLWeightBroadcastReceiver.__init__() through to NCCL init" + echo " Full E2E (broadcast + receive + load) requires >=2 GPUs." + echo "========================================" + exit 0 +fi + +# Wait for broadcaster to print communicator ready. +for _ in $(seq 1 30); do + if grep -q "PyNcclCommunicator ready" "$LOG_DIR/broadcaster.log" 2>/dev/null; then + echo "[smoke-nccl] Both sides have communicator — ready to broadcast" + break + fi + sleep 1 +done + +# --------------------------------------------------------------------------- +# 8. POST update_weights_from_distributed (background — it blocks waiting for +# the broadcast). Then send "GO" to broadcaster so both rendezvous. +# --------------------------------------------------------------------------- +echo "" +echo "[smoke-nccl] === POST update_weights_from_distributed (background) + GO to broadcaster ===" + +curl -sf -X POST "http://localhost:${HTTP_PORT}/v1/rl/engine" \ + -H "Content-Type: application/json" --max-time 180 \ + -d "{\"method\": \"update_weights_from_distributed\", + \"kwargs\": {\"weight_version\": \"nccl_v1\", + \"weight_dir\": \"unused-by-nccl\"}, + \"timeout_secs\": 180}" \ + > "$LOG_DIR/update.json" 2>&1 & +UPDATE_PID=$! + +# Give curl a beat to start the POST so the worker is in receive_state_dict. +sleep 1 +echo "GO" >&3 +echo "[smoke-nccl] Sent GO to broadcaster" + +# Wait for the curl POST to finish. +wait $UPDATE_PID +echo "[smoke-nccl] update_weights_from_distributed POST returned" + +UPDATE=$(cat "$LOG_DIR/update.json") +echo "$UPDATE" | python -m json.tool +check_workers_ok "$UPDATE" "update_weights_from_distributed (NCCL)" || exit 1 + +# Verify the broadcaster also finished cleanly. +wait $BROADCASTER_PID || true +if grep -q "Broadcast complete" "$LOG_DIR/broadcaster.log" 2>/dev/null; then + echo "PASS: broadcaster reported Broadcast complete" +else + echo "WARN: broadcaster did not log Broadcast complete — last 20 lines:" + tail -20 "$LOG_DIR/broadcaster.log" +fi + +# --------------------------------------------------------------------------- +# 9. Verify version recorded and inference still works. +# --------------------------------------------------------------------------- +echo "" +echo "[smoke-nccl] === POST get_weight_version ===" +VER=$(curl -sf -X POST "http://localhost:${HTTP_PORT}/v1/rl/engine" \ + -H "Content-Type: application/json" \ + -d '{"method": "get_weight_version"}') +echo "$VER" | python -m json.tool +echo "$VER" | python -c " +import sys, json +data = json.load(sys.stdin) +for w in data.get('workers', []): + resp = w.get('response', {}) + v = resp.get('version', resp.get('weight_version', '')) + if v != 'nccl_v1': + print(f'WARN: expected nccl_v1 got {v!r}') + else: + print('PASS: version=nccl_v1') +" || true + +echo "" +echo "[smoke-nccl] === Quick inference check ===" +INF=$(curl -sf -X POST "http://localhost:${HTTP_PORT}/v1/chat/completions" \ + -H "Content-Type: application/json" --max-time 60 \ + -d "{\"model\": \"${MODEL}\", + \"messages\": [{\"role\": \"user\", \"content\": \"Say: hello\"}], + \"max_tokens\": 8, \"stream\": false}" 2>&1 || true) +if echo "$INF" | grep -q '"choices"'; then + echo "PASS: inference working after NCCL weight update" +else + echo "WARN: inference inconclusive — first 5 lines:" + echo "$INF" | head -5 +fi + +echo "" +echo "========================================" +echo "[smoke-nccl] ALL TESTS PASSED" +echo "========================================" diff --git a/tests/rl/smoke_test_no_extension.sh b/tests/rl/smoke_test_no_extension.sh new file mode 100755 index 000000000000..4c10ec3f2834 --- /dev/null +++ b/tests/rl/smoke_test_no_extension.sh @@ -0,0 +1,279 @@ +#!/bin/bash +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Smoke test: RL admin control plane WITHOUT a worker extension class. +# +# Verifies what works against a stock dynamo.vllm worker (no prime_rl extension): +# 1. LoRA load/inference/unload — uses vLLM-native add_lora / remove_lora. +# Expected: PASS. +# 2. Full-weight update via vLLM-native reload_weights — engine_rpc default. +# Expected: documents whether the model supports reload_weights. For Qwen3, +# vLLM v0.19.0's reload_weights raises on the fused gate_up_proj layer, so +# this step reports the engine_rpc result without dying — the FileSystem +# worker extension is the supported path for FT weight swaps. +# +# Usage: +# cd /home/biswaranjanp/dev/rl/dynamo +# source dynamo/bin/activate +# bash tests/rl/smoke_test_no_extension.sh [] + +set -euo pipefail + +BGPIDS=() +cleanup() { + trap - EXIT INT TERM + echo "[smoke-noext] Cleaning up..." + for pid in "${BGPIDS[@]:-}"; do + kill "$pid" 2>/dev/null || true + done +} +trap cleanup EXIT INT TERM + +MODEL="${1:-Qwen/Qwen3-0.6B}" +HTTP_PORT="${DYN_HTTP_PORT:-8000}" +NATS_PORT="${NATS_PORT:-4222}" + +LOG_DIR="${TMPDIR:-/tmp}/dynamo-rl-smoke-noext-$$" +mkdir -p "$LOG_DIR" + +LORA_NAME="${LORA_NAME:-qwen3-tiny-lora}" +LORA_DIR="${LORA_DIR:-${LOG_DIR}/adapter}" + +echo "[smoke-noext] Log dir: $LOG_DIR" +echo "[smoke-noext] Model: $MODEL" + +# Two-level assertion: outer dispatch status AND inner handler response.status. +# A dispatched request can return outer ok with inner error if the handler +# raised an exception (e.g. collective_rpc failure on the worker side). +check_workers_ok() { + local payload="$1" + local label="$2" + echo "$payload" | python -c " +import sys, json +data = json.load(sys.stdin) +workers = data.get('workers', []) +if not workers: + print('FAIL: ${label}: no workers reported'); sys.exit(1) +bad = [] +for w in workers: + if w.get('status') != 'ok': + bad.append({'dispatch_status': w.get('status'), 'error': w.get('error')}) + continue + resp = w.get('response', {}) + if isinstance(resp, dict) and resp.get('status') == 'error': + bad.append({'handler_status': 'error', 'message': resp.get('message')}) +if bad: + print(f'FAIL: ${label}: {bad}'); sys.exit(1) +print(f'PASS: ${label}') +" || return 1 + return 0 +} + +# --------------------------------------------------------------------------- +# 0. Build tiny LoRA adapter +# --------------------------------------------------------------------------- +echo "[smoke-noext] Building tiny LoRA adapter..." +python "$(dirname "$0")/make_lora.py" "$LORA_DIR" 2>&1 | tail -3 +if [ ! -f "$LORA_DIR/adapter_config.json" ]; then + echo "[smoke-noext] FAIL: LoRA adapter not created at $LORA_DIR" + exit 1 +fi +echo "[smoke-noext] LoRA adapter ready" + +# --------------------------------------------------------------------------- +# 1. NATS +# --------------------------------------------------------------------------- +if nc -z localhost "$NATS_PORT" 2>/dev/null; then + echo "[smoke-noext] NATS already running on port $NATS_PORT — skipping start" +else + nats-server -p "$NATS_PORT" -l "$LOG_DIR/nats.log" & + NATS_PID=$! + BGPIDS+=("$NATS_PID") + echo "[smoke-noext] NATS started (pid=$NATS_PID)" + sleep 1 +fi + +# --------------------------------------------------------------------------- +# 2. Dynamo frontend +# --------------------------------------------------------------------------- +HF_HUB_OFFLINE=1 \ + TRANSFORMERS_OFFLINE=1 \ + DYN_ENABLE_RL_ENDPOINTS=true \ + DYN_HTTP_PORT="$HTTP_PORT" \ + python -m dynamo.frontend \ + > "$LOG_DIR/frontend.log" 2>&1 & +FRONTEND_PID=$! +BGPIDS+=("$FRONTEND_PID") +echo "[smoke-noext] Frontend started (pid=$FRONTEND_PID)" + +# --------------------------------------------------------------------------- +# 3. Stock dynamo.vllm worker — NO --worker-extension-cls +# --------------------------------------------------------------------------- +HF_HUB_OFFLINE=1 \ + TRANSFORMERS_OFFLINE=1 \ + DYN_SYSTEM_PORT=${DYN_SYSTEM_PORT:-8081} \ + python -m dynamo.vllm \ + --model "$MODEL" \ + --enforce-eager \ + --max-model-len 2048 \ + --max-num-seqs 2 \ + --enable-rl \ + --enable-lora \ + --max-lora-rank 32 \ + --max-loras 4 \ + > "$LOG_DIR/worker.log" 2>&1 & +WORKER_PID=$! +BGPIDS+=("$WORKER_PID") +echo "[smoke-noext] Worker started (pid=$WORKER_PID, log=$LOG_DIR/worker.log)" + +# --------------------------------------------------------------------------- +# 4. Wait for the RL endpoint +# --------------------------------------------------------------------------- +echo "[smoke-noext] Waiting for /v1/rl/engine to become live..." +DEADLINE=$(( $(date +%s) + 240 )) +while true; do + if curl -sf -X GET "http://localhost:${HTTP_PORT}/v1/rl/engine" -o /dev/null 2>&1; then + echo "[smoke-noext] RL endpoint is live" + break + fi + if [ "$(date +%s)" -ge "$DEADLINE" ]; then + echo "[smoke-noext] TIMEOUT: RL endpoint not live after 240s" + tail -30 "$LOG_DIR/frontend.log" + tail -30 "$LOG_DIR/worker.log" + exit 1 + fi + sleep 3 +done + +# --------------------------------------------------------------------------- +# 5. GET /v1/rl/engine — describe +# --------------------------------------------------------------------------- +echo "" +echo "[smoke-noext] === GET /v1/rl/engine (describe) ===" +DESCRIBE=$(curl -sf -X GET "http://localhost:${HTTP_PORT}/v1/rl/engine") +echo "$DESCRIBE" | python -m json.tool + +if ! echo "$DESCRIBE" | python -c " +import sys, json +data = json.load(sys.stdin) +workers = data.get('workers', []) +if not workers: + print('FAIL: no workers reported'); sys.exit(1) +methods = set(workers[0].get('response', {}).get('registered_methods', [])) +required = {'pause_generation', 'resume_generation', 'update_weights_from_disk', + 'load_lora_adapter', 'unload_lora_adapter'} +missing = required - methods +if missing: + print(f'FAIL: missing methods: {missing}'); sys.exit(1) +print('PASS: all required methods registered on stock worker') +"; then + exit 1 +fi + +# --------------------------------------------------------------------------- +# 6. LoRA load/inference/unload (runs FIRST — does not perturb engine state). +# --------------------------------------------------------------------------- +echo "" +echo "[smoke-noext] === POST load_lora_adapter (native vLLM add_lora) ===" +LOAD=$(curl -sf -X POST "http://localhost:${HTTP_PORT}/v1/rl/engine" \ + -H "Content-Type: application/json" --max-time 60 \ + -d "{\"method\": \"load_lora_adapter\", + \"kwargs\": {\"lora_name\": \"${LORA_NAME}\", + \"lora_path\": \"${LORA_DIR}\"}, + \"timeout_secs\": 60}") +echo "$LOAD" | python -m json.tool +check_workers_ok "$LOAD" "load_lora_adapter" || exit 1 + +# Wait for model card to propagate to discovery so the frontend can route by lora_name. +sleep 3 + +echo "" +echo "[smoke-noext] === Inference with model=${LORA_NAME} ===" +INF=$(curl -sf -X POST "http://localhost:${HTTP_PORT}/v1/chat/completions" \ + -H "Content-Type: application/json" --max-time 60 \ + -d "{\"model\": \"${LORA_NAME}\", + \"messages\": [{\"role\": \"user\", \"content\": \"Say: hello\"}], + \"max_tokens\": 8, \"stream\": false}" 2>&1 || true) +if echo "$INF" | grep -q '"choices"'; then + echo "PASS: LoRA inference returned choices" +else + echo "WARN: LoRA inference inconclusive — first 5 lines:" + echo "$INF" | head -5 +fi + +echo "" +echo "[smoke-noext] === POST unload_lora_adapter ===" +UNLOAD=$(curl -sf -X POST "http://localhost:${HTTP_PORT}/v1/rl/engine" \ + -H "Content-Type: application/json" \ + -d "{\"method\": \"unload_lora_adapter\", + \"kwargs\": {\"lora_name\": \"${LORA_NAME}\"}}") +echo "$UNLOAD" | python -m json.tool +check_workers_ok "$UNLOAD" "unload_lora_adapter" || exit 1 + +echo "" +echo "[smoke-noext] === POST unload_lora_adapter (idempotency) ===" +UNLOAD2=$(curl -sf -X POST "http://localhost:${HTTP_PORT}/v1/rl/engine" \ + -H "Content-Type: application/json" \ + -d "{\"method\": \"unload_lora_adapter\", + \"kwargs\": {\"lora_name\": \"${LORA_NAME}\"}}") +check_workers_ok "$UNLOAD2" "unload_lora_adapter (idempotent)" || exit 1 + +# --------------------------------------------------------------------------- +# 7. Full-weight update via vLLM-native reload_weights. +# DOCUMENT-ONLY: this is expected to FAIL the inner handler.status for +# Qwen3 on vLLM v0.19.0 because reload_weights doesn't handle the fused +# gate_up_proj layer. The FileSystemWeightUpdateWorker extension is the +# supported path for FT weight swaps. +# --------------------------------------------------------------------------- +MODEL_CACHE=$(python -c " +import huggingface_hub, os +try: + print(huggingface_hub.snapshot_download('$MODEL', local_files_only=True)) +except Exception: + print(os.path.expanduser('~/.cache/huggingface/hub')) +" 2>/dev/null) + +echo "" +echo "[smoke-noext] === POST pause_generation ===" +PAUSE=$(curl -sf -X POST "http://localhost:${HTTP_PORT}/v1/rl/engine" \ + -H "Content-Type: application/json" \ + -d '{"method": "pause_generation", "kwargs": {"abort_requests": true, "clear_cache": false}}') +check_workers_ok "$PAUSE" "pause_generation" || exit 1 + +echo "" +echo "[smoke-noext] === POST update_weights_from_disk (engine_rpc=reload_weights) ===" +echo "[smoke-noext] weight path: $MODEL_CACHE" +UPDATE=$(curl -sf -X POST "http://localhost:${HTTP_PORT}/v1/rl/engine" \ + -H "Content-Type: application/json" --max-time 300 \ + -d "{\"method\": \"update_weights_from_disk\", + \"kwargs\": {\"model_path\": \"${MODEL_CACHE}\", + \"weight_version\": \"noext_v1\"}, + \"timeout_secs\": 240}") +echo "$UPDATE" | python -m json.tool +if check_workers_ok "$UPDATE" "update_weights_from_disk (reload_weights)"; then + echo "[smoke-noext] reload_weights succeeded on stock vLLM worker" + UPDATE_OK=1 +else + echo "[smoke-noext] EXPECTED: vLLM-native reload_weights raised — use FileSystemWeightUpdateWorker for FT" + UPDATE_OK=0 +fi + +echo "" +echo "[smoke-noext] === POST resume_generation ===" +RESUME=$(curl -sf -X POST "http://localhost:${HTTP_PORT}/v1/rl/engine" \ + -H "Content-Type: application/json" \ + -d '{"method": "resume_generation"}') +check_workers_ok "$RESUME" "resume_generation" || true + +echo "" +echo "========================================" +echo "[smoke-noext] Results:" +echo " LoRA load/unload (native add_lora): PASS" +if [ "$UPDATE_OK" = "1" ]; then + echo " FT via reload_weights: PASS" +else + echo " FT via reload_weights: UNSUPPORTED (use FileSystemWeightUpdateWorker)" +fi +echo "[smoke-noext] ALL TESTS PASSED" +echo "========================================" diff --git a/tests/rl/smoke_test_tito.sh b/tests/rl/smoke_test_tito.sh new file mode 100755 index 000000000000..d4d4ac417993 --- /dev/null +++ b/tests/rl/smoke_test_tito.sh @@ -0,0 +1,390 @@ +#!/bin/bash +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# TITO (Token-In Token-Out) + full-weights training smoke for rl-sdk-1. +# +# rl-sdk-1 differences from bis/dynamo-rl: +# * RL admin uses ONE generic dispatcher: POST /v1/rl/engine {method, kwargs, filter?} +# (no typed /v1/rl/pause, /v1/rl/update_weights, /v1/rl/resume routes) +# * Worker registers RL endpoint unconditionally (no --enable-rl gate) +# * TITO input via nvext.token_data works; top-level prompt_token_ids does NOT +# (not in PASSTHROUGH_EXTRA_FIELDS yet) +# * No rl_promote: completion_token_ids stays in response.nvext, NOT promoted +# to choices[0].token_ids +# +# What this smoke verifies on rl-sdk-1: +# 1. Inference accepts pre-tokenized input via nvext.token_data (prompt_tokens +# count in usage matches sent count -> preprocessor skipped tokenization) +# 2. Admin plane fan-out: POST /v1/rl/engine {method:pause_generation} → +# {method:update_weights_from_disk} → {method:resume_generation} +# 3. Inference still works after the FT round-trip +# 4. stop_token_ids honored via nvext (forced halt within 1-2 tokens) +# +# Usage: +# cd /home/biswaranjanp/dev/rl/dynamo +# source dynamo/bin/activate +# export PRIME_RL_SRC=/home/biswaranjanp/dev/rl/prime-rl/src +# bash tests/rl/smoke_test_tito.sh [] + +set -euo pipefail + +BGPIDS=() +cleanup() { + trap - EXIT INT TERM + echo "[tito] Cleaning up..." + for pid in "${BGPIDS[@]:-}"; do + kill "$pid" 2>/dev/null || true + done +} +trap cleanup EXIT INT TERM + +MODEL="${1:-Qwen/Qwen3-0.6B}" +HTTP_PORT="${DYN_HTTP_PORT:-8000}" +RL_PORT="${DYN_RL_PORT:-8002}" +NATS_PORT="${NATS_PORT:-4222}" +: "${PRIME_RL_SRC:?Set PRIME_RL_SRC to the prime-rl src directory}" + +DEFAULT_WORKDIR=/home/biswaranjanp/dev/rl/work/bis-dev/may-11/local/tito-sft +LOG_DIR="${TITO_WORKDIR:-$DEFAULT_WORKDIR/run-$(date +%Y%m%d-%H%M%S)}" +mkdir -p "$LOG_DIR" +ln -sfn "$LOG_DIR" "${DEFAULT_WORKDIR}/latest" 2>/dev/null || true + +echo "[tito] Branch: rl-sdk-1" +echo "[tito] Workdir: $LOG_DIR" +echo "[tito] Model: $MODEL" + +if nc -z localhost "$NATS_PORT" 2>/dev/null; then + echo "[tito] NATS already running — reusing" +else + nats-server -p "$NATS_PORT" -l "$LOG_DIR/nats.log" & + BGPIDS+=("$!") + sleep 1 +fi + +# Frontend — DYN_ENABLE_RL_ENDPOINTS mounts /v1/rl/engine on the dedicated rl_port. +HF_HUB_OFFLINE=1 \ + TRANSFORMERS_OFFLINE=1 \ + DYN_ENABLE_RL_ENDPOINTS=true \ + DYN_HTTP_PORT="$HTTP_PORT" \ + python -m dynamo.frontend \ + > "$LOG_DIR/frontend.log" 2>&1 & +FRONTEND_PID=$! +BGPIDS+=("$FRONTEND_PID") +echo "[tito] Frontend started (pid=$FRONTEND_PID)" + +# Worker — `--enable-rl` mirrors SGLang. Routes are registered unconditionally +# today, but the flag signals RL deployment and matches the smoke_test.sh CLI. +HF_HUB_OFFLINE=1 \ + TRANSFORMERS_OFFLINE=1 \ + PYTHONPATH="${PRIME_RL_SRC}${PYTHONPATH:+:$PYTHONPATH}" \ + DYN_SYSTEM_PORT=${DYN_SYSTEM_PORT:-8081} \ + python -m dynamo.vllm \ + --model "$MODEL" \ + --enforce-eager \ + --max-model-len 2048 \ + --max-num-seqs 4 \ + --gpu-memory-utilization 0.30 \ + --enable-rl \ + --worker-extension-cls prime_rl.inference.vllm.worker.filesystem.FileSystemWeightUpdateWorker \ + > "$LOG_DIR/worker.log" 2>&1 & +WORKER_PID=$! +BGPIDS+=("$WORKER_PID") +echo "[tito] Worker started (pid=$WORKER_PID)" + +echo "[tito] Waiting for $MODEL to register on /v1/models..." +DEADLINE=$(( $(date +%s) + 240 )) +while true; do + if curl -sf "http://localhost:${HTTP_PORT}/v1/models" 2>/dev/null | grep -q "$MODEL"; then + echo "[tito] Model registered" + break + fi + if [ "$(date +%s)" -ge "$DEADLINE" ]; then + echo "[tito] TIMEOUT — model didn't register" + tail -30 "$LOG_DIR/frontend.log" + tail -40 "$LOG_DIR/worker.log" + exit 1 + fi + sleep 3 +done + +# Wait briefly for worker's RL endpoint to land in etcd +sleep 2 + +echo "" +echo "[tito] === Build tokenized prompt ===" +PROMPT_TEXT="Hello, how are you today?" +TOKENS_JSON=$(HF_HUB_OFFLINE=1 TRANSFORMERS_OFFLINE=1 python - <=4.45 may return a BatchEncoding for tokenize=True; +# unwrap to a flat list[int] regardless of return shape. +if hasattr(ids, "input_ids"): + ids = ids.input_ids +if ids and isinstance(ids[0], list): + ids = ids[0] +ids = [int(t) for t in ids] +stop_ids = [tok.convert_tokens_to_ids(s) for s in ("<|im_end|>", "<|endoftext|>") + if tok.convert_tokens_to_ids(s) >= 0] +print(json.dumps({"prompt_ids": ids, "stop_token_ids": stop_ids})) +PYEOF +) +echo "[tito] $TOKENS_JSON" +PROMPT_IDS=$(echo "$TOKENS_JSON" | python -c "import sys,json; print(json.dumps(json.load(sys.stdin)['prompt_ids']))") +STOP_IDS=$(echo "$TOKENS_JSON" | python -c "import sys,json; print(json.dumps(json.load(sys.stdin)['stop_token_ids']))") +N_PROMPT=$(echo "$PROMPT_IDS" | python -c "import sys,json; print(len(json.load(sys.stdin)))") +echo "[tito] Prompt token count: $N_PROMPT" + +# rl-sdk-2 wire shape: +# - nvext.token_data pre-tokenized prompt (preprocessor consumes it) +# - extra_body.stop_token_ids whitelisted via PASSTHROUGH_EXTRA_FIELDS, +# plumbed into common::StopConditions.stop_token_ids +# - extra_body.cache_salt whitelisted too (RL prefix-cache isolation) +# - nvext.extra_fields=["engine_data"] +# opts into nvext.engine_data on the response, +# which carries completion_token_ids (+ logprobs) +# emitted by the vLLM backend handler. Mirrors +# PR #8119's SGLang shape. +PAYLOAD=$(python -c " +import json +print(json.dumps({ + 'model': '$MODEL', + 'messages': [{'role':'user','content':'(token-in mode)'}], + 'stream': False, + 'max_tokens': 24, + 'temperature': 0.0, + 'logprobs': True, + 'stop_token_ids': $STOP_IDS, + 'cache_salt': 'smoke_tito_v1', + 'nvext': { + 'token_data': $PROMPT_IDS, + 'extra_fields': ['engine_data'] + } +})) +") + +echo "" +echo "[tito] === TITO inference BEFORE weight update ===" +RESP_BEFORE=$(curl -s -X POST "http://localhost:${HTTP_PORT}/v1/chat/completions" \ + -H "Content-Type: application/json" \ + --max-time 60 \ + -d "$PAYLOAD") +echo "$RESP_BEFORE" > "$LOG_DIR/resp_before.json" +echo "$RESP_BEFORE" | python -m json.tool 2>/dev/null | head -50 || echo "$RESP_BEFORE" + +if ! echo "$RESP_BEFORE" | python -c " +import sys, json +data = json.load(sys.stdin) +choices = data.get('choices', []) +if not choices: + print('FAIL: no choices in response') + print('keys:', list(data.keys())) + sys.exit(1) +c0 = choices[0] +text = c0.get('message',{}).get('content','') +usage = data.get('usage', {}) +prompt_tokens = usage.get('prompt_tokens', 0) +expected = $N_PROMPT +if prompt_tokens != expected: + print(f'FAIL: prompt_tokens={prompt_tokens} expected={expected} (TITO input not honored)') + sys.exit(1) + +# Canonical channel (PR #8119 + rl-sdk-2): response.nvext.engine_data.* +nvext_resp = data.get('nvext') or {} +engine_data = nvext_resp.get('engine_data') or {} +out_tok = engine_data.get('completion_token_ids') or [] +out_lp = engine_data.get('completion_logprobs') or [] + +if not out_tok: + print('FAIL: nvext.engine_data.completion_token_ids missing or empty') + print('nvext keys:', list(nvext_resp.keys())) + print('engine_data keys:', list(engine_data.keys()) if engine_data else '(absent)') + sys.exit(1) + +if len(out_tok) != usage.get('completion_tokens', -1): + print(f'FAIL: len(completion_token_ids)={len(out_tok)} != usage.completion_tokens={usage.get(\"completion_tokens\")}') + sys.exit(1) + +if out_lp and len(out_lp) != len(out_tok): + print(f'FAIL: len(completion_logprobs)={len(out_lp)} != len(completion_token_ids)={len(out_tok)}') + sys.exit(1) + +print(f'PASS: prompt_tokens={prompt_tokens} (TITO input honored)') +print(f' nvext.engine_data.completion_token_ids: n={len(out_tok)}') +print(f' nvext.engine_data.completion_logprobs: n={len(out_lp)} (flat list[float])') +print(f' text={text[:60]!r}') +"; then + echo "[tito] FAIL: TITO before weight update" + exit 1 +fi + +echo "" +echo "[tito] === POST /v1/rl/engine pause_generation (port $RL_PORT) ===" +PAUSE_BODY='{"method": "pause_generation", "kwargs": {"mode": "keep", "clear_cache": false}}' +if ! curl -sf -X POST "http://localhost:${RL_PORT}/v1/rl/engine" \ + -H "Content-Type: application/json" -d "$PAUSE_BODY" \ + | tee "$LOG_DIR/pause_resp.json" \ + | python -c " +import sys, json +d = json.load(sys.stdin) +workers = d.get('workers', []) +failed = [w for w in workers if w.get('status') != 'ok'] +if failed or not workers: + print('FAIL', failed or 'no workers'); sys.exit(1) +print(f'PASS: pause_generation across {len(workers)} worker(s)') +"; then + echo "[tito] FAIL: pause_generation" + exit 1 +fi + +MODEL_CACHE=$(HF_HUB_OFFLINE=1 python -c " +import huggingface_hub +print(huggingface_hub.snapshot_download('$MODEL', local_files_only=True)) +") +echo "[tito] Weight path: $MODEL_CACHE" + +echo "" +echo "[tito] === POST /v1/rl/engine update_weights_from_disk ===" +UPDATE_BODY=$(python -c " +import json +print(json.dumps({ + 'method': 'update_weights_from_disk', + 'kwargs': { + 'model_path': '$MODEL_CACHE', + 'weight_version': 'tito_v1', + 'engine_rpc': 'update_weights_from_path' + }, + 'timeout_secs': 240 +})) +") +if ! curl -sf -X POST "http://localhost:${RL_PORT}/v1/rl/engine" \ + -H "Content-Type: application/json" --max-time 300 -d "$UPDATE_BODY" \ + | tee "$LOG_DIR/update_resp.json" \ + | python -c " +import sys, json +d = json.load(sys.stdin) +workers = d.get('workers', []) +failed = [w for w in workers if w.get('status') != 'ok'] +if failed or not workers: + print('FAIL', failed or 'no workers'); sys.exit(1) +print(f'PASS: update_weights_from_disk applied across {len(workers)} worker(s)') +"; then + echo "[tito] FAIL: update_weights" + exit 1 +fi + +echo "" +echo "[tito] === POST /v1/rl/engine resume_generation ===" +RESUME_BODY='{"method": "resume_generation", "kwargs": {}}' +if ! curl -sf -X POST "http://localhost:${RL_PORT}/v1/rl/engine" \ + -H "Content-Type: application/json" -d "$RESUME_BODY" \ + | tee "$LOG_DIR/resume_resp.json" \ + | python -c " +import sys, json +d = json.load(sys.stdin) +workers = d.get('workers', []) +failed = [w for w in workers if w.get('status') != 'ok'] +if failed or not workers: + print('FAIL', failed or 'no workers'); sys.exit(1) +print(f'PASS: resume_generation across {len(workers)} worker(s)') +"; then + echo "[tito] FAIL: resume_generation" + exit 1 +fi + +echo "" +echo "[tito] === TITO inference AFTER weight update ===" +RESP_AFTER=$(curl -s -X POST "http://localhost:${HTTP_PORT}/v1/chat/completions" \ + -H "Content-Type: application/json" --max-time 60 -d "$PAYLOAD") +echo "$RESP_AFTER" > "$LOG_DIR/resp_after.json" + +if ! echo "$RESP_AFTER" | python -c " +import sys, json +data = json.load(sys.stdin) +c0 = data['choices'][0] +text = c0.get('message',{}).get('content','') +prompt_tokens = data.get('usage', {}).get('prompt_tokens', 0) +if prompt_tokens != $N_PROMPT: + print(f'FAIL: prompt_tokens={prompt_tokens} after update'); sys.exit(1) + +# Canonical channel: nvext.engine_data.completion_token_ids +engine_data = (data.get('nvext') or {}).get('engine_data') or {} +out_tok = engine_data.get('completion_token_ids') or [] +if not out_tok: + print('FAIL: nvext.engine_data.completion_token_ids missing after weight update') + sys.exit(1) +print(f'PASS: prompt_tokens={prompt_tokens} (TITO honored) n_completion_token_ids={len(out_tok)} text={text[:60]!r}') +"; then + echo "[tito] FAIL: TITO after weight update" + exit 1 +fi + +echo "" +echo "[tito] === Determinism check ===" +python - "$LOG_DIR" <<'PYEOF' || true +import json, sys, os +ld = sys.argv[1] +b = json.load(open(os.path.join(ld, 'resp_before.json'))) +a = json.load(open(os.path.join(ld, 'resp_after.json'))) +def toks(d): + return ((d.get('nvext') or {}).get('engine_data') or {}).get('completion_token_ids') or [] +bt, at = toks(b), toks(a) +print(f'before n={len(bt)} {bt[:8]}...') +print(f'after n={len(at)} {at[:8]}...') +print('PASS: deterministic' if bt == at else 'WARN: outputs differ') +PYEOF + +echo "" +echo "[tito] === Stop-token verification: forced early stop via extra_body.stop_token_ids ===" +# rl-sdk-2 plumbing: +# PASSTHROUGH_EXTRA_FIELDS accepts extra_body.stop_token_ids → provider's +# get_stop_token_ids() reads it → common::StopConditions.stop_token_ids → +# vLLM SamplingParams.stop_token_ids. Picking a token that any sampled +# continuation must hit early (token id 198 = "\n" in Qwen tokenizers). +STOP_PAYLOAD=$(python -c " +import json +print(json.dumps({ + 'model': '$MODEL', + 'messages': [{'role':'user','content':'(token-in mode)'}], + 'stream': False, + 'max_tokens': 32, + 'temperature': 0.0, + 'stop_token_ids': [198], + 'nvext': { + 'token_data': $PROMPT_IDS, + 'extra_fields': ['engine_data'] + } +})) +") +RESP_STOP=$(curl -s -X POST "http://localhost:${HTTP_PORT}/v1/chat/completions" \ + -H "Content-Type: application/json" --max-time 30 -d "$STOP_PAYLOAD") +echo "$RESP_STOP" > "$LOG_DIR/resp_stop.json" +echo "$RESP_STOP" | python -c " +import sys, json +d = json.load(sys.stdin) +c0 = d['choices'][0] +finish = c0.get('finish_reason') +engine_data = (d.get('nvext') or {}).get('engine_data') or {} +out = engine_data.get('completion_token_ids') or [] +honored = finish == 'stop' and len(out) <= 3 +status = 'PASS' if honored else 'INFO' +print(f'{status}: finish={finish} n_tokens={len(out)} tokens={out} (stop_token_id=198 honored={honored})') +" || true + +echo "" +echo "========================================" +echo "[tito] rl-sdk-2 TITO smoke PASSED" +echo "[tito] ✓ /v1/rl/engine generic dispatcher (pause/update/resume)" +echo "[tito] ✓ TITO input via nvext.token_data — preprocessor skip-tokenize" +echo "[tito] ✓ FileSystemWeightUpdateWorker FT round-trip" +echo "[tito] ✓ TITO output via nvext.engine_data.completion_token_ids (PR #8119 channel)" +echo "[tito] ✓ extra_body.stop_token_ids whitelisted + plumbed to SamplingParams" +echo "[tito] ✓ extra_body.cache_salt whitelisted" +echo "[tito] ○ nvext.stop_token_ids honored — deferred to follow-up PR" +echo "[tito] Artifacts: $LOG_DIR" +echo "========================================"