diff --git a/deps/renderers b/deps/renderers index d4707862ac..ab603543ec 160000 --- a/deps/renderers +++ b/deps/renderers @@ -1 +1 @@ -Subproject commit d4707862ac83aa3773c21f4096aec72bd17b91e4 +Subproject commit ab603543ecd45f14d684a512f5b1af1853aaa287 diff --git a/deps/verifiers b/deps/verifiers index d30a3f48e5..333052c0bd 160000 --- a/deps/verifiers +++ b/deps/verifiers @@ -1 +1 @@ -Subproject commit d30a3f48e5f14b06b3081b2102ec32cc3149b849 +Subproject commit 333052c0bd6e443b1485481ae9c5d48a3f5a8352 diff --git a/packages/prime-rl-configs/src/prime_rl/configs/env_server.py b/packages/prime-rl-configs/src/prime_rl/configs/env_server.py index 8e5ed3f3a0..935227652b 100644 --- a/packages/prime-rl-configs/src/prime_rl/configs/env_server.py +++ b/packages/prime-rl-configs/src/prime_rl/configs/env_server.py @@ -3,7 +3,7 @@ import verifiers.v1 as vf from pydantic import SerializeAsAny, model_validator -from prime_rl.configs.shared import LogConfig +from prime_rl.configs.shared import LogConfig, MultimodalConfig from prime_rl.utils.config import BaseConfig @@ -26,6 +26,9 @@ class EnvServerConfig(BaseConfig): output_dir: Path = Path("outputs") """Directory to write outputs to — logs and any generated artifacts are written as subdirectories.""" + multimodal: MultimodalConfig = MultimodalConfig() + """Raw multimodal image offload settings; with ``output_dir``, resolves the run image-asset dir this server offloads into.""" + @model_validator(mode="before") @classmethod def _resolve_env(cls, data): diff --git a/packages/prime-rl-configs/src/prime_rl/configs/inference.py b/packages/prime-rl-configs/src/prime_rl/configs/inference.py index baab060f95..e9924c711d 100644 --- a/packages/prime-rl-configs/src/prime_rl/configs/inference.py +++ b/packages/prime-rl-configs/src/prime_rl/configs/inference.py @@ -5,7 +5,7 @@ from pydantic import Field, model_validator from pydantic_config import BaseConfig -from prime_rl.configs.shared import BaseModelConfig, EnvVars, LogConfig, SlurmConfig +from prime_rl.configs.shared import BaseModelConfig, EnvVars, LogConfig, MultimodalConfig, SlurmConfig from prime_rl.utils.config import find_package_resource, rgetattr, rsetattr from prime_rl.utils.parsers import resolve_reasoning_parser, resolve_tool_call_parser @@ -423,6 +423,9 @@ class InferenceConfig(BaseConfig): dry_run: bool = False """Only validate and dump resolved configs, then exit early.""" + multimodal: MultimodalConfig = MultimodalConfig() + """Raw multimodal image offload settings shared with trainer and orchestrator.""" + @model_validator(mode="after") def validate_multi_node_requires_slurm(self): if self.deployment.type in ("multi_node", "disaggregated") and self.slurm is None: diff --git a/packages/prime-rl-configs/src/prime_rl/configs/orchestrator.py b/packages/prime-rl-configs/src/prime_rl/configs/orchestrator.py index 8f5f0e859f..bcb9acaefb 100644 --- a/packages/prime-rl-configs/src/prime_rl/configs/orchestrator.py +++ b/packages/prime-rl-configs/src/prime_rl/configs/orchestrator.py @@ -17,6 +17,7 @@ FileMonitorConfig, HeartbeatConfig, LogConfig, + MultimodalConfig, PrimeMonitorConfig, TransportConfig, WandbWithExtrasConfig, @@ -571,6 +572,9 @@ class OrchestratorConfig(BaseConfig): heartbeat: HeartbeatConfig | None = None """BetterStack heartbeat configuration for monitoring training progress.""" + multimodal: MultimodalConfig = MultimodalConfig() + """Raw multimodal image offload settings shared with trainer and inference.""" + @model_validator(mode="after") def auto_setup_tokenizer(self): if self.tokenizer.name is None: diff --git a/packages/prime-rl-configs/src/prime_rl/configs/rl.py b/packages/prime-rl-configs/src/prime_rl/configs/rl.py index 8a4a02ae21..a889be4581 100644 --- a/packages/prime-rl-configs/src/prime_rl/configs/rl.py +++ b/packages/prime-rl-configs/src/prime_rl/configs/rl.py @@ -21,6 +21,7 @@ from prime_rl.configs.shared import ( EnvVars, FileMonitorConfig, + MultimodalConfig, SlurmConfig, TransportConfig, VLMConfig, @@ -261,6 +262,8 @@ class RLConfig(BaseConfig): weight_broadcast: SharedWeightBroadcastConfig | None = None + multimodal: MultimodalConfig = MultimodalConfig() + """Shared raw multimodal image offload settings. Propagated to trainer, orchestrator, and inference.""" rollout_transport: TransportConfig | None = None bench: bool = False diff --git a/packages/prime-rl-configs/src/prime_rl/configs/sft.py b/packages/prime-rl-configs/src/prime_rl/configs/sft.py index 941c531deb..e9506efb7c 100644 --- a/packages/prime-rl-configs/src/prime_rl/configs/sft.py +++ b/packages/prime-rl-configs/src/prime_rl/configs/sft.py @@ -248,6 +248,16 @@ def normalize_deployment(cls, data): ### Validate configs (e.g. raise for unsupported (combinations of) configs) + @model_validator(mode="after") + def renderer_emits_processed_multimodal(self): + """SFT needs processed MM tensors; default ``multimodal_output`` to that and reject ``raw``.""" + if "multimodal_output" in self.renderer.model_fields_set: + if self.renderer.multimodal_output == "raw": + raise ValueError("multimodal_output='raw' is unsupported for SFT (defaults to 'processed')") + else: + self.renderer = self.renderer.model_copy(update={"multimodal_output": "processed"}) + return self + @model_validator(mode="after") def deepep_disables_grad_clipping(self): if self.model.ep_comm_backend == "deepep" and self.optim.max_norm is not None: diff --git a/packages/prime-rl-configs/src/prime_rl/configs/shared.py b/packages/prime-rl-configs/src/prime_rl/configs/shared.py index 74f2a859b8..a8ce6f22f1 100644 --- a/packages/prime-rl-configs/src/prime_rl/configs/shared.py +++ b/packages/prime-rl-configs/src/prime_rl/configs/shared.py @@ -10,7 +10,13 @@ # and the single shared W&B run. The launcher always sets these last, so allowing them in # `env_vars` would be a silent no-op (or, on multi-node, a footgun) — reject them instead. PROTECTED_ENV_VARS = frozenset( - {"CUDA_VISIBLE_DEVICES", "WANDB_SHARED_MODE", "WANDB_SHARED_RUN_ID", "WANDB_SHARED_LABEL"} + { + "CUDA_VISIBLE_DEVICES", + "VF_RENDERER_IMAGE_OFFLOAD_DIR", + "WANDB_SHARED_MODE", + "WANDB_SHARED_RUN_ID", + "WANDB_SHARED_LABEL", + } ) @@ -86,6 +92,11 @@ def resolve_project_dir(self): ServerType = Literal["vllm", "openai"] +class MultimodalConfig(BaseConfig): + offload_dir: Path | None = None + """Directory for offloaded image assets. Supports environment expansion such as ``/data/outputs/run_${RUN_ID}/assets/images``. When unset, prime-rl resolves a run-scoped default.""" + + class VLMConfig(BaseConfig): vision_encoder_attr: str """Dotted attribute path to the vision encoder module (e.g. ``model.visual``).""" diff --git a/packages/prime-rl-configs/src/prime_rl/configs/trainer.py b/packages/prime-rl-configs/src/prime_rl/configs/trainer.py index f3fa658a44..eb6edf9c70 100644 --- a/packages/prime-rl-configs/src/prime_rl/configs/trainer.py +++ b/packages/prime-rl-configs/src/prime_rl/configs/trainer.py @@ -10,6 +10,7 @@ FileMonitorConfig, HeartbeatConfig, MetricsServerConfig, + MultimodalConfig, TrainerLogConfig, TransportConfig, WandbConfig, @@ -628,6 +629,9 @@ class TrainerConfig(BaseConfig): max_concurrent_runs: int = Field(1, ge=1) """Maximum number of concurrent runs to allow. If 1, only one run may run at a time.""" + multimodal: MultimodalConfig = MultimodalConfig() + """Raw multimodal image offload settings shared with orchestrator and inference.""" + enable_token_export: bool = False """Opt-in per-token JSONL export for rollout debugging. When enabled, writes token ids and aligned trainer metrics after each forward pass.""" diff --git a/packages/prime-rl-configs/src/prime_rl/utils/validation.py b/packages/prime-rl-configs/src/prime_rl/utils/validation.py index 6dac87c12e..c1b04cf9af 100644 --- a/packages/prime-rl-configs/src/prime_rl/utils/validation.py +++ b/packages/prime-rl-configs/src/prime_rl/utils/validation.py @@ -132,6 +132,7 @@ def propagate(shared_path: str, *targets: str) -> None: # Top-level scalars. propagate("max_steps", "trainer.max_steps", "orchestrator.max_steps") propagate("seq_len", "trainer.model.seq_len", "orchestrator.seq_len") + propagate("multimodal", "trainer.multimodal", "orchestrator.multimodal", "inference.multimodal") # [slurm] → inference: a multi-node RL run drives its inference deployment under # the same SLURM allocation, so the nested inference inherits [slurm]. This is diff --git a/pyproject.toml b/pyproject.toml index 2274aa3aa0..7a468e3a45 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -133,9 +133,12 @@ members = [ ] # tau2-synth-v1 and tau3-bench-v1 pin `tau2` forks/revs that conflict with the # canonical tau2-bench-v1; a workspace shares one lock, so these can't coexist. +# nemo-gym-weather-v1 needs verifiers[nemo-gym], whose nemo-gym==0.4.0 pins +# openai<=2.7.2 against verifiers' own openai>=2.9.0. exclude = [ "deps/research-environments/environments/tool_use/tau2_synth_v1", "deps/research-environments/environments/tool_use/tau3_bench_v1", + "deps/verifiers/environments/nemo_gym_weather_v1", ] # ModelExpress 0.3.0 publishes protobuf<6 metadata, but its generated proto is diff --git a/src/prime_rl/entrypoints/env_server.py b/src/prime_rl/entrypoints/env_server.py index 05951ba031..af245fc087 100644 --- a/src/prime_rl/entrypoints/env_server.py +++ b/src/prime_rl/entrypoints/env_server.py @@ -1,3 +1,4 @@ +import os from functools import partial from verifiers.v1 import pool_serve_kwargs @@ -7,11 +8,17 @@ from prime_rl.orchestrator.utils import setup_env_server_logging from prime_rl.utils.config import cli from prime_rl.utils.process import set_proc_title +from prime_rl.utils.run_assets import IMAGE_OFFLOAD_DIR_ENV, resolve_image_offload_dir from prime_rl.utils.utils import clean_exit @clean_exit def run_server(config: EnvServerConfig): + # Renderers offload images to the dir named by this env var; resolve it from + # this server's own config, letting an operator-set env win (multi-node). + os.environ.setdefault( + IMAGE_OFFLOAD_DIR_ENV, str(resolve_image_offload_dir(config.output_dir, config.multimodal, os.environ)) + ) # ``serve.pool`` (static or elastic) sizes the server; a v0/legacy env runs through # the bridge, a v1 env is a native env block — both speak the same serve protocol, # so the orchestrator is agnostic. serve_env applies the logging setup in this process diff --git a/src/prime_rl/entrypoints/rl.py b/src/prime_rl/entrypoints/rl.py index 279f91e00a..2aef461cd0 100644 --- a/src/prime_rl/entrypoints/rl.py +++ b/src/prime_rl/entrypoints/rl.py @@ -105,6 +105,8 @@ def write_subconfigs(config: RLConfig, output_dir: Path) -> None: "serve": {**source_dict.get("serve", {}), "address": address}, "legacy": source_dict.get("legacy", {}), "log": {"level": config.orchestrator.log.vf_level, "json_logging": config.orchestrator.log.json_logging}, + "output_dir": str(config.orchestrator.output_dir), + "multimodal": to_toml_dict(config.multimodal), } with open(env_dir / f"{source.resolved_name}.toml", "wb") as f: tomli_w.dump(env_server_dict, f) @@ -157,6 +159,7 @@ def rl_local(config: RLConfig): "WANDB_SHARED_MODE": "1", "WANDB_SHARED_RUN_ID": os.environ.get("WANDB_SHARED_RUN_ID", uuid.uuid4().hex), } + inherited_env = dict(os.environ) # Validate client port matches inference server port if config.inference is not None and not config.orchestrator.model.client.is_elastic: @@ -200,7 +203,7 @@ def sigterm_handler(signum, frame): inference_process = Popen( inference_cmd, env={ - **os.environ, + **inherited_env, **DEFAULT_COMMON_ENV_VARS, **DEFAULT_INFERENCE_ENV_VARS, **config.env_vars, @@ -290,7 +293,7 @@ def sigterm_handler(signum, frame): stdout=log_file, stderr=log_file, env={ - **os.environ, + **inherited_env, **DEFAULT_COMMON_ENV_VARS, "LOGURU_FORCE_COLORS": "1", "WANDB_PROGRAM": "uv run rl", @@ -339,7 +342,7 @@ def sigterm_handler(signum, frame): trainer_process = Popen( trainer_cmd, env={ - **os.environ, + **inherited_env, **DEFAULT_COMMON_ENV_VARS, **DEFAULT_TRAINER_ENV_VARS, "LOGURU_FORCE_COLORS": "1", @@ -437,6 +440,9 @@ def write_slurm_script(config: RLConfig, config_dir: Path, script_path: Path) -> kv_offload_disk_path=str(offload.disk.path) if (is_mooncake and offload.disk is not None) else "", kv_offload_device_name=offload.device_name if is_mooncake else "", ) + image_offload_dir = ( + os.path.expanduser(str(config.multimodal.offload_dir)) if config.multimodal.offload_dir is not None else "" + ) # Per-component env vars: launcher defaults (shared + multi-node-specific) with the # user's config merged on top. Runtime wiring stays in the template. @@ -463,6 +469,7 @@ def write_slurm_script(config: RLConfig, config_dir: Path, script_path: Path) -> **config.slurm.template_vars, config_path=config_dir / RL_TOML, output_dir=config.output_dir, + image_offload_dir=image_offload_dir, gpus_per_node=config.deployment.gpus_per_node, ) elif config.inference is not None and config.inference.deployment.type == "disaggregated": @@ -474,6 +481,7 @@ def write_slurm_script(config: RLConfig, config_dir: Path, script_path: Path) -> config_dir=config_dir, output_dir=config.output_dir, orchestrator_output_dir=config.orchestrator.output_dir, + image_offload_dir=image_offload_dir, num_train_nodes=config.deployment.num_train_nodes, num_infer_nodes=infer_deploy.num_nodes * config.deployment.num_infer_replicas, nodes_per_infer_replica=infer_deploy.num_nodes, @@ -515,6 +523,7 @@ def write_slurm_script(config: RLConfig, config_dir: Path, script_path: Path) -> config_dir=config_dir, # TODO: should prob have each subconfig path separately output_dir=config.output_dir, orchestrator_output_dir=config.orchestrator.output_dir, + image_offload_dir=image_offload_dir, num_train_nodes=config.deployment.num_train_nodes, num_infer_nodes=config.deployment.total_infer_nodes, nodes_per_infer_replica=config.deployment.infer_nodes_per_replica, diff --git a/src/prime_rl/inference/vllm/serving_tokens.py b/src/prime_rl/inference/vllm/serving_tokens.py index e14a5ac83e..ad75460eda 100644 --- a/src/prime_rl/inference/vllm/serving_tokens.py +++ b/src/prime_rl/inference/vllm/serving_tokens.py @@ -4,7 +4,8 @@ ``vllm.entrypoints.scale_out.token_in_token_out.serving.ServingTokens`` that covers prefix-cache salting, lora dispatch, multimodal features, prompt logprobs, priority, ``data_parallel_rank`` header routing and server-side ``max_tokens`` -defaulting. We subclass it for the bits still missing from the upstream handler: +defaulting. We subclass it for prime-RL behavior that is still missing or +customized: 1. ``data_parallel_rank`` routing — read from the ``X-data-parallel-rank`` header and forwarded to ``engine_client.generate``. Upstream ``ServingTokens`` @@ -15,7 +16,12 @@ decisions, surface them as base64 raw-byte payloads without requiring a vLLM source fork. -3. Server-side ``max_tokens`` defaulting — upstream ``ServingTokens`` now applies +3. Raw image refs for multimodal rollouts — renderers send a lightweight raw + descriptor ref at every image slot (current and prior turns alike). This + handler materializes every ref through multimodal adapters before vLLM sees + the request. + +4. Server-side ``max_tokens`` defaulting — upstream ``ServingTokens`` now applies this itself (via ``GenerateRequest.is_sampling_param_provided`` + ``get_max_tokens``); we keep an equivalent guard so callers that omit ``max_tokens`` don't truncate at vLLM's 16-token ``SamplingParams`` default. @@ -26,8 +32,16 @@ from __future__ import annotations -from collections.abc import AsyncGenerator, AsyncIterable -from functools import cached_property +import asyncio +import hashlib +import logging +import os +from collections import OrderedDict +from collections.abc import AsyncGenerator, AsyncIterable, Callable +from dataclasses import dataclass +from functools import cached_property, lru_cache +from http import HTTPStatus +from io import BytesIO from typing import Any from fastapi import Request @@ -44,10 +58,23 @@ ) from vllm.entrypoints.scale_out.token_in_token_out.serving import ServingTokens from vllm.entrypoints.serve.utils.api_utils import get_max_tokens +from vllm.multimodal.cache import MultiModalCache from vllm.outputs import RequestOutput from vllm.sampling_params import RequestOutputKind, SamplingParams from prime_rl.inference.vllm.routed_experts import RoutedExpertsCapture +from prime_rl.multimodal.registry import get_multimodal_adapter +from prime_rl.multimodal.schema import RawMMItem +from prime_rl.utils.mm import load_image_processor, read_verified_image_bytes + +logger = logging.getLogger(__name__) + + +@dataclass +class _MMImageRefError(Exception): + message: str + err_type: str = "invalid_mm_image_ref" + status_code: HTTPStatus = HTTPStatus.BAD_REQUEST class PrimeRlGenerateResponseChoice(GenerateResponseChoice): @@ -147,8 +174,245 @@ async def _client_set_max_tokens(raw_request: Request | None) -> bool: return isinstance(sp, dict) and "max_tokens" in sp +@lru_cache(maxsize=8) +def _load_image_processor(model_name: str, trust_remote_code: bool): + # Named def so tests can monkeypatch ``serving_tokens._load_image_processor``. + return load_image_processor(model_name, trust_remote_code) + + +def _materialize_raw_image_ref_sync( + raw_ref: str, + *, + feature_modality: str, + mm_hash: str, + expected_placeholder_length: int, + processor_model_name: str, + trust_remote_code: bool, +): + from PIL import Image + from renderers.mm_store import split_raw_mm_ref + + try: + ref = split_raw_mm_ref(raw_ref) + except ValueError as exc: + raise _MMImageRefError(str(exc)) from exc + if ref.modality != feature_modality: + raise _MMImageRefError(f"Expected feature modality {feature_modality!r}, got {ref.modality!r}") + if ref.mm_hash != mm_hash: + raise _MMImageRefError(f"Expected image hash {mm_hash}, got {ref.mm_hash}") + + try: + raw = read_verified_image_bytes(ref.raw_image_uri, ref.mm_hash) + except OSError as exc: + raise _MMImageRefError(f"Unable to read raw image asset {ref.raw_image_uri!r}: {exc}") from exc + except ValueError as exc: + raise _MMImageRefError(str(exc)) from exc + + try: + image = Image.open(BytesIO(raw)).convert("RGB") + except OSError as exc: + raise _MMImageRefError(f"Unable to decode raw image asset {ref.raw_image_uri!r}: {exc}") from exc + + image_processor = _load_image_processor(processor_model_name, trust_remote_code) + item = RawMMItem( + family=ref.family, + raw_image_uri=ref.raw_image_uri, + payload=dict(ref.payload), + ) + try: + adapter = get_multimodal_adapter(ref.family) + return adapter.materialize_for_vllm( + image_processor, + item, + image, + expected_placeholder_length, + ) + except (TypeError, ValueError) as exc: + raise _MMImageRefError(str(exc)) from exc + + +# (raw_ref_digest, feature_modality, mm_hash, expected_placeholder_length, +# processor_model_name, trust_remote_code). Digest the ref so forged outer +# hash/placeholder metadata cannot alias a cached entry. +_MaterializeKey = tuple[str, str, str, int, str, bool] + +_MM_MATERIALIZE_CACHE_GB_ENV = "PRIME_RL_MM_MATERIALIZE_CACHE_GB" +_MM_MATERIALIZE_LOG_EVERY = 1000 + + +class _MaterializedRefCache: + """Byte-bounded LRU of materialized raw image refs, with single-flight misses. + + All mutation happens on the event loop thread (materialization itself runs in a + worker thread via ``asyncio.to_thread``, but lookup/insertion/eviction happen in + the async caller) — so no lock is needed. Do not touch this cache from sync code. + """ + + def __init__(self, max_bytes: int) -> None: + self.max_bytes = max_bytes + self._items: OrderedDict[_MaterializeKey, tuple[Any, int]] = OrderedDict() + self._inflight: dict[_MaterializeKey, asyncio.Future] = {} + # asyncio only weakly refs tasks; keep strong refs until they finish. + self._tasks: set[asyncio.Task] = set() + self._bytes = 0 + self.hits = 0 + self.misses = 0 + self.evictions = 0 + + async def get_or_materialize(self, key: _MaterializeKey, materialize_fn: Callable[[], Any]) -> Any: + if self.max_bytes == 0: + return await asyncio.to_thread(materialize_fn) + if key in self._items: + self._items.move_to_end(key) + self.hits += 1 + self._maybe_log() + return self._items[key][0] + if (inflight := self._inflight.get(key)) is not None: + self.hits += 1 + self._maybe_log() + return await asyncio.shield(inflight) + self.misses += 1 + self._maybe_log() + future: asyncio.Future = asyncio.get_running_loop().create_future() + self._inflight[key] = future + + async def _materialize_and_resolve() -> None: + try: + item = await asyncio.to_thread(materialize_fn) + except BaseException as exc: + # Never cache a failure: the exception propagates to every awaiter + # and the next request for this key retries cleanly. + future.set_exception(exc) + # A never-awaited errored future logs a spurious warning at GC. + future.exception() + else: + future.set_result(item) + self._insert(key, item) + finally: + del self._inflight[key] + + # The work runs as its own task and awaiters shield the shared future: + # a cancelled request (client disconnect) must neither poison the future + # for deduped awaiters nor cancel it out from under them. + task = asyncio.get_running_loop().create_task(_materialize_and_resolve()) + self._tasks.add(task) + task.add_done_callback(self._tasks.discard) + return await asyncio.shield(future) + + def _insert(self, key: _MaterializeKey, item: Any) -> None: + nbytes = MultiModalCache.get_item_size(item) + if nbytes > self.max_bytes: + # Don't churn the whole cache for one oversized entry. + return + self._items[key] = (item, nbytes) + self._bytes += nbytes + while self._bytes > self.max_bytes: + _, (_, evicted_bytes) = self._items.popitem(last=False) + self._bytes -= evicted_bytes + self.evictions += 1 + + def _maybe_log(self) -> None: + total = self.hits + self.misses + if total % _MM_MATERIALIZE_LOG_EVERY == 0: + logger.info( + "mm materialize cache: hits=%d misses=%d hit_rate=%.1f%% bytes=%d/%d evictions=%d", + self.hits, + self.misses, + 100.0 * self.hits / total, + self._bytes, + self.max_bytes, + self.evictions, + ) + + +def _raw_ref_payloads_for_feature( + features: Any, feature_modality: str, hashes: list[str] +) -> tuple[list[str], list[Any]]: + from renderers.mm_store import is_raw_mm_ref + + kwargs_data = features.kwargs_data + if kwargs_data is None or feature_modality not in kwargs_data: + raise _MMImageRefError(f"v1 raw multimodal: modality {feature_modality!r} arrived with no raw refs") + + raw_refs = kwargs_data[feature_modality] + placeholders = features.mm_placeholders.get(feature_modality) + if placeholders is None: + raise _MMImageRefError(f"v1 raw multimodal: modality {feature_modality!r} arrived with no placeholders") + if len(raw_refs) != len(hashes): + raise _MMImageRefError( + f"Multimodal kwargs/hash length mismatch for {feature_modality}: {len(raw_refs)} != {len(hashes)}" + ) + if len(placeholders) != len(hashes): + raise _MMImageRefError( + f"Multimodal placeholder/hash length mismatch for {feature_modality}: {len(placeholders)} != {len(hashes)}" + ) + if not all(is_raw_mm_ref(item) for item in raw_refs): + raise _MMImageRefError("v1 multimodal inference accepts raw descriptor refs only") + return raw_refs, placeholders + + +async def _decode_raw_mm_kwargs( + features: Any, + *, + processor_model_name: str, + trust_remote_code: bool, + cache: _MaterializedRefCache, +) -> dict[str, list[Any]]: + # Flatten across modalities so every image in the request materializes + # concurrently; single-flight in the cache dedupes identical refs across + # (and within) requests. Fresh lists per request — vLLM's cache-injection + # path replaces list elements in place, so never hand it cache-owned lists. + flat: list[tuple[str, str, str, Any]] = [] + for feature_modality, hashes in features.mm_hashes.items(): + raw_refs, placeholders = _raw_ref_payloads_for_feature(features, feature_modality, hashes) + flat.extend(zip([feature_modality] * len(hashes), raw_refs, hashes, placeholders, strict=True)) + + def _materialize(feature_modality: str, raw_ref: str, mm_hash: str, placeholder: Any): + return _materialize_raw_image_ref_sync( + raw_ref, + feature_modality=feature_modality, + mm_hash=mm_hash, + expected_placeholder_length=placeholder.length, + processor_model_name=processor_model_name, + trust_remote_code=trust_remote_code, + ) + + decoded = await asyncio.gather( + *( + cache.get_or_materialize( + ( + hashlib.sha256(raw_ref.encode("utf-8")).hexdigest(), + feature_modality, + mm_hash, + placeholder.length, + processor_model_name, + trust_remote_code, + ), + lambda fm=feature_modality, r=raw_ref, h=mm_hash, p=placeholder: _materialize(fm, r, h, p), + ) + for feature_modality, raw_ref, mm_hash, placeholder in flat + ) + ) + + mm_kwargs: dict[str, list[Any]] = {feature_modality: [] for feature_modality in features.mm_hashes} + for (feature_modality, _, _, _), item in zip(flat, decoded, strict=True): + mm_kwargs[feature_modality].append(item) + return mm_kwargs + + class PrimeRlServingTokens(ServingTokens): - """ServingTokens + DP-rank routing + compact routed experts + max_tokens defaulting.""" + """ServingTokens + DP-rank routing + routed experts + raw image refs + max_tokens defaulting.""" + + @cached_property + def _mm_materialize_cache(self) -> _MaterializedRefCache: + """Sized by ``PRIME_RL_MM_MATERIALIZE_CACHE_GB`` (GiB, default 2.0, ``0`` disables). + + A ``cached_property`` because ``custom_init_app_state`` grafts this + subclass via ``object.__new__`` + ``__dict__.update``, so ``__init__`` + never runs (see ``_max_tokens_defaults``). + """ + gb = float(os.environ.get(_MM_MATERIALIZE_CACHE_GB_ENV, "2.0")) + return _MaterializedRefCache(max_bytes=int(gb * (1 << 30))) @cached_property def _max_tokens_defaults(self) -> tuple[dict, int | None]: @@ -198,27 +462,29 @@ async def serve_tokens( raw_request.state.request_metadata = request_metadata # Build the engine input — features-aware (MM) or text-only fallback. - # Identical to upstream so we keep tracking it. if features := request.features: - from vllm.entrypoints.scale_out.token_in_token_out.mm_serde import decode_mm_kwargs_item from vllm.inputs import mm_input - from vllm.multimodal.inputs import ( - MultiModalKwargsItem, - MultiModalKwargsItems, - PlaceholderRange, - ) + from vllm.multimodal.inputs import MultiModalKwargsItems, PlaceholderRange mm_placeholders = { modality: [PlaceholderRange(offset=p.offset, length=p.length) for p in ranges] for modality, ranges in features.mm_placeholders.items() } - mm_kwargs: dict[str, list[MultiModalKwargsItem | None]] = {} - if features.kwargs_data is not None: - for modality, items in features.kwargs_data.items(): - mm_kwargs[modality] = [decode_mm_kwargs_item(item) if item is not None else None for item in items] - else: - for modality, hashes in features.mm_hashes.items(): - mm_kwargs[modality] = [None] * len(hashes) + processor_model_name = getattr(self.model_config, "model", None) or model_name + trust_remote_code = bool(getattr(self.model_config, "trust_remote_code", False)) + try: + mm_kwargs = await _decode_raw_mm_kwargs( + features, + processor_model_name=processor_model_name, + trust_remote_code=trust_remote_code, + cache=self._mm_materialize_cache, + ) + except _MMImageRefError as exc: + return self.create_error_response( + message=exc.message, + err_type=exc.err_type, + status_code=exc.status_code, + ) engine_input = mm_input( prompt_token_ids=request.token_ids, mm_kwargs=MultiModalKwargsItems(mm_kwargs), diff --git a/src/prime_rl/multimodal/__init__.py b/src/prime_rl/multimodal/__init__.py new file mode 100644 index 0000000000..961b12ada0 --- /dev/null +++ b/src/prime_rl/multimodal/__init__.py @@ -0,0 +1,12 @@ +from prime_rl.multimodal.adapters.base import ForwardPolicy, MaterializedMM, MultimodalAdapter +from prime_rl.multimodal.registry import get_multimodal_adapter +from prime_rl.multimodal.schema import RawMMItem, parse_raw_mm_item + +__all__ = [ + "ForwardPolicy", + "MaterializedMM", + "MultimodalAdapter", + "RawMMItem", + "get_multimodal_adapter", + "parse_raw_mm_item", +] diff --git a/src/prime_rl/multimodal/adapters/__init__.py b/src/prime_rl/multimodal/adapters/__init__.py new file mode 100644 index 0000000000..453f5c39f7 --- /dev/null +++ b/src/prime_rl/multimodal/adapters/__init__.py @@ -0,0 +1,4 @@ +from prime_rl.multimodal.adapters.kimi_k25 import KimiK25Adapter +from prime_rl.multimodal.adapters.qwen_vl import QwenVLAdapter + +__all__ = ["KimiK25Adapter", "QwenVLAdapter"] diff --git a/src/prime_rl/multimodal/adapters/base.py b/src/prime_rl/multimodal/adapters/base.py new file mode 100644 index 0000000000..c6217e13dc --- /dev/null +++ b/src/prime_rl/multimodal/adapters/base.py @@ -0,0 +1,44 @@ +from __future__ import annotations + +from dataclasses import dataclass +from typing import TYPE_CHECKING, Any, Protocol + +if TYPE_CHECKING: + import torch + from PIL.Image import Image + + from prime_rl.multimodal.schema import RawMMItem + + +@dataclass(frozen=True) +class ForwardPolicy: + pass_position_ids_with_mm: bool = True + requires_mm_token_type_ids: bool = False + + +@dataclass(frozen=True) +class MaterializedMM: + kwargs: dict[str, "torch.Tensor"] + forward_policy: ForwardPolicy + + +class MultimodalAdapter(Protocol): + family: str + forward_policy: ForwardPolicy + + def validate_item(self, item: "RawMMItem") -> None: ... + + def materialize_for_trainer( + self, + image_processor: Any, + items: list["RawMMItem"], + images: list["Image"], + ) -> MaterializedMM: ... + + def materialize_for_vllm( + self, + image_processor: Any, + item: "RawMMItem", + image: "Image", + expected_placeholder_length: int, + ) -> Any: ... diff --git a/src/prime_rl/multimodal/adapters/kimi_k25.py b/src/prime_rl/multimodal/adapters/kimi_k25.py new file mode 100644 index 0000000000..cf3410b16f --- /dev/null +++ b/src/prime_rl/multimodal/adapters/kimi_k25.py @@ -0,0 +1,92 @@ +from __future__ import annotations + +from typing import Any + +from prime_rl.multimodal.adapters.base import ForwardPolicy, MaterializedMM +from prime_rl.multimodal.schema import RawMMItem + + +def _tensorize(value: Any): + import torch + + if isinstance(value, torch.Tensor): + return value.contiguous() + return torch.as_tensor(value).contiguous() + + +def _grid_payload(item: RawMMItem) -> list[int]: + grid = item.payload.get("grid_thws") + if grid is None: + raise ValueError("Kimi raw descriptor payload is missing grid_thws") + if not isinstance(grid, list | tuple): + raise ValueError(f"Invalid Kimi grid_thws: {grid!r}") + if len(grid) == 1 and isinstance(grid[0], list | tuple): + grid = grid[0] + if not isinstance(grid, list | tuple) or len(grid) != 3: + raise ValueError(f"Invalid Kimi grid_thws: {grid!r}") + out = [int(v) for v in grid] + if any(v <= 0 for v in out): + raise ValueError(f"Invalid Kimi grid_thws: {grid!r}") + return out + + +def _process_images(image_processor: Any, images: list[Any], *, return_tensors: str): + medias = [{"type": "image", "image": image} for image in images] + preprocess = getattr(image_processor, "preprocess", None) + if preprocess is None: + raise ValueError("Kimi image processor is missing preprocess") + return preprocess(medias, return_tensors=return_tensors) + + +class KimiK25Adapter: + family = "kimi_k25" + forward_policy = ForwardPolicy(pass_position_ids_with_mm=True) + + def validate_item(self, item: RawMMItem) -> None: + if item.family != self.family: + raise ValueError(f"Kimi adapter cannot handle family {item.family!r}") + _grid_payload(item) + + def materialize_for_trainer( + self, + image_processor: Any, + items: list[RawMMItem], + images: list[Any], + ) -> MaterializedMM: + for item in items: + self.validate_item(item) + processed = _process_images(image_processor, images, return_tensors="pt") + tensors = {str(k): _tensorize(v) for k, v in dict(processed).items()} + if "grid_thws" not in tensors: + raise ValueError("Kimi processor did not return grid_thws") + actual_grids = tensors["grid_thws"].reshape(-1, 3).tolist() + for idx, (item, actual_grid) in enumerate(zip(items, actual_grids, strict=True)): + expected = _grid_payload(item) + if actual_grid != expected: + raise ValueError(f"Kimi grid mismatch at index {idx}: expected {expected}, got {actual_grid}") + return MaterializedMM(kwargs=tensors, forward_policy=self.forward_policy) + + def materialize_for_vllm( + self, + image_processor: Any, + item: RawMMItem, + image: Any, + expected_placeholder_length: int, + ) -> Any: + from vllm.multimodal.inputs import MultiModalFieldConfig, MultiModalKwargsItems + + self.validate_item(item) + hf_inputs = _process_images(image_processor, [image], return_tensors="pt") + tensors = {str(k): _tensorize(v) for k, v in dict(hf_inputs).items()} + expected_grid = _grid_payload(item) + actual_grid = tensors["grid_thws"].reshape(-1, 3).tolist()[0] + if actual_grid != expected_grid: + raise ValueError(f"Kimi grid mismatch: expected {expected_grid}, got {actual_grid}") + if expected_placeholder_length != 1: + raise ValueError(f"Kimi image placeholder length mismatch: expected {expected_placeholder_length}, got 1") + grid_sizes = tensors["grid_thws"].reshape(-1, 3).prod(-1) + config_by_key = { + "pixel_values": MultiModalFieldConfig.flat_from_sizes("vision_chunk", grid_sizes), + "grid_thws": MultiModalFieldConfig.batched("vision_chunk"), + } + return MultiModalKwargsItems.from_hf_inputs(tensors, config_by_key)["vision_chunk"][0] diff --git a/src/prime_rl/multimodal/adapters/qwen_vl.py b/src/prime_rl/multimodal/adapters/qwen_vl.py new file mode 100644 index 0000000000..d793419d2b --- /dev/null +++ b/src/prime_rl/multimodal/adapters/qwen_vl.py @@ -0,0 +1,88 @@ +from __future__ import annotations + +from typing import Any + +from prime_rl.multimodal.adapters.base import ForwardPolicy, MaterializedMM +from prime_rl.multimodal.schema import RawMMItem + + +def _tensorize(value: Any): + import torch + + if isinstance(value, torch.Tensor): + return value.contiguous() + return torch.as_tensor(value).contiguous() + + +def _grid_payload(item: RawMMItem) -> list[int]: + grid = item.payload.get("image_grid_thw") + if grid is None: + raise ValueError("Qwen raw descriptor payload is missing image_grid_thw") + if not isinstance(grid, list | tuple): + raise ValueError(f"Invalid Qwen image_grid_thw: {grid!r}") + if len(grid) == 1 and isinstance(grid[0], list | tuple): + grid = grid[0] + if not isinstance(grid, list | tuple) or len(grid) != 3: + raise ValueError(f"Invalid Qwen image_grid_thw: {grid!r}") + out = [int(v) for v in grid] + if any(v <= 0 for v in out): + raise ValueError(f"Invalid Qwen image_grid_thw: {grid!r}") + return out + + +class QwenVLAdapter: + family = "qwen_vl" + forward_policy = ForwardPolicy( + pass_position_ids_with_mm=False, + requires_mm_token_type_ids=True, + ) + + def validate_item(self, item: RawMMItem) -> None: + if item.family != self.family: + raise ValueError(f"Qwen adapter cannot handle family {item.family!r}") + _grid_payload(item) + + def materialize_for_trainer( + self, + image_processor: Any, + items: list[RawMMItem], + images: list[Any], + ) -> MaterializedMM: + for item in items: + self.validate_item(item) + processed = image_processor(images=images, return_tensors="pt") + tensors = {str(k): _tensorize(v) for k, v in dict(processed).items()} + if "image_grid_thw" not in tensors: + raise ValueError("Qwen processor did not return image_grid_thw") + actual_grids = tensors["image_grid_thw"].tolist() + for idx, (item, actual_grid) in enumerate(zip(items, actual_grids, strict=True)): + expected = _grid_payload(item) + if actual_grid != expected: + raise ValueError(f"Image grid mismatch at index {idx}: expected {expected}, got {actual_grid}") + return MaterializedMM(kwargs=tensors, forward_policy=self.forward_policy) + + def materialize_for_vllm( + self, + image_processor: Any, + item: RawMMItem, + image: Any, + expected_placeholder_length: int, + ) -> Any: + from vllm.model_executor.models.qwen2_vl import _create_qwen2vl_field_factory + from vllm.multimodal.inputs import MultiModalKwargsItems + + self.validate_item(item) + hf_inputs = image_processor(images=[image], return_tensors="pt") + merge_size = int(image_processor.merge_size) + config_by_key = _create_qwen2vl_field_factory(merge_size)(hf_inputs) + mm_item = MultiModalKwargsItems.from_hf_inputs(hf_inputs, config_by_key)["image"][0] + expected_grid = _grid_payload(item) + actual_grid = mm_item["image_grid_thw"].data.tolist() + if actual_grid != expected_grid: + raise ValueError(f"Image grid mismatch: expected {expected_grid}, got {actual_grid}") + num_image_tokens = int(expected_grid[0] * expected_grid[1] * expected_grid[2] // (merge_size * merge_size)) + if expected_placeholder_length != num_image_tokens: + raise ValueError( + f"Image placeholder length mismatch: expected {expected_placeholder_length}, got {num_image_tokens}" + ) + return mm_item diff --git a/src/prime_rl/multimodal/registry.py b/src/prime_rl/multimodal/registry.py new file mode 100644 index 0000000000..88ccab849f --- /dev/null +++ b/src/prime_rl/multimodal/registry.py @@ -0,0 +1,17 @@ +from __future__ import annotations + +from prime_rl.multimodal.adapters.base import MultimodalAdapter +from prime_rl.multimodal.adapters.kimi_k25 import KimiK25Adapter +from prime_rl.multimodal.adapters.qwen_vl import QwenVLAdapter + +_ADAPTERS: dict[str, MultimodalAdapter] = { + QwenVLAdapter.family: QwenVLAdapter(), + KimiK25Adapter.family: KimiK25Adapter(), +} + + +def get_multimodal_adapter(family: str) -> MultimodalAdapter: + try: + return _ADAPTERS[family] + except KeyError as exc: + raise NotImplementedError(f"No multimodal adapter registered for family {family!r}") from exc diff --git a/src/prime_rl/multimodal/schema.py b/src/prime_rl/multimodal/schema.py new file mode 100644 index 0000000000..f26ec1673b --- /dev/null +++ b/src/prime_rl/multimodal/schema.py @@ -0,0 +1,49 @@ +from __future__ import annotations + +from collections.abc import Mapping +from dataclasses import dataclass +from typing import Any + +from renderers.mm_store import RAW_MM_ITEM_KIND + + +@dataclass(frozen=True) +class RawMMItem: + family: str + raw_image_uri: str + payload: dict[str, Any] + + +def _descriptor_mapping(value: Any) -> Mapping[str, Any]: + if isinstance(value, Mapping): + return value + raise TypeError(f"v1 multimodal sidecars must be raw descriptor dicts, got {type(value).__name__}") + + +def _required_str(value: Mapping[str, Any], field: str) -> str: + item = value.get(field) + if isinstance(item, str) and item: + return item + raise ValueError(f"raw multimodal descriptor is missing {field}") + + +def _payload(value: Mapping[str, Any]) -> dict[str, Any]: + payload = value.get("payload") + if not isinstance(payload, Mapping): + raise ValueError("raw multimodal descriptor payload must be a dict") + return {str(k): v for k, v in payload.items()} + + +def _validate_envelope(value: Mapping[str, Any]) -> None: + if value.get("kind") != RAW_MM_ITEM_KIND: + raise ValueError("raw multimodal descriptor is missing the common envelope kind") + + +def parse_raw_mm_item(value: Any) -> RawMMItem: + descriptor = _descriptor_mapping(value) + _validate_envelope(descriptor) + return RawMMItem( + family=_required_str(descriptor, "family"), + raw_image_uri=_required_str(descriptor, "raw_image_uri"), + payload=_payload(descriptor), + ) diff --git a/src/prime_rl/orchestrator/trajectories.py b/src/prime_rl/orchestrator/trajectories.py index 8053453ac8..e62b47e13d 100644 --- a/src/prime_rl/orchestrator/trajectories.py +++ b/src/prime_rl/orchestrator/trajectories.py @@ -9,8 +9,8 @@ Training is renderer-only across every mode (RL/OPD student, SFT teacher), so every node always carries its tokens — no backfill needed. For multimodal rollouts the branch also carries -the images it introduced (`branch.multi_modal_data`), rebuilt here into the flat `mm_kwargs` / -`mm_token_type_ids` the trainer forwards. +raw image refs and renderer descriptors (`branch.multi_modal_data`), preserved here as +`mm_refs` for trainer-side materialization. """ from __future__ import annotations @@ -21,32 +21,9 @@ import verifiers.v1 as vf from prime_rl.transport import TrainingSample -from prime_rl.transport.types import EncodedTensor, RoutedExperts +from prime_rl.transport.types import MMRefs, RoutedExperts from prime_rl.utils.logger import get_logger - - -def _to_numpy(val) -> np.ndarray: - """A renderer mm item value (torch tensor or numpy array) -> a contiguous numpy array.""" - if hasattr(val, "detach"): # torch tensor - val = val.detach().cpu().numpy() - return np.ascontiguousarray(val) - - -def _encode_mm_kwargs(mm_items: dict[str, list[dict]]) -> dict[str, EncodedTensor] | None: - """Concatenate the branch's per-image renderer items into the flat `mm_kwargs` the trainer - forwards — one `EncodedTensor` per kwarg key (e.g. `pixel_values`, `image_grid_thw`), images - cat'd along dim 0 in branch token order. Model-agnostic: the keys are whatever the processor - emits. Returns None when there are no items.""" - bins: dict[str, list[np.ndarray]] = {} - for items in mm_items.values(): # per modality - for item in items: # per image - for key, val in item.items(): - bins.setdefault(key, []).append(_to_numpy(val)) - encoded: dict[str, EncodedTensor] = {} - for key, arrs in bins.items(): - arr = np.concatenate(arrs, axis=0) - encoded[key] = EncodedTensor(dtype=str(arr.dtype), shape=list(arr.shape), data=arr.tobytes()) - return encoded or None +from prime_rl.utils.mm import build_mm_refs def _encode_routed_experts(arr: np.ndarray | None, num_tokens: int) -> RoutedExperts | None: @@ -66,6 +43,23 @@ def _encode_routed_experts(arr: np.ndarray | None, num_tokens: int) -> RoutedExp return RoutedExperts(data=arr.tobytes(), shape=list(arr.shape), dtype=str(arr.dtype)) +def _validate_image_spans(mm_refs: MMRefs, mm_token_type_ids: list[int]) -> None: + """Every image ref's placeholder span must land on image-typed tokens. + + Placeholder offsets flow through the renderer, bridge extension, and node + attribution before arriving here — and the trainer truncates on them — so + any drift anywhere upstream must fail loudly before the sample ships. + """ + for image in mm_refs.images: + span = mm_token_type_ids[image.offset : image.offset + image.length] + if len(span) != image.length or any(t != 1 for t in span): + raise ValueError( + f"Raw image placeholder [{image.offset}, {image.offset + image.length}) does not " + f"cover image-typed tokens (branch length {len(mm_token_type_ids)}) — placeholder " + "offsets have drifted from the branch token stream" + ) + + def iter_trainable_branches(trace: vf.Trace) -> Iterator[tuple[vf.Branch, list[bool]]]: """Yield each branch that yields a training sample, with its trainable-token mask. @@ -101,21 +95,23 @@ def trace_to_samples( `branch.sampled_mask` / `branch.logprobs`), so a sample carries it directly: `mask` marks the trainable (model-sampled) tokens, the context tokens between completions stay masked out. Errored rollouts are dropped upstream (`TrainSink.process_rollout`), so no error - handling happens here. A branch carrying images also gets `mm_kwargs` (the concatenated - pixel tensors) and `mm_token_type_ids` (the renderer's `mm_token_type_id_map` applied to - the branch tokens). Branches with no sampled tokens (e.g. an openai client carrying none) - yield nothing. + handling happens here. A branch carrying images also gets `mm_refs` (raw image URIs + + JSON-safe renderer metadata) and `mm_token_type_ids` (the renderer's + `mm_token_type_id_map` applied to the branch tokens). Branches with no sampled tokens + (e.g. an openai client carrying none) yield nothing. """ samples: list[TrainingSample] = [] for branch, mask in iter_trainable_branches(trace): token_ids = branch.token_ids - mm_kwargs: dict[str, EncodedTensor] | None = None + mm_refs: MMRefs | None = None mm_token_type_ids: list[int] | None = None mmd = branch.multi_modal_data if mmd is not None: - mm_kwargs = _encode_mm_kwargs(mmd.mm_items) + mm_refs = build_mm_refs(mmd) mapping = mm_token_type_ids_mapping or {} mm_token_type_ids = [mapping.get(t, 0) for t in token_ids] + if mm_refs is not None and mapping: + _validate_image_spans(mm_refs, mm_token_type_ids) samples.append( TrainingSample( token_ids=token_ids, @@ -123,7 +119,7 @@ def trace_to_samples( logprobs=branch.logprobs, temperatures=[], # filled by TrainSink.process_group env_name=env_name, - mm_kwargs=mm_kwargs, + mm_refs=mm_refs, mm_token_type_ids=mm_token_type_ids, routed_experts=_encode_routed_experts(branch.routed_experts, len(token_ids)), ) diff --git a/src/prime_rl/templates/multi_node_rl.sbatch.j2 b/src/prime_rl/templates/multi_node_rl.sbatch.j2 index fe73a77373..70c7ef2f47 100755 --- a/src/prime_rl/templates/multi_node_rl.sbatch.j2 +++ b/src/prime_rl/templates/multi_node_rl.sbatch.j2 @@ -59,6 +59,18 @@ export PROJECT_DIR={{ project_dir }} export CONFIG_DIR={{ config_dir }} export OUTPUT_DIR={{ output_dir }} export ORCHESTRATOR_OUTPUT_DIR={{ orchestrator_output_dir }} +{% if image_offload_dir %} +export VF_RENDERER_IMAGE_OFFLOAD_DIR="{{ image_offload_dir }}" +{% else %} +if [ -z "${VF_RENDERER_IMAGE_OFFLOAD_DIR:-}" ]; then + if [ -n "${RUN_ID:-}" ]; then + RUN_ASSET_ID="${RUN_ID#run_}" + export VF_RENDERER_IMAGE_OFFLOAD_DIR="/data/outputs/run_${RUN_ASSET_ID}/assets/images" + else + export VF_RENDERER_IMAGE_OFFLOAD_DIR="$ORCHESTRATOR_OUTPUT_DIR/assets/images" + fi +fi +{% endif %} mkdir -p $OUTPUT_DIR/logs/trainer $OUTPUT_DIR/logs/inference rm -f $OUTPUT_DIR/logs/inference/*.log ln -sfn trainer/node_0.log $OUTPUT_DIR/logs/trainer.log diff --git a/src/prime_rl/trainer/batch.py b/src/prime_rl/trainer/batch.py index 4eb8ff4baa..a5be16d8b0 100644 --- a/src/prime_rl/trainer/batch.py +++ b/src/prime_rl/trainer/batch.py @@ -2,10 +2,11 @@ from collections.abc import Callable, Sequence from dataclasses import dataclass +import msgspec import numpy as np from prime_rl.trainer.utils import balanced_partition -from prime_rl.transport.types import EncodedTensor, MicroBatch, RoutedExperts, TrainingSample +from prime_rl.transport.types import EncodedTensor, MicroBatch, MMImageRef, MMRefs, RoutedExperts, TrainingSample # Backfill value per component weight stream when a packed sample doesn't # carry it: absent rl means weight 1.0 on the loss mask, absent ce/ref_kl @@ -42,56 +43,23 @@ def _pad_routed_experts(micro_batch: MicroBatch, padding_size: int) -> None: routed_experts.shape[0] += padding_size -def _slice_encoded(tensor: EncodedTensor, n_rows: int) -> EncodedTensor: - """First `n_rows` rows of a dim-0-stacked encoded tensor (e.g. pixel_values, image_grid_thw).""" - row = int(np.prod(tensor.shape[1:])) if len(tensor.shape) > 1 else 1 - itemsize = np.dtype(tensor.dtype).itemsize - return EncodedTensor( - dtype=tensor.dtype, - shape=[n_rows, *tensor.shape[1:]], - data=tensor.data[: n_rows * row * itemsize], - ) - - -def _truncate_mm( - mm_token_type_ids: list[int], mm_kwargs: dict[str, EncodedTensor], seq_len: int -) -> tuple[int, dict[str, EncodedTensor] | None]: - """Truncating a sample must not split an image's placeholder block, else the surviving image - token count no longer matches the image embeddings in `mm_kwargs`. Returns the cut point - (<= seq_len, never inside an image block) and `mm_kwargs` sliced to the images whose - placeholders fully survive (None if no image survives).""" - grid = np.frombuffer(bytearray(mm_kwargs["image_grid_thw"].data), dtype=mm_kwargs["image_grid_thw"].dtype).reshape( - mm_kwargs["image_grid_thw"].shape - ) - patches_per_image = [int(g.prod()) for g in grid] - total_patches = mm_kwargs["pixel_values"].shape[0] - total_tokens = sum(1 for t in mm_token_type_ids if t) - ppt = total_patches // total_tokens if total_tokens else 1 # patches per token (merge^2) - tokens_per_image = [p // ppt for p in patches_per_image] - - surviving = sum(1 for t in mm_token_type_ids[:seq_len] if t) - kept = acc = 0 - for n in tokens_per_image: - if acc + n > surviving: - break - acc += n - kept += 1 - if acc == surviving: - cut = seq_len # surviving image tokens are exactly `kept` whole images - else: - # `surviving` lands inside image `kept`; cut to its first placeholder, dropping it. - seen, cut = 0, seq_len - for i, t in enumerate(mm_token_type_ids): - if t: - seen += 1 - if seen == acc + 1: - cut = i - break - if not kept: +def _truncate_mm_refs(mm_refs: MMRefs, seq_len: int) -> tuple[int, MMRefs | None]: + """Return a token cut that never splits an image placeholder, plus the surviving refs.""" + cut, kept = seq_len, 0 + for image in mm_refs.images: # token order, non-overlapping — enforced by build_mm_refs + if image.offset + image.length <= seq_len: + kept += 1 + continue + if image.offset < seq_len: + cut = image.offset + break + if cut == 0: + raise ValueError(f"Cannot truncate multimodal sample: leading image does not fit in seq_len={seq_len}") + if kept == len(mm_refs.images): + return seq_len, mm_refs + if kept == 0: return cut, None - kept_patches = sum(patches_per_image[:kept]) - sliced = {k: _slice_encoded(v, kept if k == "image_grid_thw" else kept_patches) for k, v in mm_kwargs.items()} - return cut, sliced + return cut, MMRefs(images=mm_refs.images[:kept]) def multimodal_sample_error(sample: TrainingSample) -> str | None: @@ -101,6 +69,10 @@ def multimodal_sample_error(sample: TrainingSample) -> str | None: "mm_token_type_ids length must match token_ids length " f"({len(mm_token_type_ids)} != {len(sample.token_ids)})" ) + if sample.mm_refs is not None and mm_token_type_ids is None: + # The orchestrator always stamps mm_token_type_ids alongside mm_refs; the + # trainer's image-boundary truncation and forward policies rely on them. + return "raw multimodal samples require mm_token_type_ids" if sample.mm_kwargs is not None and "image_grid_thw" in sample.mm_kwargs and mm_token_type_ids is None: return "image_grid_thw multimodal samples require mm_token_type_ids" return None @@ -134,7 +106,9 @@ def prepare_sample(training_example: TrainingSample, seq_len: int) -> MicroBatch ref_kl_weights = list(training_example.ref_kl_weights) if training_example.ref_kl_weights is not None else None position_ids = list(range(len(input_ids))) mm_token_type_ids = training_example.mm_token_type_ids - mm_kwargs = training_example.mm_kwargs + if training_example.mm_kwargs is not None: + raise ValueError("Processed multimodal mm_kwargs are unsupported in v1; use raw mm_refs") + mm_refs = training_example.mm_refs assert training_example.env_name != "all", "env_name='all' is reserved for aggregate metric keys" env_names = [training_example.env_name] * len(input_ids) @@ -149,11 +123,9 @@ def prepare_sample(training_example: TrainingSample, seq_len: int) -> MicroBatch ) if len(input_ids) > seq_len: - # Multimodal: never split an image's placeholder block — cut to a whole-image boundary - # and slice mm_kwargs to match, so image-token count == image-embedding count. cut = seq_len - if mm_token_type_ids is not None and mm_kwargs is not None: - cut, mm_kwargs = _truncate_mm(mm_token_type_ids, mm_kwargs, seq_len) + if mm_refs is not None: + cut, mm_refs = _truncate_mm_refs(mm_refs, seq_len) input_ids = input_ids[:cut] loss_mask = loss_mask[:cut] inference_logprobs = inference_logprobs[:cut] @@ -214,7 +186,7 @@ def prepare_sample(training_example: TrainingSample, seq_len: int) -> MicroBatch routed_experts=routed_experts, mm_token_type_ids=mm_token_type_ids, env_names=env_names, - mm_kwargs=mm_kwargs, + mm_refs=mm_refs, rl_weights=rl_weights, ce_weights=ce_weights, ref_kl_weights=ref_kl_weights, @@ -224,7 +196,16 @@ def prepare_sample(training_example: TrainingSample, seq_len: int) -> MicroBatch def _is_multimodal_sample(sample: MicroBatch) -> bool: """Check if a sample contains multimodal data (images).""" - return sample.mm_kwargs is not None + return sample.mm_refs is not None or sample.mm_kwargs is not None + + +def _mm_refs_family(mm_refs: MMRefs) -> str | None: + """The adapter family of a sample's raw image refs. + + ``build_mm_refs`` validates every descriptor and the materializer enforces + exactly one family per micro batch, so the first image's family stands for + the sample.""" + return mm_refs.images[0].item.get("family") if mm_refs.images else None @dataclass @@ -268,6 +249,12 @@ def can_add(self, sample: MicroBatch, max_seq_len: int, lora_idx: int) -> bool: if self.first_lora_idx != lora_idx: return False if existing_mm_sample is not None and sample_is_mm: + if (existing_mm_sample.mm_refs is None) != (sample.mm_refs is None): + return False + if sample.mm_refs is not None: + # Raw refs materialize downstream through exactly one adapter family + # per micro batch; within a family, grids and image sizes vary freely. + return _mm_refs_family(existing_mm_sample.mm_refs) == _mm_refs_family(sample.mm_refs) dst = existing_mm_sample.mm_kwargs src = sample.mm_kwargs assert dst is not None and src is not None, "multimodal samples must carry mm_kwargs" @@ -327,10 +314,12 @@ def _materialize_bin(bin_content: _MicroBatchBin, num_loras: int) -> MicroBatch: streams: dict[str, list[float] | None] = {name: ([] if has_stream[name] else None) for name in STREAM_FILL} seq_lens: list[int] = [] routed_experts: RoutedExperts | None = None + mm_ref_images: list[MMImageRef] = [] lora_num_tokens = [0] * num_loras for lora_idx, sample in bin_content.samples: sample_len = len(sample.input_ids) + sample_start = len(input_ids) input_ids.extend(sample.input_ids) loss_mask.extend(sample.loss_mask) advantages.extend(sample.advantages) @@ -364,6 +353,11 @@ def _materialize_bin(bin_content: _MicroBatchBin, num_loras: int) -> MicroBatch: for key in mm_kwargs: mm_kwargs[key].data += sample.mm_kwargs[key].data mm_kwargs[key].shape[0] += sample.mm_kwargs[key].shape[0] + if sample.mm_refs is not None: + # Placeholder offsets are sample-relative; rebase them to the packed stream. + mm_ref_images.extend( + msgspec.structs.replace(image, offset=image.offset + sample_start) for image in sample.mm_refs.images + ) seq_lens.extend(sample.seq_lens) lora_num_tokens[lora_idx] += sample_len @@ -384,6 +378,7 @@ def _materialize_bin(bin_content: _MicroBatchBin, num_loras: int) -> MicroBatch: routed_experts=routed_experts, mm_token_type_ids=mm_token_type_ids, env_names=env_names, + mm_refs=MMRefs(images=mm_ref_images) if mm_ref_images else None, mm_kwargs=mm_kwargs, rl_weights=streams["rl_weights"], ce_weights=streams["ce_weights"], @@ -422,8 +417,9 @@ def packed_samples_into_micro_bs( With per-token temperatures, samples can be packed together regardless of their temperature values. Multimodal samples pack with text spans from the same run/LoRA and with - compatible eager ``mm_kwargs`` samples. Packed batches preserve sample - boundaries in ``seq_lens``. + same-family raw-ref samples (image ref offsets are rebased to the packed + stream in ``_materialize_bin``). Packed batches preserve sample boundaries + in ``seq_lens``. """ # Sort by (lora_idx, -length) for packing efficiency samples.sort(key=lambda x: (x[0], -len(x[1].input_ids))) diff --git a/src/prime_rl/trainer/model.py b/src/prime_rl/trainer/model.py index d2d6703819..34e03dd9a8 100644 --- a/src/prime_rl/trainer/model.py +++ b/src/prime_rl/trainer/model.py @@ -34,6 +34,7 @@ MXFP8Config, TokenizerConfig, ) +from prime_rl.multimodal.adapters.base import ForwardPolicy from prime_rl.trainer.distributed import DeepEPExpertParallel, MXFP8AllToAllExpertParallel from prime_rl.trainer.lora import apply_lora_to_model, freeze_all_except_lora_and_specified, strip_lora_from_state_dict from prime_rl.trainer.models import ( @@ -1373,12 +1374,14 @@ def forward( labels: Int[Tensor, "batch seq"] | None = None, temperature: Tensor | None = None, routed_experts: Int[Tensor, "batch seq layers topk"] | None = None, - # Generic multimodal kwargs (e.g. {"pixel_values": ..., - # "image_grid_thw": ...} for Qwen3-VL; just {"pixel_values": ...} - # for Gemma3). Passed straight through to ``model(**kwargs)`` so - # the model's HF forward signature is the schema. ``mm_token_type_ids`` - # is split out because it comes from the renderer rather than the processor. + # Generic multimodal kwargs materialized by the trainer's model processor + # (e.g. {"pixel_values": ..., "image_grid_thw": ...} for Qwen3-VL; just + # {"pixel_values": ...} for Gemma3). Passed straight through to + # ``model(**kwargs)`` so the model's HF forward signature is the schema. + # ``mm_token_type_ids`` is split out because it's prime-rl-computed from + # renderer token ids, not part of the processor's own output. mm_kwargs: dict[str, Tensor] | None = None, + mm_forward_policy: ForwardPolicy | None = None, mm_token_type_ids: Int[Tensor, "batch seq"] | None = None, # True when seq_lens holds the full pre-CP-shard document boundaries # (kept global because documents can straddle the shard cut). @@ -1390,14 +1393,19 @@ def forward( "temperature": temperature, } - if mm_kwargs: + if mm_kwargs is not None: # Forward the per-model multimodal tensors verbatim, plus the - # renderer-supplied ``mm_token_type_ids`` (renderer owns the - # token→modality mapping via ``mm_token_type_id_map``). + # token→modality map derived from renderer token ids. kwargs.update(mm_kwargs) if mm_token_type_ids is not None: kwargs["mm_token_type_ids"] = mm_token_type_ids - if "image_grid_thw" not in mm_kwargs: + # Callers without an adapter policy (SFT) fall back to key presence: + # models whose kwargs carry image_grid_thw (Qwen-VL family) build + # their own MRoPE position ids and must not receive packed 1D ones. + policy = mm_forward_policy or ForwardPolicy(pass_position_ids_with_mm="image_grid_thw" not in mm_kwargs) + if policy.requires_mm_token_type_ids and mm_token_type_ids is None: + raise ValueError("Multimodal forward policy requires mm_token_type_ids") + if policy.pass_position_ids_with_mm: kwargs["position_ids"] = position_ids else: kwargs["position_ids"] = position_ids diff --git a/src/prime_rl/trainer/rl/data.py b/src/prime_rl/trainer/rl/data.py index 1ad250cf72..0f09e00fea 100644 --- a/src/prime_rl/trainer/rl/data.py +++ b/src/prime_rl/trainer/rl/data.py @@ -1,3 +1,4 @@ +import time from collections.abc import Callable, Sequence from pathlib import Path from typing import TypedDict @@ -7,6 +8,7 @@ from torch import Tensor from prime_rl.configs.trainer import FakeDataLoaderConfig +from prime_rl.multimodal.adapters.base import ForwardPolicy, MaterializedMM from prime_rl.trainer.rl.packer import BasePacker, setup_packer from prime_rl.trainer.runs import get_multi_run_manager from prime_rl.trainer.world import get_world @@ -16,6 +18,9 @@ TransportConfig, setup_micro_batch_receiver, ) +from prime_rl.transport.types import MMRefs +from prime_rl.utils.logger import get_logger +from prime_rl.utils.mm import RawImageMaterializer class TensorMicroBatch(TypedDict): @@ -39,12 +44,9 @@ class TensorMicroBatch(TypedDict): # MoE router replay routed_experts: Int[Tensor, "batch seq layers topk"] | None - # Generic multimodal kwargs — flat dict matching the model's forward - # signature (e.g. ``{"pixel_values": ..., "image_grid_thw": ...}`` for - # Qwen3-VL; ``{"pixel_values": ...}`` for Gemma3-VL). The trainer - # ``**`` -unpacks this into the forward call, so any HF VLM whose - # processor and forward agree on kwarg names works out of the box. + # Generic multimodal kwargs materialized by the trainer's model processor. mm_kwargs: dict[str, Tensor] | None + mm_forward_policy: ForwardPolicy | None # mm_token_type_ids: token type per token [batch seq], int64 (0=text, 1=image, 2=video) mm_token_type_ids: Int[Tensor, "batch seq"] | None @@ -59,6 +61,9 @@ class TensorMicroBatch(TypedDict): run_step: int | None +MaterializedMMParts = tuple[dict[str, Tensor] | None, ForwardPolicy | None] + + class FakeDataLoader: def __init__(self, config: FakeDataLoaderConfig, seq_len: int, dp_world_size: int): self.world = get_world() @@ -72,6 +77,8 @@ def __init__(self, config: FakeDataLoaderConfig, seq_len: int, dp_world_size: in self.generate_samples = config.generate_samples self.batch_counter = 0 self.multi_run_manager = get_multi_run_manager() + self.last_mm_materialize_time = 0.0 + self.last_mm_images_materialized = 0 def wait_for_batch(self) -> None: return @@ -133,6 +140,7 @@ def _get_sample_micro_batch(self, generator: torch.Generator) -> TensorMicroBatc "seq_lens": torch.tensor(sequence_lengths, dtype=torch.long), "routed_experts": None, "mm_kwargs": None, + "mm_forward_policy": None, "mm_token_type_ids": None, "rl_weights": None, "ce_weights": None, @@ -166,6 +174,7 @@ def _get_micro_batch(self, generator: torch.Generator) -> TensorMicroBatch: "seq_lens": torch.tensor([self.seq_len], dtype=torch.long), "routed_experts": None, "mm_kwargs": None, + "mm_forward_policy": None, "mm_token_type_ids": None, "rl_weights": None, "ce_weights": None, @@ -187,6 +196,8 @@ def __init__( pad_to_multiple_of: int, bin_cost: Callable[[Sequence[int]], int], config: TransportConfig, + model_name: str, + model_trust_remote_code: bool, ): self.world = get_world() @@ -205,6 +216,8 @@ def __init__( self.multi_run_manager = get_multi_run_manager() self.receiver: MicroBatchReceiver = setup_micro_batch_receiver(output_dir, dp_rank, start_step, config) + self.mm_materializer = RawImageMaterializer(model_name, trust_remote_code=model_trust_remote_code) + self._reset_mm_stats() def wait_for_batch(self) -> None: if self.world.is_master: @@ -218,22 +231,50 @@ def wait_for_batch(self) -> None: def get_batch(self) -> list[TensorMicroBatch]: micro_batches = self.receiver.receive() + self._reset_mm_stats() return [self._micro_batch_to_tensor(mb) for mb in micro_batches] + def _reset_mm_stats(self) -> None: + self.last_mm_materialize_time = 0.0 + self.last_mm_images_materialized = 0 + + @staticmethod + def _materialized_mm_parts(materialized: MaterializedMM | None) -> MaterializedMMParts: + if materialized is None: + return None, None + return materialized.kwargs, materialized.forward_policy + + @staticmethod + def _mm_run_context(micro_batch: MicroBatch) -> str: + run_idx = next((i for i, n in enumerate(micro_batch.lora_num_tokens or []) if n > 0), None) + return f"run_idx={run_idx}, run_id={micro_batch.run_id}, run_step={micro_batch.run_step}" + + def _materialize_mm_refs(self, micro_batch: MicroBatch, refs: MMRefs) -> MaterializedMMParts: + materialize_start = time.perf_counter() + try: + materialized = self.mm_materializer.materialize(refs) + except FileNotFoundError as exc: + self.last_mm_materialize_time += time.perf_counter() - materialize_start + get_logger().error( + f"raw image materialization failed ({self._mm_run_context(micro_batch)}, uris={refs.uris}): {exc!r}" + ) + raise + + self.last_mm_materialize_time += time.perf_counter() - materialize_start + self.last_mm_images_materialized += len(refs.uris) + return self._materialized_mm_parts(materialized) + def _micro_batch_to_tensor(self, micro_batch: MicroBatch) -> TensorMicroBatch: """Convert a MicroBatch (msgspec struct with lists) to a TensorMicroBatch (dict with tensors).""" if micro_batch.lora_num_tokens is None: micro_batch.lora_num_tokens = [0] * self.multi_run_manager.max_runs micro_batch.lora_num_tokens[0] = len(micro_batch.input_ids) mm_kwargs: dict[str, Tensor] | None = None - if micro_batch.mm_kwargs: - # Each value is an EncodedTensor (dtype, shape, raw bytes). - # No batch dim — the orchestrator concatenates per-image along - # dim=0 generically, matching what each HF VLM's forward expects. - mm_kwargs = { - key: torch.frombuffer(bytearray(payload.data), dtype=_torch_dtype(payload.dtype)).reshape(payload.shape) - for key, payload in micro_batch.mm_kwargs.items() - } + mm_forward_policy: ForwardPolicy | None = None + if micro_batch.mm_kwargs is not None: + raise ValueError("Processed multimodal mm_kwargs are unsupported in v1; use raw mm_refs") + if micro_batch.mm_refs is not None: + mm_kwargs, mm_forward_policy = self._materialize_mm_refs(micro_batch, micro_batch.mm_refs) routed_experts = None packed_routed_experts = micro_batch.routed_experts if packed_routed_experts is not None: @@ -261,6 +302,7 @@ def _micro_batch_to_tensor(self, micro_batch: MicroBatch) -> TensorMicroBatch: lora_num_tokens=torch.tensor(micro_batch.lora_num_tokens, dtype=torch.int32), seq_lens=torch.tensor(micro_batch.seq_lens, dtype=torch.long), mm_kwargs=mm_kwargs, + mm_forward_policy=mm_forward_policy, mm_token_type_ids=torch.tensor(micro_batch.mm_token_type_ids, dtype=torch.long).unsqueeze(0) if micro_batch.mm_token_type_ids is not None else None, diff --git a/src/prime_rl/trainer/rl/train.py b/src/prime_rl/trainer/rl/train.py index 295d3ad15e..0b5149a9a2 100644 --- a/src/prime_rl/trainer/rl/train.py +++ b/src/prime_rl/trainer/rl/train.py @@ -251,6 +251,8 @@ def load_run_checkpoint(_optimizer, idx: int) -> None: config.model.cp, build_bin_cost(model.config), config.rollout_transport, + config.model.name, + config.model.trust_remote_code, ) token_exporter = setup_token_exporter(config, parallel_dims, world, logger) @@ -365,12 +367,10 @@ def load_run_checkpoint(_optimizer, idx: int) -> None: # we could've gotten routed experts from the inference server, but we didn't enable router replay routed_experts = None - # Multimodal kwargs are an opaque per-model dict (e.g. - # {"pixel_values": ..., "image_grid_thw": ...} for Qwen3-VL, - # just {"pixel_values": ...} for Gemma3-VL) — we move every - # tensor to CUDA and let the model's forward sort them. + # Multimodal kwargs are materialized by the trainer model processor. mm_kwargs_raw = micro_batch.get("mm_kwargs") mm_kwargs = {k: v.to("cuda") for k, v in mm_kwargs_raw.items()} if mm_kwargs_raw else None + mm_forward_policy = micro_batch.get("mm_forward_policy") if mm_kwargs is not None and config.model.vlm is None: raise ValueError( "Received multimodal samples but [model.vlm] is not set. " @@ -433,6 +433,7 @@ def load_run_checkpoint(_optimizer, idx: int) -> None: labels=labels, temperature=temperatures, mm_kwargs=mm_kwargs, + mm_forward_policy=mm_forward_policy, mm_token_type_ids=mm_token_type_ids, seq_lens=seq_lens, seq_lens_are_pre_shard=seq_lens_are_pre_shard, @@ -717,9 +718,11 @@ def load_run_checkpoint(_optimizer, idx: int) -> None: "time/step": step_time, "time/wait_for_batch": wait_for_batch_time, "time/load_data": load_data_time, + "time/mm_materialize": dataloader.last_mm_materialize_time, "time/broadcast_weights": broadcast_weights_time, "time/save_ckpt": save_ckpt_time, "time/forward_backward": forward_backward_time, + "mm/images_materialized": dataloader.last_mm_images_materialized, "step": progress.step, } monitor.log(time_metrics, step=progress.step) diff --git a/src/prime_rl/transport/types.py b/src/prime_rl/transport/types.py index 121da68d92..c6d8ed490a 100644 --- a/src/prime_rl/transport/types.py +++ b/src/prime_rl/transport/types.py @@ -18,6 +18,36 @@ class RoutedExperts(msgspec.Struct, array_like=True, gc=False, omit_defaults=Tru dtype: str +class MMImageRef(msgspec.Struct, array_like=True, gc=False, omit_defaults=True): + """One raw image reference for trainer-side materialization. + + ``item`` is the JSON-safe renderer descriptor (adapter family + payload, + parseable via ``parse_raw_mm_item``), ``hash``/``uri`` identify the raw image + file, and ``offset``/``length`` are the image's placeholder range in token + space (needed to truncate at image boundaries). + """ + + item: dict + hash: str + uri: str + offset: int + length: int + + +class MMRefs(msgspec.Struct, array_like=True, gc=False, omit_defaults=True): + """Raw multimodal sidecar references for one sample, in token order. + + The trainer materializes the referenced image files with its own processor. + Processed tensors are intentionally not part of this transport. + """ + + images: list[MMImageRef] + + @property + def uris(self) -> list[str]: + return [image.uri for image in self.images] + + # Orchestrator -> Packer class TrainingSample(msgspec.Struct, array_like=True, gc=False, omit_defaults=True): """A single training example — one branch of a rollout as a flat token sequence. @@ -34,16 +64,11 @@ class TrainingSample(msgspec.Struct, array_like=True, gc=False, omit_defaults=Tr env_name: str ref_logprobs: list[float] | None = None # reference-model logprobs (ref_kl component) - # Generic multimodal kwargs: flat dict keyed by the kwarg names the - # model's forward expects (e.g. {"pixel_values": ..., "image_grid_thw": - # ...} for Qwen3-VL; just {"pixel_values": ...} for Gemma3). The - # orchestrator batches per-image renderer items by torch.cat along - # dim=0 generically — no model-specific knowledge in prime-rl. The - # trainer ``**`` -unpacks this into the model forward, so any VLM - # whose HF processor / forward agree on kwarg names works without - # touching this transport. + # Processed multimodal payloads are rejected by the v1 raw-image-ref path. mm_kwargs: dict[str, EncodedTensor] | None = None + mm_refs: MMRefs | None = None + routed_experts: RoutedExperts | None = None # mm_token_type_ids: token type ids per token [batch seq], int64 (0=text, 1=image, 2=video) @@ -94,8 +119,9 @@ class MicroBatch(msgspec.Struct, array_like=True, gc=False, omit_defaults=True): lora_num_tokens: list[int] | None = None routed_experts: RoutedExperts | None = None - # See TrainingSample.mm_kwargs. + # Processed multimodal payloads are rejected by the v1 raw-image-ref path. mm_kwargs: dict[str, EncodedTensor] | None = None + mm_refs: MMRefs | None = None # mm_token_type_ids: token type ids per token [batch seq], int64 (0=text, 1=image, 2=video) mm_token_type_ids: list[int] | None = None diff --git a/src/prime_rl/utils/mm.py b/src/prime_rl/utils/mm.py new file mode 100644 index 0000000000..d46cf93c75 --- /dev/null +++ b/src/prime_rl/utils/mm.py @@ -0,0 +1,168 @@ +from __future__ import annotations + +import hashlib +from collections.abc import Iterable, Mapping +from io import BytesIO +from pathlib import Path +from typing import Any +from urllib.parse import unquote, urlparse + +from prime_rl.multimodal.adapters.base import MaterializedMM, MultimodalAdapter +from prime_rl.multimodal.registry import get_multimodal_adapter +from prime_rl.multimodal.schema import RawMMItem, parse_raw_mm_item +from prime_rl.transport.types import MMImageRef, MMRefs + +IMAGE_MODALITY = "image" +SUPPORTED_MODALITIES = {IMAGE_MODALITY} + + +def _field(value: Any, name: str, default: Any = None) -> Any: + if isinstance(value, Mapping): + return value.get(name, default) + return getattr(value, name, default) + + +def file_uri_to_path(uri: str) -> Path: + parsed = urlparse(uri) + if parsed.scheme != "file": + raise ValueError(f"Raw multimodal image refs must be file:// URIs, got {uri!r}") + if parsed.netloc not in ("", "localhost"): + raise ValueError(f"file:// multimodal refs must be local paths, got {uri!r}") + return Path(unquote(parsed.path)) + + +def _validate_modalities(mm_items: Mapping[str, list[Any]]) -> None: + unsupported = sorted( + modality for modality, items in mm_items.items() if items and modality not in SUPPORTED_MODALITIES + ) + if unsupported: + raise NotImplementedError( + "v1 multimodal training currently supports raw image refs only; " + f"unsupported modalities: {', '.join(unsupported)}" + ) + + +def _raw_item_dicts(items: Iterable[Mapping[str, Any]]) -> tuple[list[dict[str, Any]], list[str]]: + item_dicts: list[dict[str, Any]] = [] + uris: list[str] = [] + for item in items: + item_dict = dict(item) + parsed = parse_raw_mm_item(item_dict) + file_uri_to_path(parsed.raw_image_uri) + item_dicts.append(item_dict) + uris.append(parsed.raw_image_uri) + return item_dicts, uris + + +def _placeholder_bounds(placeholder: Any) -> tuple[int, int]: + offset = _field(placeholder, "offset") + length = _field(placeholder, "length") + if not isinstance(offset, int) or not isinstance(length, int): + raise ValueError(f"Raw image placeholder must have integer offset/length, got {placeholder!r}") + if offset < 0 or length <= 0: + raise ValueError(f"Raw image placeholder must have offset >= 0 and length > 0, got {placeholder!r}") + return offset, length + + +def build_mm_refs(multi_modal_data: Any) -> MMRefs | None: + mm_items = _field(multi_modal_data, "mm_items", None) + if not mm_items: + return None + _validate_modalities(mm_items) + + image_items, uris = _raw_item_dicts(mm_items.get(IMAGE_MODALITY, [])) + if not image_items: + return None + + mm_hashes = _field(multi_modal_data, "mm_hashes", {}) or {} + image_hashes = list(mm_hashes.get(IMAGE_MODALITY, [])) + mm_placeholders = _field(multi_modal_data, "mm_placeholders", {}) or {} + image_placeholders = list(mm_placeholders.get(IMAGE_MODALITY, [])) + if len(image_hashes) != len(image_items) or len(image_placeholders) != len(image_items): + raise ValueError( + "Raw image descriptor/hash/placeholder mismatch: " + f"descriptors={len(image_items)}, hashes={len(image_hashes)}, placeholders={len(image_placeholders)}" + ) + + images: list[MMImageRef] = [] + prev_end = 0 + for item, uri, image_hash, placeholder in zip(image_items, uris, image_hashes, image_placeholders, strict=True): + offset, length = _placeholder_bounds(placeholder) + # Truncation cuts at a prefix of ``images``, which is only sound if + # placeholders arrive in token order without overlap. + if offset < prev_end: + raise ValueError(f"Raw image placeholders must be sorted and non-overlapping, got {image_placeholders!r}") + prev_end = offset + length + images.append(MMImageRef(item=item, hash=image_hash, uri=uri, offset=offset, length=length)) + return MMRefs(images=images) + + +def sha256_32(data: bytes) -> str: + return hashlib.sha256(data).hexdigest()[:32] + + +def read_verified_image_bytes(uri: str, expected_hash: str) -> bytes: + """Read a local ``file://`` image and verify its content hash (``sha256_32``).""" + raw = file_uri_to_path(uri).read_bytes() + actual_hash = sha256_32(raw) + if actual_hash != expected_hash: + raise ValueError(f"Raw image hash mismatch for {uri}: expected {expected_hash}, got {actual_hash}") + return raw + + +def load_image_processor(model_name: str, trust_remote_code: bool = False): + """Load the HF ``image_processor`` for ``model_name``.""" + from transformers import AutoProcessor + + processor = AutoProcessor.from_pretrained(model_name, trust_remote_code=trust_remote_code) + image_processor = getattr(processor, "image_processor", None) + if image_processor is None: + raise ValueError(f"{model_name!r} does not expose an image_processor") + return image_processor + + +def _parse_image_refs(refs: MMRefs) -> list[RawMMItem]: + return [parse_raw_mm_item(image.item) for image in refs.images] + + +def _single_family_adapter(items: list[RawMMItem]) -> MultimodalAdapter: + families = {item.family for item in items} + if len(families) != 1: + raise ValueError(f"Raw multimodal refs must use exactly one adapter family, got {sorted(families)}") + return get_multimodal_adapter(next(iter(families))) + + +def _load_verified_images(image_refs: list[MMImageRef]) -> list[Any]: + from PIL import Image + + images = [] + for ref in image_refs: + raw = read_verified_image_bytes(ref.uri, ref.hash) + with Image.open(BytesIO(raw)) as image: + images.append(image.convert("RGB")) + return images + + +class RawImageMaterializer: + """Materialize raw image refs with the trainer model's HF image processor.""" + + def __init__(self, model_name: str, *, trust_remote_code: bool): + self.model_name = model_name + self.trust_remote_code = trust_remote_code + self._image_processor = None + + @property + def image_processor(self): + if self._image_processor is None: + self._image_processor = load_image_processor(self.model_name, self.trust_remote_code) + return self._image_processor + + def materialize(self, refs: MMRefs) -> MaterializedMM | None: + image_items = _parse_image_refs(refs) + if not image_items: + return None + + image_processor = self.image_processor + adapter = _single_family_adapter(image_items) + images = _load_verified_images(refs.images) + return adapter.materialize_for_trainer(image_processor, image_items, images) diff --git a/src/prime_rl/utils/run_assets.py b/src/prime_rl/utils/run_assets.py new file mode 100644 index 0000000000..d15ede0cf1 --- /dev/null +++ b/src/prime_rl/utils/run_assets.py @@ -0,0 +1,43 @@ +from __future__ import annotations + +import os +from collections.abc import Mapping +from pathlib import Path +from string import Template + +from prime_rl.configs.shared import MultimodalConfig + +# Contract: must match renderers.mm_store.IMAGE_OFFLOAD_DIR_ENV. +IMAGE_OFFLOAD_DIR_ENV = "VF_RENDERER_IMAGE_OFFLOAD_DIR" +RUN_ID_ENV = "RUN_ID" + +RUN_OUTPUT_ROOT = Path("/data/outputs") +IMAGE_ASSET_SUBDIR = Path("assets/images") + + +def _expand_path(path: Path, env: Mapping[str, str]) -> Path: + expanded = Template(os.path.expanduser(str(path))).safe_substitute(env) + return Path(expanded).resolve() + + +def _run_id_dir(env: Mapping[str, str]) -> Path | None: + raw_run_id = env.get(RUN_ID_ENV, "").strip() + if not raw_run_id: + return None + run_id = raw_run_id.removeprefix("run_") + return RUN_OUTPUT_ROOT / f"run_{run_id}" + + +def resolve_image_offload_dir( + output_dir: Path, + multimodal: MultimodalConfig, + env: Mapping[str, str], +) -> Path: + """Resolve image asset dir by precedence: offload_dir, RUN_ID hosted path, then output_dir/assets/images.""" + explicit = multimodal.offload_dir + if explicit is not None: + return _expand_path(explicit, env) + hosted_run_dir = _run_id_dir(env) + if hosted_run_dir is not None: + return (hosted_run_dir / IMAGE_ASSET_SUBDIR).resolve() + return (output_dir.resolve() / IMAGE_ASSET_SUBDIR).resolve() diff --git a/tests/unit/inference/test_serving_tokens.py b/tests/unit/inference/test_serving_tokens.py index 951ef5c9d8..9b01b84a26 100644 --- a/tests/unit/inference/test_serving_tokens.py +++ b/tests/unit/inference/test_serving_tokens.py @@ -11,6 +11,7 @@ from __future__ import annotations import asyncio +import hashlib import numpy as np import pybase64 @@ -267,3 +268,257 @@ def test_client_set_max_tokens_assumes_set_when_body_unreadable(): # non-dict body → can't tell, don't override. assert asyncio.run(_client_set_max_tokens(_FakeRawRequest([1, 2, 3]))) is True + + +def test_materialize_raw_image_ref_uses_generic_family_payload(tmp_path, monkeypatch): + from PIL import Image + from renderers.mm_store import raw_mm_ref + + from prime_rl.inference.vllm import serving_tokens + + image_dir = tmp_path / "run_serving" / "assets" / "images" + image_dir.mkdir(parents=True) + image_path = image_dir / "image.png" + Image.new("RGB", (8, 6), color=(32, 64, 128)).save(image_path) + + mm_hash = hashlib.sha256(image_path.read_bytes()).hexdigest()[:32] + raw_ref = raw_mm_ref( + family="test_family", + modality="image", + mm_hash=mm_hash, + raw_image_uri=image_path.as_uri(), + payload={"adapter_owned": [1, 2, 3]}, + ) + processor = object() + captured = {} + + class _Adapter: + def materialize_for_vllm(self, image_processor, item, image, expected_placeholder_length): + captured["image_processor"] = image_processor + captured["item"] = item + captured["image_size"] = image.size + captured["expected_placeholder_length"] = expected_placeholder_length + return {"materialized": True} + + def _get_adapter(family): + captured["family"] = family + return _Adapter() + + monkeypatch.setattr(serving_tokens, "_load_image_processor", lambda _model, _trust: processor) + monkeypatch.setattr(serving_tokens, "get_multimodal_adapter", _get_adapter) + + out = serving_tokens._materialize_raw_image_ref_sync( + raw_ref, + feature_modality="image", + mm_hash=mm_hash, + expected_placeholder_length=7, + processor_model_name="model", + trust_remote_code=True, + ) + + assert out == {"materialized": True} + assert captured["family"] == "test_family" + assert captured["image_processor"] is processor + assert captured["image_size"] == (8, 6) + assert captured["expected_placeholder_length"] == 7 + item = captured["item"] + assert item.family == "test_family" + assert item.raw_image_uri == image_path.as_uri() + assert item.payload == {"adapter_owned": [1, 2, 3]} + + +def test_materialize_raw_image_ref_maps_adapter_validation_error(tmp_path, monkeypatch): + import pytest + + from prime_rl.inference.vllm import serving_tokens + + features = _mm_features(tmp_path) + raw_ref = features.kwargs_data["image"][0] + mm_hash = features.mm_hashes["image"][0] + + class _InvalidAdapter: + def materialize_for_vllm(self, image_processor, item, image, expected_placeholder_length): + raise ValueError("image grid mismatch") + + monkeypatch.setattr(serving_tokens, "_load_image_processor", lambda _model, _trust: object()) + monkeypatch.setattr(serving_tokens, "get_multimodal_adapter", lambda _family: _InvalidAdapter()) + + with pytest.raises(serving_tokens._MMImageRefError, match="image grid mismatch") as exc_info: + serving_tokens._materialize_raw_image_ref_sync( + raw_ref, + feature_modality="image", + mm_hash=mm_hash, + expected_placeholder_length=7, + processor_model_name="model", + trust_remote_code=False, + ) + + assert exc_info.value.status_code == 400 + + +def _mm_features(tmp_path, *, image_seed: int = 0, placeholder_length: int = 7): + """A minimal ``GenerateRequest.features``-shaped object carrying one real raw ref.""" + from types import SimpleNamespace + + from PIL import Image + from renderers.mm_store import raw_mm_ref + + image_dir = tmp_path / "assets" / "images" + image_dir.mkdir(parents=True, exist_ok=True) + image_path = image_dir / f"image-{image_seed}.png" + Image.new("RGB", (8, 6), color=(image_seed % 255, 64, 128)).save(image_path) + + mm_hash = hashlib.sha256(image_path.read_bytes()).hexdigest()[:32] + raw_ref = raw_mm_ref( + family="test_family", + modality="image", + mm_hash=mm_hash, + raw_image_uri=image_path.as_uri(), + payload={}, + ) + return SimpleNamespace( + mm_hashes={"image": [mm_hash]}, + kwargs_data={"image": [raw_ref]}, + mm_placeholders={"image": [SimpleNamespace(offset=0, length=placeholder_length)]}, + ) + + +def _patch_adapter(monkeypatch, calls: list, materialize=None): + from prime_rl.inference.vllm import serving_tokens + + class _Adapter: + def materialize_for_vllm(self, image_processor, item, image, expected_placeholder_length): + calls.append(item.raw_image_uri) + if materialize is not None: + return materialize(item) + return {"materialized": item.raw_image_uri} + + monkeypatch.setattr(serving_tokens, "_load_image_processor", lambda _model, _trust: object()) + monkeypatch.setattr(serving_tokens, "get_multimodal_adapter", lambda _family: _Adapter()) + + +def test_mm_materialize_cache_hit_skips_work(tmp_path, monkeypatch): + from prime_rl.inference.vllm import serving_tokens + + calls: list = [] + _patch_adapter(monkeypatch, calls) + cache = serving_tokens._MaterializedRefCache(max_bytes=1 << 20) + features = _mm_features(tmp_path) + + async def _run(): + first = await serving_tokens._decode_raw_mm_kwargs( + features, processor_model_name="model", trust_remote_code=False, cache=cache + ) + second = await serving_tokens._decode_raw_mm_kwargs( + features, processor_model_name="model", trust_remote_code=False, cache=cache + ) + return first, second + + first, second = asyncio.run(_run()) + assert len(calls) == 1 + assert first["image"][0] is second["image"][0] + assert (cache.hits, cache.misses) == (1, 1) + + +def test_mm_materialize_cache_does_not_alias_distinct_raw_refs(tmp_path, monkeypatch): + import pytest + + from prime_rl.inference.vllm import serving_tokens + + calls: list = [] + _patch_adapter(monkeypatch, calls) + cache = serving_tokens._MaterializedRefCache(max_bytes=1 << 20) + valid = _mm_features(tmp_path, image_seed=1) + forged = _mm_features(tmp_path, image_seed=2) + forged.mm_hashes["image"][0] = valid.mm_hashes["image"][0] + + async def _run(): + await serving_tokens._decode_raw_mm_kwargs( + valid, processor_model_name="model", trust_remote_code=False, cache=cache + ) + await serving_tokens._decode_raw_mm_kwargs( + forged, processor_model_name="model", trust_remote_code=False, cache=cache + ) + + with pytest.raises(serving_tokens._MMImageRefError, match="Expected image hash"): + asyncio.run(_run()) + assert len(calls) == 1 + assert (cache.hits, cache.misses) == (0, 2) + + +def test_mm_materialize_cache_byte_budget_evicts_oldest(tmp_path, monkeypatch): + from vllm.multimodal.cache import MultiModalCache + + from prime_rl.inference.vllm import serving_tokens + + calls: list = [] + _patch_adapter(monkeypatch, calls) + monkeypatch.setattr(MultiModalCache, "get_item_size", classmethod(lambda _cls, _item: 60)) + cache = serving_tokens._MaterializedRefCache(max_bytes=100) + features_a = _mm_features(tmp_path, image_seed=1) + features_b = _mm_features(tmp_path, image_seed=2) + + async def _decode(features): + return await serving_tokens._decode_raw_mm_kwargs( + features, processor_model_name="model", trust_remote_code=False, cache=cache + ) + + asyncio.run(_decode(features_a)) + asyncio.run(_decode(features_b)) # 60 + 60 > 100: evicts a + assert cache.evictions == 1 + asyncio.run(_decode(features_a)) # miss again: re-materializes + assert len(calls) == 3 + assert cache.misses == 3 + + +def test_mm_materialize_cache_failure_not_cached(tmp_path, monkeypatch): + import pytest + + from prime_rl.inference.vllm import serving_tokens + + calls: list = [] + fail_first = {"remaining": 1} + + def _materialize(item): + if fail_first["remaining"]: + fail_first["remaining"] -= 1 + raise serving_tokens._MMImageRefError("transient failure") + return {"materialized": True} + + _patch_adapter(monkeypatch, calls, materialize=_materialize) + cache = serving_tokens._MaterializedRefCache(max_bytes=1 << 20) + features = _mm_features(tmp_path) + + async def _decode(): + return await serving_tokens._decode_raw_mm_kwargs( + features, processor_model_name="model", trust_remote_code=False, cache=cache + ) + + with pytest.raises(serving_tokens._MMImageRefError): + asyncio.run(_decode()) + out = asyncio.run(_decode()) + assert out["image"][0] == {"materialized": True} + assert len(calls) == 2 + + +def test_mm_materialize_cache_single_flight(tmp_path, monkeypatch): + from prime_rl.inference.vllm import serving_tokens + + calls: list = [] + _patch_adapter(monkeypatch, calls) + cache = serving_tokens._MaterializedRefCache(max_bytes=1 << 20) + features = _mm_features(tmp_path) + + async def _run(): + return await asyncio.gather( + serving_tokens._decode_raw_mm_kwargs( + features, processor_model_name="model", trust_remote_code=False, cache=cache + ), + serving_tokens._decode_raw_mm_kwargs( + features, processor_model_name="model", trust_remote_code=False, cache=cache + ), + ) + + first, second = asyncio.run(_run()) + assert len(calls) == 1 + assert first["image"][0] is second["image"][0] diff --git a/tests/unit/orchestrator/test_batch.py b/tests/unit/orchestrator/test_batch.py index 8b4b31f7ed..c85115d0eb 100644 --- a/tests/unit/orchestrator/test_batch.py +++ b/tests/unit/orchestrator/test_batch.py @@ -5,7 +5,7 @@ from prime_rl.trainer.batch import pad_micro_batch, prepare_batch, prepare_sample from prime_rl.trainer.utils import build_bin_cost -from prime_rl.transport.types import EncodedTensor, MicroBatch, RoutedExperts, TrainingSample +from prime_rl.transport.types import MicroBatch, MMImageRef, MMRefs, RoutedExperts, TrainingSample def _routed_experts(data, dtype=np.uint8): @@ -17,11 +17,6 @@ def _routed_experts(data, dtype=np.uint8): ) -def _encoded(arr) -> EncodedTensor: - a = np.asarray(arr) - return EncodedTensor(data=a.tobytes(), shape=list(a.shape), dtype=str(a.dtype)) - - @pytest.fixture def make_training_example(): def _make_training_example( @@ -416,52 +411,67 @@ def test_prepare_sample_truncates_routed_experts(): assert micro_batch.env_names == ["test-env"] * 3 -def test_prepare_sample_truncates_mm_at_image_boundary(): - """Truncation never splits an image's placeholder block: it cuts to a whole-image boundary - and slices mm_kwargs to match, so image-token count stays == image-embedding count.""" - # Two 2-token images (patches-per-token = 1): image-pad at indices 1,2 (img0) and 4,5 (img1). - mm_token_type_ids = [0, 1, 1, 0, 1, 1, 0] - pixel_values = np.array([[1.0], [1.0], [2.0], [2.0]], dtype=np.float32) # img0=1.0, img1=2.0 - grid = np.array([[1, 2, 1], [1, 2, 1]], dtype=np.int64) +def _image_ref(uri: str, offset: int, length: int) -> MMImageRef: + return MMImageRef( + item={ + "kind": "prime_raw_mm_item", + "family": "qwen_vl", + "raw_image_uri": uri, + "payload": {"image_grid_thw": [[1, 1, 1]]}, + }, + hash="a" * 32, + uri=uri, + offset=offset, + length=length, + ) + + +def test_prepare_sample_truncates_raw_mm_refs_at_image_boundary(): + first_image = _image_ref("file:///tmp/image-0.png", offset=1, length=2) + second_image = _image_ref("file:///tmp/image-1.png", offset=4, length=2) sample = TrainingSample( token_ids=[10, 11, 12, 13, 14, 15, 16], - mask=[False, False, False, False, False, True, True], + mask=[False, False, True, True, False, True, True], logprobs=[0.0] * 7, temperatures=[1.0] * 7, - advantages=[0.0] * 6 + [1.0], + advantages=[0.0, 0.0, 1.0, 1.0, 0.0, 1.0, 1.0], env_name="test-env", - mm_token_type_ids=mm_token_type_ids, - mm_kwargs={"pixel_values": _encoded(pixel_values), "image_grid_thw": _encoded(grid)}, + mm_token_type_ids=[0, 1, 1, 0, 1, 1, 0], + mm_refs=MMRefs(images=[first_image, second_image]), ) - # seq_len=5 falls inside img1 (one of its two placeholders survives) -> drop img1 entirely. - mb = prepare_sample(sample, seq_len=5) - assert len(mb.input_ids) == 4 # cut back to img1's first placeholder (index 4) - assert len(mb.mm_token_type_ids) == len(mb.input_ids) - n_placeholders = sum(1 for t in mb.mm_token_type_ids if t) - assert n_placeholders == 2 # only img0's two placeholders remain - # No mismatch: placeholders == image embeddings, and only img0's pixels are kept. - assert mb.mm_kwargs["pixel_values"].shape == [2, 1] - assert mb.mm_kwargs["image_grid_thw"].shape == [1, 3] - kept = np.frombuffer(bytearray(mb.mm_kwargs["pixel_values"].data), dtype=np.float32) - assert kept.tolist() == [1.0, 1.0] - assert n_placeholders == mb.mm_kwargs["pixel_values"].shape[0] # ppt == 1 here - - -def test_prepare_batch_packs_multimodal_with_text(): - mm_sample = TrainingSample( - token_ids=[10, 11, 12], - mask=[False, True, True], - logprobs=[0.0, -0.1, -0.2], - temperatures=[1.0, 1.0, 1.0], - advantages=[0.0, 1.0, 1.0], - env_name="mm-env", - mm_token_type_ids=[0, 1, 0], - mm_kwargs={ - "pixel_values": _encoded(np.array([[1.0, 2.0]], dtype=np.float32)), - "image_grid_thw": _encoded(np.array([[1, 2, 2]], dtype=np.int64)), - }, - ) + # seq_len=5 splits the second image: cut at its start, keep refs for the first. + micro_batch = prepare_sample(sample, seq_len=5) + assert micro_batch.input_ids == [10, 11, 12, 13] + assert micro_batch.mm_token_type_ids == [0, 1, 1, 0] + assert micro_batch.mm_refs == MMRefs(images=[first_image]) + + # seq_len=2 splits the first image: no image survives. + micro_batch = prepare_sample(sample, seq_len=2) + assert micro_batch.input_ids == [10] + assert micro_batch.mm_token_type_ids == [0] + assert micro_batch.mm_refs is None + + +def test_prepare_batch_packs_raw_mm_samples_with_rebased_offsets(): + """Same-family raw-ref samples pack with each other and with text from the same + run/LoRA; merged image ref offsets are rebased to the packed token stream. A + different adapter family never joins the bin.""" + + def mm_sample(uri: str, family: str = "qwen_vl") -> TrainingSample: + ref = _image_ref(uri, offset=1, length=1) + ref.item["family"] = family + return TrainingSample( + token_ids=[10, 11, 12], + mask=[False, True, True], + logprobs=[0.0, -0.1, -0.2], + temperatures=[1.0, 1.0, 1.0], + advantages=[0.0, 1.0, 1.0], + env_name="mm-env", + mm_token_type_ids=[0, 1, 0], + mm_refs=MMRefs(images=[ref]), + ) + text_sample = TrainingSample( token_ids=[20, 21], mask=[False, True], @@ -472,25 +482,33 @@ def test_prepare_batch_packs_multimodal_with_text(): ) batches_per_gpu = prepare_batch( - rollouts=[mm_sample, text_sample], - seq_len=8, + rollouts=[ + mm_sample("file:///tmp/image-0.png"), + mm_sample("file:///tmp/image-1.png"), + mm_sample("file:///tmp/image-2.png", family="kimi_k25"), + text_sample, + ], + seq_len=16, num_train_workers=2, - idxs=[0, 0], + idxs=[0, 0, 0, 0], num_loras=1, bin_cost=build_bin_cost(None), ) real_batches = [batch for batch in _flatten_batches(batches_per_gpu) if _has_loss_tokens(batch)] - assert len(real_batches) == 1 - batch = real_batches[0] - assert batch.seq_lens == [3, 2] - assert batch.sequence_lengths == [3, 2] - assert batch.position_ids == [0, 1, 2, 0, 1] - assert batch.mm_token_type_ids == [0, 1, 0, 0, 0] - assert batch.mm_kwargs is not None - assert batch.mm_kwargs["pixel_values"].shape == [1, 2] - assert batch.mm_kwargs["image_grid_thw"].shape == [1, 3] - assert batch.env_names == ["mm-env"] * 3 + ["text-env"] * 2 + mm_batches = [batch for batch in real_batches if batch.mm_refs is not None] + assert len(mm_batches) == 2 + packed = max(mm_batches, key=lambda batch: len(batch.mm_refs.images)) + assert len(packed.mm_refs.images) == 2 + assert packed.seq_lens[:2] == [3, 3] + # Sample-relative offset 1, rebased by each sample's start in the packed stream. + assert [image.offset for image in packed.mm_refs.images] == [1, 4] + assert packed.mm_token_type_ids[:6] == [0, 1, 0, 0, 1, 0] + assert packed.env_names[:6] == ["mm-env"] * 6 + # The kimi-family sample stayed out of the qwen bin. + (other,) = [batch for batch in mm_batches if batch is not packed] + assert len(other.mm_refs.images) == 1 + assert other.mm_refs.images[0].item["family"] == "kimi_k25" def test_prepare_sample_none_routed_experts(): diff --git a/tests/unit/orchestrator/test_qwen3_vl_e2e.py b/tests/unit/orchestrator/test_qwen3_vl_e2e.py deleted file mode 100644 index d125e8a469..0000000000 --- a/tests/unit/orchestrator/test_qwen3_vl_e2e.py +++ /dev/null @@ -1,190 +0,0 @@ -"""End-to-end integration test for the Qwen3-VL renderer path. - -Walks a multimodal request through the full client stack — RendererClient -→ renderers.client.generate → /inference/v1/generate features payload — -with the HTTP layer mocked, and verifies that vLLM can deserialize the -features back into engine inputs identical to what its own server-side -processor would have produced for the same messages. - -This is the strongest end-to-end check we can run without a GPU. The -remaining missing piece (vLLM actually consuming the engine input, -sampling tokens, and returning them) is exercised in real rollouts. -""" - -from __future__ import annotations - -import asyncio -import json -from pathlib import Path -from typing import Any -from unittest.mock import MagicMock - -import httpx -import pytest - -_HF_CACHE = Path("~/.cache/huggingface/hub").expanduser() -_MODEL = "Qwen/Qwen3-VL-4B-Instruct" - - -def _model_cached() -> bool: - safe = "models--" + _MODEL.replace("/", "--") - snapshots = _HF_CACHE / safe / "snapshots" - if not snapshots.is_dir(): - return False - return any(p.is_dir() for p in snapshots.iterdir()) - - -pytestmark = pytest.mark.skipif( - not _model_cached(), - reason=f"{_MODEL}: HF snapshot not cached locally", -) - - -class _FakeOpenAI: - """Minimal AsyncOpenAI stand-in that captures POST bodies. - - The renderer client calls ``client.post(absolute_url, body=...)``; - we capture the body for assertions and return a canned generate - response so the parse-side of the flow runs. - """ - - def __init__(self): - self.calls: list[dict[str, Any]] = [] - self.base_url = "http://fake-host:8000/v1" - - async def post(self, path, *, cast_to=dict, body=None, options=None): - self.calls.append({"path": path, "body": body, "options": options}) - # Reply with two sampled tokens + <|im_end|>. The renderer's - # parse_response slices the content tokens. - payload = { - "request_id": "qwen-vl-e2e", - "choices": [ - { - "index": 0, - "token_ids": [50, 60, 151645], - "logprobs": { - "content": [ - {"token": "t1", "logprob": -0.1}, - {"token": "t2", "logprob": -0.2}, - {"token": "t3", "logprob": -0.3}, - ] - }, - "finish_reason": "stop", - }, - ], - } - return httpx.Response(200, content=json.dumps(payload).encode()) - - -def test_renderer_client_qwen3_vl_e2e_features_payload_roundtrips_through_vllm(): - """Walk a Qwen3-VL multimodal turn through the renderer client and - verify the resulting ``/inference/v1/generate`` body has a valid - ``features`` payload that: - - 1. parses through vLLM's ``GenerateRequest`` pydantic model, - 2. decodes back to ``MultiModalKwargsItem`` instances carrying - ``pixel_values`` + ``image_grid_thw`` of the right shapes, - 3. has placeholder ranges that exactly cover the ``<|image_pad|>`` - runs in the prompt token sequence. - """ - from PIL import Image - from renderers.base import load_tokenizer - from renderers.qwen3_vl import Qwen3VLRenderer - from transformers import AutoProcessor - from verifiers.clients.renderer_client import RendererClient - from verifiers.types import ( - ClientConfig, - UserMessage, - ) - from vllm.entrypoints.scale_out.token_in_token_out.mm_serde import decode_mm_kwargs_item - from vllm.entrypoints.scale_out.token_in_token_out.protocol import GenerateRequest - - # ── Build a real Qwen3VLRenderer with a real processor. ───────────── - tokenizer = load_tokenizer(_MODEL) - processor = AutoProcessor.from_pretrained(_MODEL) - renderer = Qwen3VLRenderer(tokenizer, processor=processor) - - image_pad_id = tokenizer.convert_tokens_to_ids("<|image_pad|>") - - # ── Manually wire a RendererClient bypassing the pool factory. ────── - client_cfg = ClientConfig(client_type="renderer", base_url="http://fake-host:8000/v1") - rc = object.__new__(RendererClient) - rc._config = client_cfg - rc._renderer = renderer - rc._pool_size = 1 - rc._client = _FakeOpenAI() - rc.logger = MagicMock() - - # ── Build a verifiers-shaped user message with an image. ──────────── - img = Image.new("RGB", (224, 224), color=(64, 128, 255)) - # The renderer accepts the OpenAI ``image_url`` content-part shape — - # the same shape verifiers' UserMessage carries through. - user = UserMessage( - content=[ - {"type": "text", "text": "What's in this picture?"}, - # Embed the PIL image directly. The verifiers→renderer message - # converter forwards content unchanged for our purposes. - {"type": "image", "image": img}, - ] - ) - - # to_native_prompt converts to renderer-shaped messages. - prompt, _ = asyncio.run(rc.to_native_prompt([user])) - sampling = {"max_tokens": 16} - - response = asyncio.run( - rc.get_native_response( - prompt=prompt, - model=_MODEL, - sampling_args=sampling, - tools=None, - ) - ) - - # ── The HTTP body should carry a features payload. ────────────────── - fake = rc.client - assert isinstance(fake, _FakeOpenAI) - assert len(fake.calls) == 1 - body = fake.calls[0]["body"] - assert "features" in body, "RendererClient should ship features for image content" - features = body["features"] - - # ── Pydantic-roundtrip through vLLM's GenerateRequest model. ──────── - gen_req = GenerateRequest( - token_ids=body["token_ids"], - features=features, - sampling_params=body["sampling_params"], - ) - assert gen_req.features is not None - assert "image" in gen_req.features.mm_hashes - assert len(gen_req.features.mm_hashes["image"]) == 1 - - # ── Placeholder anchoring: the offset/length in features must land - # exactly on a run of <|image_pad|> ids in the prompt. ─────────── - placeholders = gen_req.features.mm_placeholders["image"] - assert len(placeholders) == 1 - ph = placeholders[0] - pad_slice = body["token_ids"][ph.offset : ph.offset + ph.length] - assert all(t == image_pad_id for t in pad_slice), ( - f"placeholder span ({ph.offset}, {ph.length}) does not cover image_pad tokens; slice={pad_slice[:8]}..." - ) - - # ── kwargs_data decodes to MultiModalKwargsItem with the right keys. ─ - assert gen_req.features.kwargs_data is not None - encoded_items = gen_req.features.kwargs_data["image"] - assert len(encoded_items) == 1 - item = decode_mm_kwargs_item(encoded_items[0]) - assert set(item.keys()) == {"pixel_values", "image_grid_thw"} - - # The image_grid_thw must match what the HF processor would have - # produced for the same PIL image — strongest signal that the engine - # sees the same image features the trainer will. - direct_proc_out = processor.image_processor(images=[img], return_tensors="pt") - expected_grid = direct_proc_out["image_grid_thw"][0].tolist() - assert item["image_grid_thw"].data.tolist() == expected_grid - - # ── Response parsed through renderer's parse_response. ────────────── - assert response["completion_ids"] == [50, 60, 151645] - # multi_modal_data surfaces on the result so the caller can persist it. - assert response["multi_modal_data"] is not None - assert len(response["multi_modal_data"].mm_items["image"]) == 1 diff --git a/tests/unit/train/rl/test_packer.py b/tests/unit/train/rl/test_packer.py index 93a1f04355..09a59a598a 100644 --- a/tests/unit/train/rl/test_packer.py +++ b/tests/unit/train/rl/test_packer.py @@ -1,7 +1,6 @@ from pathlib import Path from typing import Generator -import numpy as np import pytest import tomli_w import torch @@ -13,7 +12,7 @@ from prime_rl.trainer.runs import setup_multi_run_manager from prime_rl.trainer.utils import build_bin_cost from prime_rl.trainer.world import reset_world -from prime_rl.transport.types import EncodedTensor, TrainingSample +from prime_rl.transport.types import MMImageRef, MMRefs, TrainingSample @pytest.fixture(autouse=True, scope="module") @@ -52,16 +51,8 @@ def make_training_sample() -> TrainingSample: ) -def _encoded_tensor(data, dtype) -> EncodedTensor: - arr = np.asarray(data, dtype=dtype) - return EncodedTensor(dtype=str(arr.dtype), shape=list(arr.shape), data=arr.tobytes()) - - -def _decode_encoded_tensor(encoded: EncodedTensor): - return np.frombuffer(encoded.data, dtype=np.dtype(encoded.dtype)).reshape(encoded.shape).tolist() - - def _mm_sample(value: float, env_name: str = "test-env") -> TrainingSample: + uri = f"file:///tmp/image-{value}.png" return TrainingSample( token_ids=[1, 250, 2], mask=[False, True, True], @@ -70,10 +61,22 @@ def _mm_sample(value: float, env_name: str = "test-env") -> TrainingSample: env_name=env_name, advantages=[0.0, 1.0, 1.0], mm_token_type_ids=[0, 1, 0], - mm_kwargs={ - "pixel_values": _encoded_tensor([[value, value + 1]], np.float32), - "image_grid_thw": _encoded_tensor([[1, 2, 2]], np.int64), - }, + mm_refs=MMRefs( + images=[ + MMImageRef( + item={ + "kind": "prime_raw_mm_item", + "family": "qwen_vl", + "raw_image_uri": uri, + "payload": {"image_grid_thw": [[1, 1, 1]]}, + }, + hash="a" * 32, + uri=uri, + offset=1, + length=1, + ) + ] + ), ) @@ -176,7 +179,7 @@ def fake_sender(_output_dir, _data_world_size, _current_step, _config): assert micro_batch.run_step == 1 -def test_multipacker_pack_preserves_mm_kwargs_modality_and_run_tagging(tmp_path, monkeypatch): +def test_multipacker_pack_preserves_mm_modality_alignment_and_run_tagging(tmp_path, monkeypatch): """MultiPacker keeps multimodal and text microbatches aligned across ranks.""" from prime_rl.trainer.batch import _is_multimodal_sample @@ -204,7 +207,7 @@ def test_multipacker_pack_preserves_mm_kwargs_modality_and_run_tagging(tmp_path, assert mm_mbs, "no MM microbatches produced" real_run_idxs = set() for mb in mm_mbs: - assert mb.mm_kwargs is not None + assert mb.mm_refs is not None if any(mb.loss_mask): tagged = [i for i, n in enumerate(mb.lora_num_tokens) if n > 0] assert len(tagged) == 1 and mb.lora_num_tokens[tagged[0]] == len(mb.input_ids) @@ -212,8 +215,8 @@ def test_multipacker_pack_preserves_mm_kwargs_modality_and_run_tagging(tmp_path, assert real_run_idxs == {a, b}, f"both runs' MM should be tagged; got {real_run_idxs}" -def test_multipacker_packs_mm_kwargs_within_each_run(tmp_path, monkeypatch): - """Compatible eager multimodal samples pack within a run but never across runs.""" +def test_multipacker_packs_raw_mm_samples_within_each_run(tmp_path, monkeypatch): + """Same-family raw-ref samples pack within a run (offsets rebased) but never across runs.""" from prime_rl.trainer.batch import _is_multimodal_sample manager, packer, sent = _packer_with_two_runs(tmp_path, monkeypatch, dp_world_size=1, seq_len=12) @@ -230,11 +233,8 @@ def test_multipacker_packs_mm_kwargs_within_each_run(tmp_path, monkeypatch): assert len(real_mm_mbs) == 2 for mb in real_mm_mbs: assert len(mb.input_ids) == 6 - assert mb.position_ids == [0, 1, 2, 0, 1, 2] assert mb.seq_lens == [3, 3] - assert mb.mm_kwargs is not None - assert mb.mm_kwargs["pixel_values"].shape == [2, 2] - assert mb.mm_kwargs["image_grid_thw"].shape == [2, 3] - assert len(_decode_encoded_tensor(mb.mm_kwargs["pixel_values"])) == 2 + assert mb.mm_refs is not None and len(mb.mm_refs.images) == 2 + assert [image.offset for image in mb.mm_refs.images] == [1, 4] tagged = [i for i, n in enumerate(mb.lora_num_tokens) if n > 0] - assert len(tagged) == 1 + assert len(tagged) == 1 and mb.lora_num_tokens[tagged[0]] == 6 diff --git a/tests/unit/train/test_model_forward.py b/tests/unit/train/test_model_forward.py index 14fd521d5e..8071809655 100644 --- a/tests/unit/train/test_model_forward.py +++ b/tests/unit/train/test_model_forward.py @@ -1,8 +1,10 @@ from types import SimpleNamespace +import pytest import torch import torch.nn as nn +from prime_rl.multimodal.adapters.base import ForwardPolicy from prime_rl.trainer.model import forward @@ -35,11 +37,11 @@ def test_forward_passes_renderer_mm_token_type_ids_through(): position_ids, seq_lens=torch.tensor([input_ids.shape[1]]), mm_kwargs={"pixel_values": pixel_values, "image_grid_thw": image_grid_thw}, + mm_forward_policy=ForwardPolicy(pass_position_ids_with_mm=False), mm_token_type_ids=mm_token_type_ids, ) assert model.kwargs is not None - # MRoPE families (image_grid_thw present) get position_ids stripped. assert "position_ids" not in model.kwargs torch.testing.assert_close(model.kwargs["pixel_values"], pixel_values) torch.testing.assert_close(model.kwargs["image_grid_thw"], image_grid_thw) @@ -60,6 +62,7 @@ def test_forward_omits_mm_token_type_ids_when_renderer_does_not_supply(): position_ids, seq_lens=torch.tensor([input_ids.shape[1]]), mm_kwargs={"pixel_values": torch.ones(2, 3), "image_grid_thw": torch.tensor([[1, 1, 2]])}, + mm_forward_policy=ForwardPolicy(pass_position_ids_with_mm=False), ) assert model.kwargs is not None @@ -68,8 +71,7 @@ def test_forward_omits_mm_token_type_ids_when_renderer_does_not_supply(): def test_forward_keeps_position_ids_for_non_mrope_vlm(): - """Non-MRoPE VLM families (no ``image_grid_thw``) keep the trainer's - pre-computed ``position_ids``.""" + """Families whose adapter asks for position_ids keep the trainer's values.""" model = _CaptureModel(SimpleNamespace(model_type="gemma3")) input_ids = torch.tensor([[1, 10, 10, 2]]) position_ids = torch.arange(input_ids.shape[1]).unsqueeze(0) @@ -80,7 +82,24 @@ def test_forward_keeps_position_ids_for_non_mrope_vlm(): position_ids, seq_lens=torch.tensor([input_ids.shape[1]]), mm_kwargs={"pixel_values": torch.ones(2, 3)}, + mm_forward_policy=ForwardPolicy(pass_position_ids_with_mm=True), ) assert model.kwargs is not None torch.testing.assert_close(model.kwargs["position_ids"], position_ids) + + +def test_forward_policy_can_require_mm_token_type_ids(): + model = _CaptureModel(SimpleNamespace(model_type="qwen3_vl")) + input_ids = torch.tensor([[1, 10, 10, 2]]) + position_ids = torch.arange(input_ids.shape[1]).unsqueeze(0) + + with pytest.raises(ValueError, match="mm_token_type_ids"): + forward( + model, + input_ids, + position_ids, + seq_lens=torch.tensor([input_ids.shape[1]]), + mm_kwargs={"pixel_values": torch.ones(2, 3)}, + mm_forward_policy=ForwardPolicy(requires_mm_token_type_ids=True), + ) diff --git a/tests/unit/utils/test_mm.py b/tests/unit/utils/test_mm.py new file mode 100644 index 0000000000..e891435693 --- /dev/null +++ b/tests/unit/utils/test_mm.py @@ -0,0 +1,73 @@ +from types import SimpleNamespace + +import pytest + +from prime_rl.multimodal.schema import RAW_MM_ITEM_KIND +from prime_rl.trainer.rl.data import DataLoader +from prime_rl.transport.types import MicroBatch, MMImageRef, MMRefs +from prime_rl.utils.mm import build_mm_refs + + +class _MissingMaterializer: + def materialize(self, refs): + raise FileNotFoundError("missing image") + + +def _qwen_item(grid, uri: str = "file:///tmp/missing-image.png"): + return { + "kind": RAW_MM_ITEM_KIND, + "family": "qwen_vl", + "raw_image_uri": uri, + "payload": {"image_grid_thw": grid}, + } + + +def _refs(uri: str = "file:///tmp/missing-image.png") -> MMRefs: + return MMRefs( + images=[MMImageRef(item=_qwen_item([[1, 1, 1]], uri), hash="a" * 32, uri=uri, offset=1, length=1)], + ) + + +def _loader() -> DataLoader: + loader = object.__new__(DataLoader) + loader.multi_run_manager = SimpleNamespace(max_runs=1) + loader.mm_materializer = _MissingMaterializer() + loader.last_mm_materialize_time = 0.0 + loader.last_mm_images_materialized = 0 + return loader + + +def _micro_batch() -> MicroBatch: + return MicroBatch( + input_ids=[10, 11, 12], + loss_mask=[False, True, True], + advantages=[1.5, 1.5, 1.5], + inference_logprobs=[0.0, -0.1, -0.2], + position_ids=[0, 1, 2], + sequence_lengths=[3], + seq_lens=[3], + temperatures=[1.0, 1.0, 1.0], + env_names=["env", "env", "env"], + lora_num_tokens=[3], + mm_refs=_refs(), + mm_token_type_ids=[0, 1, 0], + ) + + +def test_build_mm_refs_accepts_offloaded_file_uri(tmp_path): + image_path = tmp_path / "image.png" + image_path.write_bytes(b"image") + multi_modal_data = SimpleNamespace( + mm_items={"image": [_qwen_item([[1, 1, 1]], image_path.as_uri())]}, + mm_hashes={"image": ["a" * 32]}, + mm_placeholders={"image": [SimpleNamespace(offset=1, length=1)]}, + ) + + refs = build_mm_refs(multi_modal_data) + + assert refs == _refs(image_path.as_uri()) + + +def test_dataloader_raises_on_missing_raw_image(): + with pytest.raises(FileNotFoundError, match="missing image"): + _loader()._micro_batch_to_tensor(_micro_batch()) diff --git a/uv.lock b/uv.lock index e59c12c400..7403c9927e 100644 --- a/uv.lock +++ b/uv.lock @@ -5357,7 +5357,7 @@ requires-dist = [ [[package]] name = "prime-sandboxes" -version = "0.2.33" +version = "0.2.35" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "aiofiles" }, @@ -5367,9 +5367,9 @@ dependencies = [ { name = "pydantic" }, { name = "tenacity" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/15/a9/880685cefd503aa92d2b49da46b60ba807015f8839c344a83d8c5bc02f30/prime_sandboxes-0.2.33.tar.gz", hash = "sha256:4253a3c345ccf6c07b41a81790f8515763333f094d20bc405a257432461e81d8", size = 81228, upload-time = "2026-07-22T23:23:55.715Z" } +sdist = { url = "https://files.pythonhosted.org/packages/57/cc/287c0db192ddbe003bd5882d1aab5461cdb3912d0c92db7b39734a9e828e/prime_sandboxes-0.2.35.tar.gz", hash = "sha256:e716bfccdcb51aefd5e2d48fdfc3a79bea8c6825bce1871f49f095ced03efc94", size = 86226, upload-time = "2026-08-05T14:29:53.33Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/f8/1b/10b1044b8aaef561c32a34140f4a13c7310f73f92f9546337e2e8e6a5239/prime_sandboxes-0.2.33-py3-none-any.whl", hash = "sha256:92c9647557e76191ba926ac8ed01708de9ec4450142131b673cd5cf8d9d7ea56", size = 42773, upload-time = "2026-07-22T23:23:54.381Z" }, + { url = "https://files.pythonhosted.org/packages/25/07/c186395925f780b1a18b350f562ebc4e8a64a182c37bd1e3db3948f297e2/prime_sandboxes-0.2.35-py3-none-any.whl", hash = "sha256:29d7057532b21545be5938d3aee822ae3a54085a45934d5c3ac55190e06c5bbb", size = 47189, upload-time = "2026-08-05T14:29:52.142Z" }, ] [[package]] @@ -6187,10 +6187,12 @@ requires-dist = [ { name = "numpy" }, { name = "openai", specifier = ">=1.108.1" }, { name = "openai-harmony", specifier = ">=0.0.4" }, + { name = "pillow", marker = "extra == 'vision'", specifier = ">=12.2.0" }, { name = "prime-pydantic-config", specifier = ">=0.3.0.dev83" }, { name = "tiktoken" }, { name = "transformers", specifier = ">=4.50.0" }, ] +provides-extras = ["vision"] [package.metadata.requires-dev] dev = [ @@ -6198,7 +6200,7 @@ dev = [ { name = "pre-commit" }, { name = "pytest", specifier = ">=7.0.0" }, { name = "pytest-asyncio", specifier = ">=0.21.0" }, - { name = "ruff" }, + { name = "ruff", specifier = "==0.15.12" }, { name = "torch", specifier = ">=2.11.0" }, { name = "torchvision", specifier = ">=0.26.0" }, { name = "ty", specifier = ">=0.0.1a29,<0.0.22" }, @@ -7918,6 +7920,7 @@ requires-dist = [ { name = "mcp", specifier = ">=1.24.0,<2" }, { name = "modal", marker = "extra == 'modal'", specifier = ">=1.4.0" }, { name = "msgpack", specifier = ">=1.1.2" }, + { name = "nemo-gym", marker = "python_full_version >= '3.12' and extra == 'nemo-gym'", specifier = "==0.4.0" }, { name = "nest-asyncio", marker = "extra == 'notebook'", specifier = ">=1.6.0" }, { name = "nltk", marker = "extra == 'ta'", specifier = ">=3.9.2" }, { name = "numpy", specifier = ">=2.1.0" }, @@ -7925,7 +7928,7 @@ requires-dist = [ { name = "openai-agents", specifier = ">=0.8.2" }, { name = "openenv", marker = "extra == 'openenv'", specifier = ">=0.4.1" }, { name = "prime-pydantic-config", extras = ["toml"], specifier = ">=0.4.2" }, - { name = "prime-sandboxes", specifier = ">=0.2.33" }, + { name = "prime-sandboxes", specifier = ">=0.2.35" }, { name = "prime-tunnel", specifier = ">=0.1.8" }, { name = "pydantic", specifier = ">=2.12.3" }, { name = "python-dotenv", marker = "extra == 'browser'", specifier = ">=1.0.0" }, @@ -7942,7 +7945,7 @@ requires-dist = [ { name = "typing-extensions", specifier = ">=4.12.2" }, { name = "uvloop", marker = "platform_python_implementation != 'PyPy' and sys_platform != 'cygwin' and sys_platform != 'win32'", specifier = ">=0.21.0" }, ] -provides-extras = ["browser", "harbor", "modal", "notebook", "openenv", "rg", "ta"] +provides-extras = ["browser", "harbor", "modal", "nemo-gym", "notebook", "openenv", "rg", "ta"] [package.metadata.requires-dev] dev = [