From 776ae71e05349972c3b7a23ecfaa2541c51c3ced Mon Sep 17 00:00:00 2001 From: pxc Date: Wed, 24 Jun 2026 20:08:21 +0800 Subject: [PATCH 01/84] openai api server record experience data --- trinity/common/config.py | 14 + trinity/common/models/allocator.py | 11 + trinity/common/models/vllm_model.py | 7 + trinity/common/models/vllm_patch/__init__.py | 58 ++- .../models/vllm_patch/recording/__init__.py | 43 +++ .../models/vllm_patch/recording/context.py | 38 ++ .../models/vllm_patch/recording/models.py | 141 +++++++ .../models/vllm_patch/recording/query.py | 94 +++++ .../models/vllm_patch/recording/recorder.py | 187 ++++++++++ .../models/vllm_patch/recording/server.py | 345 ++++++++++++++++++ .../models/vllm_patch/recording/store.py | 136 +++++++ trinity/explorer/proxy/app.py | 125 +------ trinity/explorer/proxy/service.py | 31 -- 13 files changed, 1077 insertions(+), 153 deletions(-) create mode 100644 trinity/common/models/vllm_patch/recording/__init__.py create mode 100644 trinity/common/models/vllm_patch/recording/context.py create mode 100644 trinity/common/models/vllm_patch/recording/models.py create mode 100644 trinity/common/models/vllm_patch/recording/query.py create mode 100644 trinity/common/models/vllm_patch/recording/recorder.py create mode 100644 trinity/common/models/vllm_patch/recording/server.py create mode 100644 trinity/common/models/vllm_patch/recording/store.py diff --git a/trinity/common/config.py b/trinity/common/config.py index 2d231490c55..1ade372c4ac 100644 --- a/trinity/common/config.py +++ b/trinity/common/config.py @@ -591,6 +591,20 @@ class InferenceModelConfig: node_rank: int = 0 enable_return_routed_experts: bool = False + # Turn on in-vLLM generation recording for the OpenAI API serving path: the + # engine wraps ``engine_client.generate`` and writes each finished turn as a + # Trinity ``Experience`` to the explorer proxy's shared SQL table + # (``proxy_history``), keyed by the request id (== OpenAI ``response.id``). + # When True, the Allocator forces ``enable_return_routed_experts`` and fills + # ``record_db_url`` from ``ExplorerConfig.db_url``. VLLMModel then mirrors + # the recording config onto the engine instance for the recorder to read. + # The capture width (top-k logprobs) reuses ``logprobs`` below (default 1). + enable_recording: bool = False + # SQL db url the recorder writes to (shared with the explorer proxy's + # HistoryRecorder). Populated by the Allocator from ExplorerConfig.db_url + # (with a cache_dir fallback) when enable_recording is on; None otherwise. + record_db_url: Optional[str] = None + # Buffer size (MB) for batched NCCL weight sync. Controls peak GPU memory during sync. weight_sync_buffer_size: int = 1024 # MB diff --git a/trinity/common/models/allocator.py b/trinity/common/models/allocator.py index 551ab09641d..0b51531c1ba 100644 --- a/trinity/common/models/allocator.py +++ b/trinity/common/models/allocator.py @@ -82,6 +82,17 @@ async def create_engine( config = deepcopy(config) config.engine_id = engine_id + # In-vLLM recording: force routed_experts capture (the engine reads + # this at build time — see vllm_model.py) and carry the db_url on the + # config so VLLMModel can mirror the recording config onto the engine + # instance for the recorder to read. No env / runtime_env involved. + if config.enable_recording: + config.enable_return_routed_experts = True + if not config.record_db_url: + config.record_db_url = self.config.db_url or ( + f"sqlite:///{self.config.buffer.cache_dir}/proxy_history.db" + ) + actor_bundle_lists = [] for node_id in range(config.nnodes): actor_name = self.get_actor_name(role, engine_id, node_id) diff --git a/trinity/common/models/vllm_model.py b/trinity/common/models/vllm_model.py index b30e83e8287..4b133c3ac64 100644 --- a/trinity/common/models/vllm_model.py +++ b/trinity/common/models/vllm_model.py @@ -186,6 +186,11 @@ async def prepare(self) -> None: engine_args.master_port = self.master_port if self.config.node_rank == 0: self.async_llm = vllm.AsyncLLMEngine.from_engine_args(engine_args) + # Expose the current checkpoint version on the engine instance so + # the in-vLLM recorder (which only sees `engine_client`) can + # attribute experiences to the right policy without an extra + # launch-time parameter. Updated in sync_model_weights. + self.async_llm.trinity_model_version = self.model_version await self._collective_rpc("apply_patches") await self.run_api_server() else: @@ -580,6 +585,7 @@ async def sync_model_weights( await self.async_llm.remove_lora(lora_id) await self.async_llm.add_lora(self.get_lora_request(self.default_lora_path)) self.model_version = model_version + self.async_llm.trinity_model_version = model_version return model_version from vllm.distributed.weight_transfer.base import WeightTransferUpdateRequest @@ -605,6 +611,7 @@ async def sync_model_weights( await self.async_llm.finish_weight_update() await self.async_llm.resume_generation() self.model_version = model_version + self.async_llm.trinity_model_version = model_version return model_version async def init_process_group( diff --git a/trinity/common/models/vllm_patch/__init__.py b/trinity/common/models/vllm_patch/__init__.py index 99670ecbce3..7e06a279bf4 100644 --- a/trinity/common/models/vllm_patch/__init__.py +++ b/trinity/common/models/vllm_patch/__init__.py @@ -106,7 +106,7 @@ def get_vllm_version(): return vllm_version -def _get_api_server_runner(vllm_version): +def _get_api_server_runner(vllm_version, *, recording: bool = False): if vllm_version == VLLM_VERSION_0120: from trinity.common.models.vllm_patch.api_patch_v12 import ( run_api_server_in_ray_actor_v12, @@ -122,6 +122,16 @@ def _get_api_server_runner(vllm_version): return run_api_server_in_ray_actor_v13 if VLLM_VERSION_0170 <= vllm_version: + # When generation recording is on, use the recording-enabled entry + # point (a superset of api_patch_v17 that wraps engine.generate and + # writes Experiences to the shared store). Otherwise stock api_patch_v17. + if recording: + from trinity.common.models.vllm_patch.recording import ( + run_api_server_with_recording as _recording_runner, + ) + + return _recording_runner + from trinity.common.models.vllm_patch.api_patch_v17 import ( run_api_server_in_ray_actor_v17, ) @@ -142,20 +152,38 @@ def get_api_server( logger: Logger, ): vllm_version = get_vllm_version() + # Recording is driven by the config field (not env, not an engine attr): + # get_api_server already receives InferenceModelConfig, so it builds the + # static RecordingConfig here and threads it into the recording runner. + recording = bool(config.enable_recording) and vllm_version >= VLLM_VERSION_0170 + if config.enable_recording and not recording: + logger.warning( + "enable_recording is on but vLLM %s < 0.17.0; recording disabled", + vllm.__version__, + ) - run_api_server_in_ray_actor = _get_api_server_runner(vllm_version) + run_api_server_in_ray_actor = _get_api_server_runner(vllm_version, recording=recording) logger.info(f"Using vLLM API patch for version {vllm.__version__}") - return asyncio.create_task( - run_api_server_in_ray_actor( - async_llm, - host=host, - port=port, - model_path=config.model_path, # type: ignore [arg-type] - logger=logger, - enable_auto_tool_choice=config.enable_auto_tool_choice, - tool_call_parser=config.tool_call_parser, - reasoning_parser=config.reasoning_parser, - enable_log_requests=config.enable_log_requests, - chat_template=config.chat_template, - ) + kwargs = dict( + host=host, + port=port, + model_path=config.model_path, # type: ignore [arg-type] + logger=logger, + enable_auto_tool_choice=config.enable_auto_tool_choice, + tool_call_parser=config.tool_call_parser, + reasoning_parser=config.reasoning_parser, + enable_log_requests=config.enable_log_requests, + chat_template=config.chat_template, ) + if recording: + from trinity.common.models.vllm_patch.recording import RecordingConfig + + kwargs["recording_config"] = RecordingConfig( + db_url=config.record_db_url, + table="proxy_history", + topk=config.logprobs if config.logprobs else 1, + ) + # The dynamic checkpoint version is read live off the engine instance + # (``async_llm.trinity_model_version``, mirrored by VLLMModel), so it is + # not part of the static config passed here. + return asyncio.create_task(run_api_server_in_ray_actor(async_llm, **kwargs)) diff --git a/trinity/common/models/vllm_patch/recording/__init__.py b/trinity/common/models/vllm_patch/recording/__init__.py new file mode 100644 index 00000000000..fa490baeb3f --- /dev/null +++ b/trinity/common/models/vllm_patch/recording/__init__.py @@ -0,0 +1,43 @@ +"""Generation-recording patch for the vLLM OpenAI server. + +Self-contained entry point that mirrors ``api_patch_v17.py``'s bootstrap flow +(``build_app`` -> ``init_app_state`` -> ``serve_http``) and additionally wires +in generation recording without touching vLLM source code. + +Designed for vllm >= 0.17.0. Drop-in alternative to +``trinity.common.models.vllm_patch.api_patch_v17``: point your launcher at +``recording.run_api_server_with_recording`` instead. +""" +from trinity.common.models.vllm_patch.recording.config import RecordingConfig +from trinity.common.models.vllm_patch.recording.context import ( + SessionMiddleware, + task_id_ctx, +) +from trinity.common.models.vllm_patch.recording.models import build_experience +from trinity.common.models.vllm_patch.recording.query import query_router +from trinity.common.models.vllm_patch.recording.recorder import ( + Recorder, + patch_engine_for_recording, +) +from trinity.common.models.vllm_patch.recording.server import ( + run_api_server_with_recording, +) +from trinity.common.models.vllm_patch.recording.store import ( + MemoryStore, + RecordStore, + SqlStore, +) + +__all__ = [ + "MemoryStore", + "RecordStore", + "RecordingConfig", + "SqlStore", + "Recorder", + "SessionMiddleware", + "build_experience", + "patch_engine_for_recording", + "query_router", + "run_api_server_with_recording", + "task_id_ctx", +] diff --git a/trinity/common/models/vllm_patch/recording/context.py b/trinity/common/models/vllm_patch/recording/context.py new file mode 100644 index 00000000000..15db3a332b0 --- /dev/null +++ b/trinity/common/models/vllm_patch/recording/context.py @@ -0,0 +1,38 @@ +"""Per-request task_id propagation. + +The OpenAI ``X-Session-ID`` request header carries the task id that scopes a +multi-turn task. We read it in an in-process ASGI middleware and stash it in a +contextvar so the engine-level wrapper (which runs in the same async task as +the serving handler) can recover it at record time. + +No ``X-Session-ID`` on a request is fine: the recorder falls back to +``request_id`` as the task id so nothing is silently dropped. +""" +from contextvars import ContextVar +from typing import Optional + +from starlette.middleware.base import BaseHTTPMiddleware +from starlette.requests import Request + +# The task id for the in-flight request, or None when the client did not send +# ``X-Session-ID`` (the recorder then uses request_id as a fallback task id). +task_id_ctx: ContextVar[Optional[str]] = ContextVar("trinity_recording_task_id", default=None) + +#: Canonical header name. Lower-cased per ASGI/httpx convention. +SESSION_ID_HEADER = "x-session-id" + + +class SessionMiddleware(BaseHTTPMiddleware): + """Capture ``X-Session-ID`` into ``task_id_ctx`` for the request's lifetime. + + Runs in-process (ASGI) — no extra network hop, no serialization cost beyond + a contextvar set/reset. + """ + + async def dispatch(self, request: Request, call_next): + task_id = request.headers.get(SESSION_ID_HEADER) + token = task_id_ctx.set(task_id) + try: + return await call_next(request) + finally: + task_id_ctx.reset(token) diff --git a/trinity/common/models/vllm_patch/recording/models.py b/trinity/common/models/vllm_patch/recording/models.py new file mode 100644 index 00000000000..d7e60b49140 --- /dev/null +++ b/trinity/common/models/vllm_patch/recording/models.py @@ -0,0 +1,141 @@ +# -*- coding: utf-8 -*- +"""Build a Trinity ``Experience`` from a finished vLLM ``RequestOutput``. + +We record into Trinity's native ``Experience`` struct (see +``trinity.common.experience``) rather than a bespoke record, so captured data +drops straight into Trinity's RL/buffer pipeline without a conversion step. + +Field mapping (captured ``RequestOutput`` fields -> ``Experience``): + request_id -> eid.suffix (``EID(suffix=...)``; this is the msg_id the + proxy/openai client sees as ``response.id`` — the key the + proxy's ``HistoryRecorder.update_reward`` looks up by). + ``eid.task``/``run``/``reward`` are left default here and + assigned later by the proxy at ``/feedback`` time, matching + ``explorer/proxy/service.record_experience`` semantics. + X-Session-ID -> info["task_id"] (traceability only; not used as a key) + prompt_token_ids -> tokens (prompt portion) + prompt_length + response_token_ids-> tokens (response portion) + logprobs -> Experience.logprobs -- but ONLY the *chosen* token's + logprob per position (flat ``[resp_length]`` tensor), per + the RL convention. vLLM's ``CompletionOutput.logprobs`` is a + top-k structure per position; we look up the actually-sampled + token id and take its ``.logprob``. + routed_experts -> Experience.routed_experts (uint8 tensor, verbatim) + model_version -> info["model_version"] (which checkpoint policy served the + turn; read in-actor by the recorder's provider) + +Plus bookkeeping (request_id / task_id / rank / timestamp / endpoint / +model_version) stashed in ``Experience.info`` so it round-trips with the +experience through serialize/deserialize. +""" +from typing import Any, Optional + +from trinity.common.experience import EID, Experience + + +def _extract_chosen_logprobs( + sample_logprobs: Any, + response_token_ids: list[int], +) -> Optional[list[float]]: + """Pull the sampled token's logprob at each response position. + + vLLM exposes ``CompletionOutput.logprobs`` as either a list of + ``dict[int, Logprob]`` or the ``FlatLogprobs`` container; both support + positional indexing returning ``dict[int, Logprob]`` for that position, so + we treat them uniformly. + + Returns a flat ``[resp_length]`` list of floats, or None when logprobs were + not requested/computed. + + Note: the sampled token is *always* present at each position's dict. vLLM + force-includes it as column 0 of the reported set + (``vllm/v1/worker/gpu/sample/logprob.py:compute_topk_logprobs``), so a + request with ``sampling_params.logprobs = N`` reports ``{sampled} ∪ + top-N`` — the chosen token is reported even when it ranks beyond N in the + model's distribution. There is therefore no "sampled token absent from + top-k" case to handle here: ``pos[tid]`` always resolves, and a length + mismatch between ``sample_logprobs`` and ``response_token_ids`` cannot + occur in normal operation (both are indexed per generated token). + + Args: + sample_logprobs: ``CompletionOutput.logprobs`` (may be None). + response_token_ids: The generated token ids. + + Returns: + Flat list of chosen-token logprobs, or None. + """ + if not sample_logprobs: + return None + # One entry per generated token; sampled token is force-included per the + # note above, so a direct lookup per position is always well-defined. + return [float(sample_logprobs[i][tid].logprob) for i, tid in enumerate(response_token_ids)] + + +def build_experience( + output: Any, + task_id: Optional[str], + *, + rank: int, + timestamp: str, + endpoint: str = "unknown", + model_version: Optional[int] = None, +) -> Optional[Experience]: + """Build a Trinity ``Experience`` from a finished ``RequestOutput``. + + Args: + output: A ``RequestOutput`` with ``finished == True``. + task_id: From ``X-Session-ID``; stored in ``info`` for traceability only. + Not used as the storage key — ``eid.suffix`` is, so a missing + ``X-Session-ID`` never drops a turn. + rank: Data-parallel serving rank. + timestamp: UTC ISO-8601 string (caller-stamped to keep this pure). + endpoint: Which OpenAI endpoint served the turn (best-effort). + model_version: Checkpoint version the serving policy was at; stamped + into ``info`` for RL attribution (read in-actor by the recorder). + + Returns: + A populated ``Experience``, or None if the turn is degenerate (no + prompt or no response tokens) and cannot form a valid experience. + """ + request_id = output.request_id + # Key by the request id (= the OpenAI response ``id`` / proxy msg_id) so the + # proxy's HistoryRecorder.update_reward can find this row at feedback time. + # task/run/reward are intentionally left default — the proxy assigns them. + + prompt_token_ids = list(output.prompt_token_ids or []) + + completion = output.outputs[0] if output.outputs else None + if completion is None: + return None + response_token_ids = list(completion.token_ids or []) + + # A valid single-turn experience needs both a prompt and a response; + # Experience.__init__ asserts len(tokens) > prompt_length otherwise. + if not prompt_token_ids or not response_token_ids: + return None + + tokens = prompt_token_ids + response_token_ids + prompt_length = len(prompt_token_ids) + + chosen_logprobs = _extract_chosen_logprobs(completion.logprobs, response_token_ids) + routed_experts = completion.routed_experts + + info = { + "request_id": request_id, + "task_id": task_id, + "rank": rank, + "timestamp": timestamp, + "endpoint": endpoint, + "model_version": model_version, + } + + return Experience( + eid=EID(suffix=request_id), + tokens=tokens, + logprobs=chosen_logprobs, + prompt_length=prompt_length, + routed_experts=routed_experts, + prompt_text=output.prompt, + response_text=completion.text, + info=info, + ) diff --git a/trinity/common/models/vllm_patch/recording/query.py b/trinity/common/models/vllm_patch/recording/query.py new file mode 100644 index 00000000000..7f53783b294 --- /dev/null +++ b/trinity/common/models/vllm_patch/recording/query.py @@ -0,0 +1,94 @@ +"""HTTP query endpoints over recorded experiences. + +Mounted on the main vLLM app (we own ``app`` in the bootstrap) via +``app.include_router(query_router)``. Routes are prefixed under ``/records`` to +avoid colliding with the OpenAI ``/v1/*`` surface — the OpenAI protocol is left +untouched. + +Listing endpoints return ``Experience.to_dict()`` (lightweight metadata, no +tensor payload). A single-turn endpoint returns the full experience as Trinity +serialized bytes (``Experience.serialize()``) so it can be fed straight back +into Trinity pipelines. +""" +from fastapi import APIRouter, HTTPException, Request, Response + +from trinity.common.experience import Experience +from trinity.common.models.vllm_patch.recording.store import RecordStore + +#: Mounted on app.state by the bootstrap wiring in ``server.py``. +_STORE_ATTR = "trinity_record_store" + +query_router = APIRouter(prefix="/records", tags=["trinity-recording"]) + + +def _store(request: Request, *, read: bool = False) -> RecordStore: + store = getattr(request.app.state, _STORE_ATTR, None) + if store is None: + raise HTTPException(status_code=503, detail="recording store not configured") + if read and not getattr(store, "supports_reads", True): + # SqlStore shares the explorer proxy's table; reads/consumes are served + # by the proxy's /feedback + /commit, not by these debug endpoints. + raise HTTPException( + status_code=503, + detail="this record store backend is read-only from the vLLM " + "side; query via the explorer proxy instead", + ) + return store + + +async def _get_exp(store: RecordStore, task_id: str, request_id: str) -> Experience: + exp = await store.get_turn(task_id, request_id) + if exp is None: + raise HTTPException(status_code=404, detail="experience not found") + return exp + + +@query_router.get("/tasks") +async def list_tasks(request: Request) -> dict: + """List all known task ids.""" + store = _store(request, read=True) + return {"task_ids": await store.list_tasks()} + + +@query_router.get("/tasks/{task_id}") +async def get_task(task_id: str, request: Request) -> dict: + """Return lightweight metadata for all experiences of a task. + + Tensor payloads (tokens/logprobs/routed_experts) are omitted here to keep + listing responses small; fetch the per-turn blob endpoint for full data. + """ + store = _store(request, read=True) + experiences = await store.get_task(task_id) + return {"task_id": task_id, "turns": [e.to_dict() for e in experiences]} + + +@query_router.get("/tasks/{task_id}/turns/{request_id}") +async def get_turn(task_id: str, request_id: str, request: Request) -> dict: + """Return lightweight metadata for a single experience.""" + store = _store(request, read=True) + exp = await _get_exp(store, task_id, request_id) + return exp.to_dict() + + +@query_router.get("/tasks/{task_id}/turns/{request_id}/blob") +async def get_turn_blob(task_id: str, request_id: str, request: Request) -> Response: + """Return the full experience as Trinity serialized bytes. + + Equivalent to ``Experience.serialize()``; deserializable via + ``Experience.deserialize(...)`` or ``deserialize_many``. + """ + store = _store(request, read=True) + exp = await _get_exp(store, task_id, request_id) + return Response( + content=Experience.serialize(exp), + media_type="application/octet-stream", + headers={"Content-Disposition": f'attachment; filename="{request_id}.bin"'}, + ) + + +@query_router.delete("/tasks/{task_id}") +async def delete_task(task_id: str, request: Request) -> dict: + """Delete all experiences for a task.""" + store = _store(request, read=True) + await store.delete_task(task_id) + return {"task_id": task_id, "deleted": True} diff --git a/trinity/common/models/vllm_patch/recording/recorder.py b/trinity/common/models/vllm_patch/recording/recorder.py new file mode 100644 index 00000000000..1db9b82f339 --- /dev/null +++ b/trinity/common/models/vllm_patch/recording/recorder.py @@ -0,0 +1,187 @@ +"""Engine-level wrap that forces top-k logprobs and records finished turns. + +This is the heart of the recording patch. It follows the same instance-level +wrap pattern as ``api_patch_v17.patch_vllm_reasoning_content_alias``: +``functools.wraps`` + a ``__patched_*__`` guard attribute to stay idempotent. + +Why wrap ``engine_client.generate`` instead of the serving layer? + * The serving layer (``OpenAIServingChat``/``OpenAIServingCompletion``) is + what decides streaming vs non-streaming and what fields to emit. vLLM does + NOT put ``routed_experts`` into streaming responses (the streaming choice + schemas omit it), so capturing at the HTTP layer misses it. + * ``RequestOutput`` / ``CompletionOutput`` carry the full data regardless of + streaming mode: ``prompt_token_ids``, ``token_ids``, ``logprobs``, + ``routed_experts`` (raw ndarray). Wrapping at the engine boundary captures + all four uniformly for chat / completion / responses endpoints. + * Forcing ``sampling_params.logprobs`` here only affects engine-internal + computation — the client response is unchanged unless the client itself + requested logprobs. Recording stays transparent. +""" +import asyncio +import functools +import logging +from typing import Optional + +from trinity.common.experience import Experience +from trinity.common.models.vllm_patch.recording.context import task_id_ctx +from trinity.common.models.vllm_patch.recording.models import build_experience +from trinity.common.models.vllm_patch.recording.store import RecordStore + +#: Guard attribute marking the wrapped generate, mirroring api_patch_v17 style. +_PATCHED_FLAG = "__patched_engine_recording__" +#: Instance attribute on the AsyncLLM engine_client holding the current serving +#: checkpoint version. Mirrored by ``VLLMModel.sync_model_weights`` (and at +#: engine creation); read live here so each experience is attributed to the +#: right policy without a launch-time parameter. +_MODEL_VERSION_ATTR = "trinity_model_version" + + +def patch_engine_for_recording( + engine_client, + recorder: "Recorder", + logger: logging.Logger, +) -> None: + """Wrap ``engine_client.generate`` in place to record finished turns. + + Instance-level: only this server's engine_client is affected, the global + class is untouched. Must run before ``init_app_state`` stores the engine + reference into the serving objects (they hold the same object, so the wrap + is inherited). + + Args: + engine_client: The AsyncLLM instance passed into the bootstrap. + recorder: The ``Recorder`` that will persist turns. + logger: Logger for the idempotency/confirmation message. + + Raises: + RuntimeError: If ``engine_client.generate`` is missing (unexpected + vLLM version drift). + """ + current = getattr(engine_client, "generate", None) + if current is None: + raise RuntimeError("vLLM patch failed: engine_client.generate not found") + if getattr(current, _PATCHED_FLAG, False): + return + + @functools.wraps(current) + async def _patched_generate(*args, **kwargs): + # generate(prompt, sampling_params, request_id, *, ...). + # ``engine_client.generate`` assigned as an instance attribute is NOT + # bound, so ``self`` is absent and args map 1:1 to the protocol. + sampling_params = kwargs.get("sampling_params") + if sampling_params is None and len(args) >= 2: + sampling_params = args[1] + + if recorder.enabled and sampling_params is not None: + desired = recorder.topk + cur = sampling_params.logprobs + sampling_params.logprobs = max(cur, desired) if cur is not None else desired + + last = None + # ``current`` is the original *bound* method captured pre-wrap, so it + # still resolves ``self`` correctly. Yields RequestOutput unchanged. + async for out in current(*args, **kwargs): + last = out + yield out + + if recorder.enabled and last is not None and getattr(last, "finished", False): + # Recover task id from the request's async context (set by + # SessionMiddleware). None when the client omitted X-Session-ID; + # the recorder then falls back to request_id. + task_id = task_id_ctx.get() + # Offload heavy serialization off the response critical path. + asyncio.create_task(recorder.record(last, task_id)) + + setattr(_patched_generate, _PATCHED_FLAG, True) + engine_client.generate = _patched_generate + logger.info("Patched vLLM engine_client.generate for generation recording") + + +class Recorder: + """Drains finished turns into a ``RecordStore`` from a background task. + + Putting records into an ``asyncio.Queue`` and flushing from a single worker + keeps the response path cheap (record == one ``queue.put``) and serializes + expensive payloads (ndarray -> .npy, json) off the serving hot loop. + """ + + def __init__( + self, + store: RecordStore, + *, + topk: int, + enabled: bool, + rank: int = 0, + engine_client=None, + ) -> None: + self.store = store + self.topk = topk + self.enabled = enabled + self.rank = rank + # The engine_client is the same AsyncLLM instance VLLMModel updates in + # sync_model_weights (``.trinity_model_version``), so we read the live + # checkpoint version off it at record time. + self.engine_client = engine_client + self._queue: "asyncio.Queue[Optional[Experience]]" = asyncio.Queue() + self._flusher: Optional[asyncio.Task] = None + + def start(self) -> None: + """Start the background flusher. Idempotent.""" + if self._flusher is not None or not self.enabled: + return + self._flusher = asyncio.create_task(self._flush_loop()) + + async def stop(self) -> None: + """Cancel the flusher and drain remaining queued turns.""" + if self._flusher is None: + return + self._flusher.cancel() + self._flusher = None + # Drain anything already queued so we don't lose in-flight turns. + while not self._queue.empty(): + exp = self._queue.get_nowait() + if exp is not None: + await self._safe_append(exp) + + async def record(self, output, task_id: Optional[str]) -> None: + """Enqueue a finished ``RequestOutput`` for recording as an Experience. + + Args: + output: A finished ``RequestOutput``. + task_id: From ``task_id_ctx``; stored in ``info`` for traceability. + """ + # Stamp now (real runtime, not a workflow sandbox): permitted here. + from datetime import datetime, timezone + + timestamp = datetime.now(timezone.utc).isoformat() + # Read the live checkpoint version the actor mirrors onto the engine. + model_version = getattr(self.engine_client, _MODEL_VERSION_ATTR, None) + exp = build_experience( + output, + task_id, + rank=self.rank, + timestamp=timestamp, + model_version=model_version, + ) + if exp is None: + # Degenerate turn (no prompt/response) — nothing to record. + return + await self._queue.put(exp) + + async def _flush_loop(self) -> None: + while True: + exp = await self._queue.get() + if exp is None: + # Sentinel for graceful shutdown. + return + await self._safe_append(exp) + + async def _safe_append(self, exp: Experience) -> None: + try: + await self.store.append_turn(exp) + except Exception: + # Never let a storage hiccup crash the flusher loop. + logging.getLogger(__name__).exception( + "recording store.append_turn failed for request %s", + exp.info.get("request_id"), + ) diff --git a/trinity/common/models/vllm_patch/recording/server.py b/trinity/common/models/vllm_patch/recording/server.py new file mode 100644 index 00000000000..c88774650e7 --- /dev/null +++ b/trinity/common/models/vllm_patch/recording/server.py @@ -0,0 +1,345 @@ +"""Self-contained bootstrap that copies api_patch_v17.py's server lifecycle +and additionally wires in generation recording. + +This module deliberately mirrors ``api_patch_v17.py`` so it can be used as a +drop-in alternative: point your launcher at +``trinity.common.models.vllm_patch.recording.run_api_server_with_recording`` +and you get the standard vLLM OpenAI server *plus* generation recording, with +no edits to vLLM source or to ``api_patch_v17.py``. + +Recording wiring (all applied between ``build_app`` and ``serve_http`` because +we own both ``app`` and ``engine_client`` at that point): + 1. ``patch_engine_for_recording`` — instance-level wrap of + ``engine_client.generate`` to force top-k logprobs and record finished + ``RequestOutput`` (covers chat/completion/responses, streaming and not). + 2. ``SessionMiddleware`` — in-process ASGI middleware reading + ``X-Session-ID`` into a contextvar. + 3. ``query_router`` — ``/records/*`` endpoints for later analysis. + +Only for vllm versions >= 0.17.0. +""" +import asyncio +import functools +import logging +from typing import Optional + +import vllm +import vllm.envs as envs +from packaging.version import parse as parse_version +from vllm.entrypoints.launcher import serve_http +from vllm.entrypoints.openai.api_server import ( + build_app, + create_server_socket, + create_server_unix_socket, + init_app_state, + validate_api_server_args, +) +from vllm.entrypoints.openai.cli_args import make_arg_parser +from vllm.reasoning import ReasoningParserManager +from vllm.tool_parsers import ToolParserManager +from vllm.utils.argparse_utils import FlexibleArgumentParser +from vllm.utils.network_utils import is_valid_ipv6_address +from vllm.utils.system_utils import set_ulimit +from vllm.version import __version__ as VLLM_VERSION + +from trinity.common.models.vllm_patch import get_vllm_version +from trinity.common.models.vllm_patch.recording.config import RecordingConfig +from trinity.common.models.vllm_patch.recording.context import SessionMiddleware +from trinity.common.models.vllm_patch.recording.query import query_router +from trinity.common.models.vllm_patch.recording.recorder import ( + Recorder, + patch_engine_for_recording, +) +from trinity.common.models.vllm_patch.recording.store import ( + MemoryStore, + RecordStore, + SqlStore, +) + +#: Attribute on app.state holding the active RecordStore. +_STORE_STATE_ATTR = "trinity_record_store" +#: Attribute on app.state holding the active Recorder. +_RECORDER_STATE_ATTR = "trinity_recorder" + + +def setup_server_in_ray(args, logger): + """Validate API server args, set up signal handler, create socket + ready to serve. + + Copied verbatim from api_patch_v17.py — identical lifecycle so the + recording entry point behaves like the stock Trinity server. + """ + + logger.info("vLLM API server version %s", VLLM_VERSION) + if get_vllm_version() < parse_version("0.23.0"): + from vllm.entrypoints.utils import log_non_default_args + + log_non_default_args(args) + else: + from vllm.entrypoints.serve.utils.api_utils import log_non_default_args + + log_non_default_args(args) + + if args.tool_parser_plugin and len(args.tool_parser_plugin) > 3: + ToolParserManager.import_tool_parser(args.tool_parser_plugin) + + if args.reasoning_parser_plugin and len(args.reasoning_parser_plugin) > 3: + ReasoningParserManager.import_reasoning_parser(args.reasoning_parser_plugin) + + validate_api_server_args(args) + + # workaround to make sure that we bind the port before the engine is set up. + # This avoids race conditions with ray. + # see https://github.com/vllm-project/vllm/issues/8204 + if args.uds: + sock = create_server_unix_socket(args.uds) + else: + sock_addr = (args.host or "", args.port) + sock = create_server_socket(sock_addr) + + # workaround to avoid footguns where uvicorn drops requests with too + # many concurrent requests active + set_ulimit() + + if args.uds: + listen_address = f"unix:{args.uds}" + else: + addr, port = sock_addr + is_ssl = args.ssl_keyfile and args.ssl_certfile + host_part = f"[{addr}]" if is_valid_ipv6_address(addr) else addr or "0.0.0.0" + listen_address = f"http{'s' if is_ssl else ''}://{host_part}:{port}" + return listen_address, sock + + +def dummy_add_signal_handler(self, *args, **kwargs): + # DO NOTHING HERE + pass + + +def _setup_recording( + args, + engine_client, + app, + logger, + recording_config: Optional[RecordingConfig] = None, +) -> Optional[Recorder]: + """Wire generation recording onto the in-construction server. + + Returns the started Recorder (for lifecycle management), or None if + recording is disabled (``recording_config`` is None). + + The static config (db_url/table/topk) arrives explicitly via + ``recording_config`` (built by ``get_api_server`` from + ``InferenceModelConfig``). The *dynamic* checkpoint version is read live + off ``engine_client.trinity_model_version`` (mirrored by VLLMModel at + engine creation and in ``sync_model_weights``). + + Args: + args: Parsed vLLM CLI args. + engine_client: AsyncLLM instance (we own it pre-init_app_state). + app: FastAPI app from ``build_app`` (we own it pre-serve_http). + logger: Logger. + recording_config: Static recording config; None disables recording. + """ + if recording_config is None: + return None + + if recording_config.db_url: + store: RecordStore = SqlStore( + db_url=recording_config.db_url, table_name=recording_config.table + ) + else: + logger.warning( + "recording enabled but recording_config.db_url is None; falling " + "back to in-process MemoryStore (no cross-process visibility)" + ) + store = MemoryStore() + + # Rank is constant per process; capture once (RequestOutput does not expose + # parallel_config, so we read it from engine_client here, mirroring + # api_patch_v17.py:148). + try: + rank = int(engine_client.vllm_config.parallel_config._api_process_rank) + except Exception: + rank = 0 + + recorder = Recorder( + store=store, + topk=recording_config.topk, + enabled=True, + rank=rank, + engine_client=engine_client, + ) + + # (1) engine-level wrap — before init_app_state so serving objects inherit + # the wrapped reference. Idempotent via the __patched_*__ guard. + patch_engine_for_recording(engine_client, recorder, logger) + + # (2) in-process middleware: X-Session-ID -> contextvar. Zero network hop. + app.add_middleware(SessionMiddleware) + + # (3) query routes mounted on the main app; OpenAI /v1/* surface untouched. + app.include_router(query_router) + + setattr(app.state, _STORE_STATE_ATTR, store) + setattr(app.state, _RECORDER_STATE_ATTR, recorder) + + logger.info( + "Generation recording enabled: topk=%d store=%s rank=%d", + recording_config.topk, + type(store).__name__, + rank, + ) + return recorder + + +async def run_server_worker_in_ray( + listen_address, + sock, + args, + engine_client, + logger, + recording_config: Optional[RecordingConfig] = None, +) -> None: + """Modified from vllm.entrypoints.openai.api_server.run_server_worker. + + Differs from api_patch_v17.py only in the recording wiring inserted between + ``build_app`` and ``init_app_state``, plus starting/stopping the recorder + flusher around ``serve_http``. + """ + if args.tool_parser_plugin and len(args.tool_parser_plugin) > 3: + ToolParserManager.import_tool_parser(args.tool_parser_plugin) + + if args.reasoning_parser_plugin and len(args.reasoning_parser_plugin) > 3: + ReasoningParserManager.import_reasoning_parser(args.reasoning_parser_plugin) + + app = build_app(args) + + # --- recording wiring: engine wrap must precede init_app_state ----------- + recorder = _setup_recording(args, engine_client, app, logger, recording_config=recording_config) + # ------------------------------------------------------------------------ + + await init_app_state(engine_client, app.state, args) + + loop = asyncio.get_event_loop() + loop.add_signal_handler = functools.partial(dummy_add_signal_handler, loop) + + logger.info( + "Starting vLLM API server %d on %s", + engine_client.vllm_config.parallel_config._api_process_rank, + listen_address, + ) + + if recorder is not None: + recorder.start() + + shutdown_task = await serve_http( + app, + sock=sock, + enable_ssl_refresh=args.enable_ssl_refresh, + host=args.host, + port=args.port, + log_level=args.uvicorn_log_level, + # NOTE: When the 'disable_uvicorn_access_log' value is True, + # no access log will be output. + access_log=not args.disable_uvicorn_access_log, + timeout_keep_alive=envs.VLLM_HTTP_TIMEOUT_KEEP_ALIVE, + ssl_keyfile=args.ssl_keyfile, + ssl_certfile=args.ssl_certfile, + ssl_ca_certs=args.ssl_ca_certs, + ssl_cert_reqs=args.ssl_cert_reqs, + h11_max_incomplete_event_size=args.h11_max_incomplete_event_size, + h11_max_header_count=args.h11_max_header_count, + ) + + # NB: Await server shutdown only after the backend context is exited + try: + await shutdown_task + finally: + if recorder is not None: + await recorder.stop() + sock.close() + + +async def run_server_in_ray( + args, + engine_client, + logger, + recording_config: Optional[RecordingConfig] = None, +): + # Modified from vllm.entrypoints.openai.api_server.run_server + listen_address, sock = setup_server_in_ray(args, logger) + logger.info("vLLM API server listening on %s", listen_address) + await run_server_worker_in_ray( + listen_address, sock, args, engine_client, logger, recording_config + ) + + +async def run_api_server_with_recording( + async_llm, + host: str, + port: int, + model_path: str, + logger: logging.Logger, + chat_template: Optional[str] = None, + enable_auto_tool_choice: bool = False, + tool_call_parser: Optional[str] = None, + reasoning_parser: Optional[str] = None, + enable_log_requests: bool = False, + recording_config: Optional[RecordingConfig] = None, +): + """Drop-in recording-enabled variant of + ``api_patch_v17.run_api_server_in_ray_actor_v17``. + + Same signature plus an optional ``recording_config`` so launchers can + switch by import path. Requires vllm >= 0.17.0. Recording is on iff + ``recording_config`` is provided (built by ``get_api_server`` from + ``InferenceModelConfig`` when ``enable_recording`` is on). The dynamic + checkpoint version is read off ``async_llm.trinity_model_version`` + (mirrored by VLLMModel), so it is not part of the static config here. + """ + vllm_version = get_vllm_version() + if vllm_version < parse_version("0.17.0"): + raise ValueError( + f"Unsupported vllm version: {vllm.__version__}. " + "This patch requires vllm version >= 0.17.0" + ) + + parser = FlexibleArgumentParser(description="Run the OpenAI API server.") + args = make_arg_parser(parser) + cli_args = [ + "--host", + str(host), + "--port", + str(port), + "--model", + model_path, + "--enable-server-load-tracking", # enable tracking for load balancing + ] + if enable_log_requests: + cli_args.append("--enable-log-requests") + if enable_auto_tool_choice: + cli_args.append("--enable-auto-tool-choice") + if tool_call_parser: + cli_args.extend(["--tool-call-parser", tool_call_parser]) + if reasoning_parser: + cli_args.extend(["--reasoning-parser", reasoning_parser]) + if chat_template: + cli_args.extend(["--chat-template", chat_template]) + + # NOTE: routed_experts capture and the logprobs cap are ENGINE-level + # ModelConfig fields (consumed by the scheduler/worker, not the API serving + # layer), so they take effect at engine build time — which in this launch + # path happens in VLLMModel (via EngineArgs), *before* this runner gets the + # already-built ``async_llm``. Adding ``--enable-return-routed-experts`` / + # ``--max-logprobs`` here would be inert (init_app_state does not read them). + # The Allocator therefore forces ``InferenceModelConfig.enable_return_routed_experts + # = True`` when recording is on, and the engine's default ``max_logprobs=20`` + # covers the recorder's top-k (``VLLM_RECORD_TOPK``, default 1). To record + # routed_experts, the engine must be built with that flag on — the launcher + # is responsible for that, not these CLI args. + + args = parser.parse_args(cli_args) + args.structured_outputs_config.reasoning_parser = reasoning_parser + logger.info(f"Starting vLLM OpenAI API server with args: {args}") + await run_server_in_ray(args, async_llm, logger, recording_config) diff --git a/trinity/common/models/vllm_patch/recording/store.py b/trinity/common/models/vllm_patch/recording/store.py new file mode 100644 index 00000000000..481059a3c3d --- /dev/null +++ b/trinity/common/models/vllm_patch/recording/store.py @@ -0,0 +1,136 @@ +"""Pluggable storage backends for recorded experiences. + +A ``RecordStore`` persists Trinity ``Experience`` objects. Backends: + +* ``MemoryStore`` — in-process, request/session-keyed; for standalone runs and + the ``/records`` debug endpoints. +* ``SqlStore`` — delegates to ``trinity.explorer.proxy.recorder.HistoryRecorder`` + so the in-vLLM recorder writes to the *same* SQL table the explorer proxy + reads (``proxy_history``). This is the online-RL path: experiences written + here by the vLLM process are later picked up by the proxy's + ``update_reward``/``submit_experiences`` via the shared ``msg_id`` key. + +Keying: experiences are identified by ``eid.suffix`` (the vLLM ``request_id``, +== the OpenAI ``response.id`` == the proxy ``msg_id``). ``eid.task``/``run``/ +``reward`` are assigned by the proxy at feedback time, not here. + +Concurrency: ``append_turn`` is called from a single background flusher task; +the async signatures keep the door open for I/O-bound backends. +""" +import abc +from collections import defaultdict +from typing import Optional + +from trinity.common.experience import Experience + +#: Attribute carrying the vLLM request id on each experience's ``info`` dict. +_REQUEST_ID_INFO_KEY = "request_id" +#: Attribute carrying the task id (X-Session-ID) on each experience's ``info``. +_TASK_ID_INFO_KEY = "task_id" + + +class RecordStore(abc.ABC): + """Abstract persistence interface for recorded experiences.""" + + @abc.abstractmethod + async def append_turn(self, exp: Experience) -> None: + """Persist one completed experience.""" + + @abc.abstractmethod + async def get_task(self, task_id: str) -> list[Experience]: + """Return all experiences for a task, in insertion order.""" + + @abc.abstractmethod + async def get_turn(self, task_id: str, request_id: str) -> Optional[Experience]: + """Return a single experience, or None if not found.""" + + @abc.abstractmethod + async def list_tasks(self) -> list[str]: + """Return all known task ids.""" + + @abc.abstractmethod + async def delete_task(self, task_id: str) -> None: + """Drop all experiences for a task.""" + + +class MemoryStore(RecordStore): + """In-process store. + + Groups experiences by session (``info["task_id"]``) when an X-Session-ID + was supplied, otherwise each turn is keyed by its own ``eid.suffix`` + (request_id) — so a missing session header never collapses distinct turns. + ``get_turn`` resolves an individual turn by ``info["request_id"]``. + + Note: per-process under data-parallel serving — each API-server rank holds + only the experiences it served. For cross-rank aggregation, use ``SqlStore``. + """ + + def __init__(self) -> None: + # group key -> [experience, ...] in insertion order. + self._records: dict[str, list[Experience]] = defaultdict(list) + + @staticmethod + def _group_key(exp: Experience) -> str: + session = exp.info.get(_TASK_ID_INFO_KEY) + return session if session else exp.eid.suffix + + async def append_turn(self, exp: Experience) -> None: + self._records[self._group_key(exp)].append(exp) + + async def get_task(self, task_id: str) -> list[Experience]: + return list(self._records.get(task_id, [])) + + async def get_turn(self, task_id: str, request_id: str) -> Optional[Experience]: + for exp in self._records.get(task_id, []): + if exp.info.get(_REQUEST_ID_INFO_KEY) == request_id: + return exp + return None + + async def list_tasks(self) -> list[str]: + return list(self._records.keys()) + + async def delete_task(self, task_id: str) -> None: + self._records.pop(task_id, None) + + +class SqlStore(RecordStore): + """SQL-backed store sharing the explorer proxy's ``proxy_history`` table. + + Writes go through ``HistoryRecorder.record_history`` (which ``prepare()``s + the engine on first use and maps ``eid.suffix`` -> ``msg_id``). The explorer + proxy's own ``HistoryRecorder`` instance reads/updates the same rows for + ``/feedback`` and ``/commit``, so the in-vLLM recorder and the proxy share + one table by ``db_url`` + ``table_name``. + + Reads (``get_task``/``get_turn``/``list_tasks``/``delete_task``) are NOT + implemented here: in the online-RL setup the proxy owns the read/consume + side. The ``/records`` query endpoints surface this as 503 when this backend + is active. + """ + + #: Marks that this backend does not serve the ``/records`` read endpoints. + supports_reads = False + + def __init__(self, db_url: str, table_name: str) -> None: + # Imported lazily so the vLLM process only pulls in the SQL/explorer + # stack when this backend is actually selected. + from trinity.explorer.proxy.recorder import HistoryRecorder + + self._recorder = HistoryRecorder(db_url=db_url, table_name=table_name) + + async def append_turn(self, exp: Experience) -> None: + # record_history() calls prepare() on first use; serializes the + # experience into the blob column and writes meta keyed by msg_id. + await self._recorder.record_history([exp]) + + async def get_task(self, task_id: str) -> list[Experience]: + raise NotImplementedError("SqlStore reads are served by the explorer proxy; use /feedback") + + async def get_turn(self, task_id: str, request_id: str) -> Optional[Experience]: + raise NotImplementedError("SqlStore reads are served by the explorer proxy; use /feedback") + + async def list_tasks(self) -> list[str]: + raise NotImplementedError("SqlStore reads are served by the explorer proxy; use /feedback") + + async def delete_task(self, task_id: str) -> None: + raise NotImplementedError("SqlStore reads are served by the explorer proxy; use /feedback") diff --git a/trinity/explorer/proxy/app.py b/trinity/explorer/proxy/app.py index 067fdc2614b..cbc97eb57f4 100644 --- a/trinity/explorer/proxy/app.py +++ b/trinity/explorer/proxy/app.py @@ -1,7 +1,6 @@ -import json import traceback from contextlib import asynccontextmanager -from typing import Any, Dict +from typing import Dict import httpx import uvicorn @@ -65,113 +64,25 @@ def _build_json_or_text_response(upstream_response: httpx.Response): ) -def _consume_sse_line(line: str, aggregate: Dict[str, Any]) -> None: - line = line.strip() - if not line or not line.startswith("data:"): - return - - payload = line[5:].strip() - if not payload or payload == "[DONE]": - return - - try: - data = json.loads(payload) - except json.JSONDecodeError: - return - - if isinstance(data.get("id"), str) and data["id"]: - aggregate["id"] = data["id"] - - prompt_token_ids = data.get("prompt_token_ids") - if isinstance(prompt_token_ids, list) and prompt_token_ids: - aggregate["prompt_token_ids"] = prompt_token_ids - - for choice in data.get("choices", []): - if not isinstance(choice, dict): - continue - - choice_index = choice.get("index", 0) - if not isinstance(choice_index, int): - choice_index = 0 - - choice_acc = aggregate["choices"].setdefault( - choice_index, - { - "index": choice_index, - "token_ids": [], - "logprobs": {"content": []}, - }, - ) - - token_ids = choice.get("token_ids") - if isinstance(token_ids, list) and token_ids: - choice_acc["token_ids"].extend(token_ids) - - logprobs = choice.get("logprobs") - if isinstance(logprobs, dict): - content = logprobs.get("content") - if isinstance(content, list) and content: - choice_acc["logprobs"]["content"].extend(content) - - -def _finalize_stream_aggregate(aggregate: Dict[str, Any]) -> Dict[str, Any] | None: - prompt_token_ids = aggregate.get("prompt_token_ids") - if not isinstance(prompt_token_ids, list) or not prompt_token_ids: - return None - - ordered_choices = [] - for _, choice in sorted(aggregate["choices"].items(), key=lambda item: item[0]): - if not choice.get("token_ids"): - continue - ordered_choices.append(choice) - - if not ordered_choices: - return None - - return { - "id": aggregate.get("id", ""), - "prompt_token_ids": prompt_token_ids, - "choices": ordered_choices, - } - - -async def _proxy_chat_stream_with_experience( +async def _proxy_chat_stream( request: Request, upstream_response: httpx.Response, - model_version: int, ): - async def iterator(): - stream_buffer = "" - aggregate = { - "id": "", - "prompt_token_ids": [], - "choices": {}, - } + """Pure passthrough: stream the upstream SSE bytes to the client unchanged. + + Experience capture is handled in-process by the vLLM recorder (wrapping + ``engine_client.generate``), so the proxy no longer parses/aggregates the + stream here. + """ + async def iterator(): try: async for chunk in upstream_response.aiter_raw(): if chunk: - stream_buffer += chunk.decode("utf-8", errors="ignore") - while "\n" in stream_buffer: - line, stream_buffer = stream_buffer.split("\n", 1) - _consume_sse_line(line.rstrip("\r"), aggregate) yield chunk finally: - if stream_buffer: - _consume_sse_line(stream_buffer.rstrip("\r"), aggregate) - await upstream_response.aclose() - experience_response = _finalize_stream_aggregate(aggregate) - if experience_response is not None: - try: - await request.app.state.service.record_experience( - experience_response, - model_version, - ) - except Exception: - pass - return StreamingResponse( content=iterator(), status_code=upstream_response.status_code, @@ -189,16 +100,15 @@ async def chat_completions(request: Request): raise HTTPException(status_code=400, detail=f"Invalid JSON: {str(e)}") forward_headers = _build_forward_headers(request) - # for experience data recording, we need to return token ids and logprobs - request_data["return_token_ids"] = True - request_data["logprobs"] = True - # temperature must be set from config, ignore user's input + # Temperature is a policy knob controlled by the explorer config; override + # the client's value. (Experience capture — token_ids/logprobs — is handled + # in-process by the vLLM recorder, so we no longer force them onto the wire.) request_data["temperature"] = request.app.state.temperature - url, model_version = await request.app.state.service.allocate_model() + url, _ = await request.app.state.service.allocate_model() if request_data.get("stream", False): - # For streaming response, we need to handle it differently to aggregate experience data + # Streaming: passthrough the upstream SSE bytes unchanged. try: upstream_request = request.app.state.http_client.build_request( method="POST", @@ -234,10 +144,9 @@ async def chat_completions(request: Request): }, ) - return await _proxy_chat_stream_with_experience( + return await _proxy_chat_stream( request=request, upstream_response=upstream_response, - model_version=model_version, ) try: @@ -282,7 +191,9 @@ async def chat_completions(request: Request): headers=_build_downstream_headers(resp.headers), ) - await request.app.state.service.record_experience(resp_data, model_version) + # Non-streaming success: forward unchanged. Experience capture happens + # in-process at the vLLM engine boundary (the recorder wraps + # engine_client.generate), so nothing to record here. return JSONResponse( status_code=resp.status_code, content=resp_data, diff --git a/trinity/explorer/proxy/service.py b/trinity/explorer/proxy/service.py index 3f3362f41f0..1fe77d1564a 100644 --- a/trinity/explorer/proxy/service.py +++ b/trinity/explorer/proxy/service.py @@ -3,8 +3,6 @@ from collections import deque from typing import Dict, List, Tuple -import torch - from trinity.common.constants import RunningStatus, SyncMethod from trinity.common.experience import Experience from trinity.common.models.model import ModelWrapper @@ -39,7 +37,6 @@ def __init__(self, explorer: Explorer, listen_address: str = "localhost", port: or f"sqlite:///{explorer.config.buffer.cache_dir}/proxy_history.db", table_name="proxy_history", ) - self.total_experience_count = 0 self.ready_experience_count = 0 async def serve(self) -> None: @@ -132,37 +129,9 @@ def collect_metrics(self) -> Dict: for i, model in enumerate(self.models): metrics[f"rollout/model_{i}/total_request_count"] = model.request_count metrics[f"rollout/model_{i}/model_version"] = model.model_version - metrics["rollout/total_experience_count"] = self.total_experience_count metrics["rollout/ready_experience_count"] = self.ready_experience_count return metrics - async def record_experience(self, response, model_version: int) -> None: - experiences = [] - for choice in response["choices"]: - exp = Experience( - tokens=torch.cat( - ( - torch.tensor(response["prompt_token_ids"], dtype=torch.int32), - torch.tensor(choice["token_ids"], dtype=torch.int32), - ) - ), - logprobs=( - torch.tensor( - [logprob["logprob"] for logprob in choice["logprobs"]["content"]], - dtype=torch.float32, - ) - if "logprobs" in choice and choice["logprobs"] is not None - else torch.tensor([], dtype=torch.float32) - ), - prompt_length=len(response["prompt_token_ids"]), - ) - exp.eid.suffix = response["id"] - exp.info["model_version"] = model_version - experiences.append(exp) - - self.total_experience_count += len(experiences) - await self.recorder.record_history(experiences) - async def submit_experiences(self) -> None: async with self.commit_lock: experiences = list(self.ready_experiences) From 5fe90531113b3cc29ec8c7ac594aebffb6b5709d Mon Sep 17 00:00:00 2001 From: pxc Date: Thu, 25 Jun 2026 11:15:26 +0800 Subject: [PATCH 02/84] simplify --- trinity/common/models/vllm_patch/__init__.py | 5 +++-- trinity/common/models/vllm_patch/recording/__init__.py | 9 +-------- trinity/common/models/vllm_patch/recording/server.py | 10 ++-------- 3 files changed, 6 insertions(+), 18 deletions(-) diff --git a/trinity/common/models/vllm_patch/__init__.py b/trinity/common/models/vllm_patch/__init__.py index 7e06a279bf4..692ae4e1cc7 100644 --- a/trinity/common/models/vllm_patch/__init__.py +++ b/trinity/common/models/vllm_patch/__init__.py @@ -10,6 +10,7 @@ VLLM_VERSION_0120 = parse_version("0.12.0") VLLM_VERSION_0170 = parse_version("0.17.0") +VLLM_VERSION_0230 = parse_version("0.23.0") def vllm_patch(): @@ -155,10 +156,10 @@ def get_api_server( # Recording is driven by the config field (not env, not an engine attr): # get_api_server already receives InferenceModelConfig, so it builds the # static RecordingConfig here and threads it into the recording runner. - recording = bool(config.enable_recording) and vllm_version >= VLLM_VERSION_0170 + recording = bool(config.enable_recording) and vllm_version >= VLLM_VERSION_0230 if config.enable_recording and not recording: logger.warning( - "enable_recording is on but vLLM %s < 0.17.0; recording disabled", + "enable_recording is on but vLLM %s < 0.23.0; recording disabled", vllm.__version__, ) diff --git a/trinity/common/models/vllm_patch/recording/__init__.py b/trinity/common/models/vllm_patch/recording/__init__.py index fa490baeb3f..280d1dcfa03 100644 --- a/trinity/common/models/vllm_patch/recording/__init__.py +++ b/trinity/common/models/vllm_patch/recording/__init__.py @@ -1,12 +1,5 @@ """Generation-recording patch for the vLLM OpenAI server. - -Self-contained entry point that mirrors ``api_patch_v17.py``'s bootstrap flow -(``build_app`` -> ``init_app_state`` -> ``serve_http``) and additionally wires -in generation recording without touching vLLM source code. - -Designed for vllm >= 0.17.0. Drop-in alternative to -``trinity.common.models.vllm_patch.api_patch_v17``: point your launcher at -``recording.run_api_server_with_recording`` instead. +Designed for vllm >= 0.23.0. """ from trinity.common.models.vllm_patch.recording.config import RecordingConfig from trinity.common.models.vllm_patch.recording.context import ( diff --git a/trinity/common/models/vllm_patch/recording/server.py b/trinity/common/models/vllm_patch/recording/server.py index c88774650e7..565ef785677 100644 --- a/trinity/common/models/vllm_patch/recording/server.py +++ b/trinity/common/models/vllm_patch/recording/server.py @@ -35,6 +35,7 @@ validate_api_server_args, ) from vllm.entrypoints.openai.cli_args import make_arg_parser +from vllm.entrypoints.serve.utils.api_utils import log_non_default_args from vllm.reasoning import ReasoningParserManager from vllm.tool_parsers import ToolParserManager from vllm.utils.argparse_utils import FlexibleArgumentParser @@ -71,14 +72,7 @@ def setup_server_in_ray(args, logger): """ logger.info("vLLM API server version %s", VLLM_VERSION) - if get_vllm_version() < parse_version("0.23.0"): - from vllm.entrypoints.utils import log_non_default_args - - log_non_default_args(args) - else: - from vllm.entrypoints.serve.utils.api_utils import log_non_default_args - - log_non_default_args(args) + log_non_default_args(args) if args.tool_parser_plugin and len(args.tool_parser_plugin) > 3: ToolParserManager.import_tool_parser(args.tool_parser_plugin) From 180b4168e37017395d8d8523ac79e2e783517f6a Mon Sep 17 00:00:00 2001 From: pxc Date: Thu, 25 Jun 2026 11:18:11 +0800 Subject: [PATCH 03/84] add config --- .../models/vllm_patch/recording/config.py | 33 +++++++++++++++++++ 1 file changed, 33 insertions(+) create mode 100644 trinity/common/models/vllm_patch/recording/config.py diff --git a/trinity/common/models/vllm_patch/recording/config.py b/trinity/common/models/vllm_patch/recording/config.py new file mode 100644 index 00000000000..30bdcd08acc --- /dev/null +++ b/trinity/common/models/vllm_patch/recording/config.py @@ -0,0 +1,33 @@ +# -*- coding: utf-8 -*- +"""Recording configuration carried explicitly through the launch chain. + +Instead of env vars or attributes bolted onto the engine instance, the static +recording config (db url / table / top-k) is passed as an explicit parameter +into ``run_api_server_with_recording``. ``get_api_server`` builds one from +``InferenceModelConfig`` when ``enable_recording`` is on. + +The *dynamic* checkpoint version is NOT here: it changes at runtime +(``sync_model_weights``), so the recorder reads it live off the engine +instance attribute ``trinity_model_version`` (mirrored by VLLMModel). +""" +from dataclasses import dataclass +from typing import Optional + + +@dataclass +class RecordingConfig: + """Static configuration for the in-vLLM recorder. + + Attributes: + db_url: SQL db url the recorder writes to (shared with the explorer + proxy's HistoryRecorder). When None, the recorder falls back to an + in-process MemoryStore (no cross-process visibility). + table: SQL table name (default ``proxy_history``, matching the proxy). + topk: How many top-k logprobs the engine computes per generated token. + Only the chosen token's logprob is stored, and vLLM force-includes + the sampled token, so 1 suffices by default. + """ + + db_url: Optional[str] = None + table: str = "proxy_history" + topk: int = 1 From f27988b9ccaa01b493d6a9ce2e8492e4530abcfa Mon Sep 17 00:00:00 2001 From: "panxuchen.pxc" Date: Thu, 25 Jun 2026 11:23:21 +0800 Subject: [PATCH 04/84] use api key as session id --- .../models/vllm_patch/recording/__init__.py | 2 + .../models/vllm_patch/recording/context.py | 57 +++++++++++++++---- .../models/vllm_patch/recording/models.py | 4 +- .../models/vllm_patch/recording/server.py | 14 +++-- .../models/vllm_patch/recording/store.py | 9 +-- 5 files changed, 65 insertions(+), 21 deletions(-) diff --git a/trinity/common/models/vllm_patch/recording/__init__.py b/trinity/common/models/vllm_patch/recording/__init__.py index 280d1dcfa03..4bdf5ebc272 100644 --- a/trinity/common/models/vllm_patch/recording/__init__.py +++ b/trinity/common/models/vllm_patch/recording/__init__.py @@ -3,6 +3,7 @@ """ from trinity.common.models.vllm_patch.recording.config import RecordingConfig from trinity.common.models.vllm_patch.recording.context import ( + RecordingIdentityMiddleware, SessionMiddleware, task_id_ctx, ) @@ -25,6 +26,7 @@ "MemoryStore", "RecordStore", "RecordingConfig", + "RecordingIdentityMiddleware", "SqlStore", "Recorder", "SessionMiddleware", diff --git a/trinity/common/models/vllm_patch/recording/context.py b/trinity/common/models/vllm_patch/recording/context.py index 15db3a332b0..195c5ad9134 100644 --- a/trinity/common/models/vllm_patch/recording/context.py +++ b/trinity/common/models/vllm_patch/recording/context.py @@ -1,11 +1,13 @@ -"""Per-request task_id propagation. +"""Per-request recording identity propagation. -The OpenAI ``X-Session-ID`` request header carries the task id that scopes a -multi-turn task. We read it in an in-process ASGI middleware and stash it in a -contextvar so the engine-level wrapper (which runs in the same async task as -the serving handler) can recover it at record time. +The OpenAI ``Authorization: Bearer `` header is the preferred identity +source for recording because it works with CLI agents that do not support +custom headers. ``X-Session-ID`` remains a compatibility fallback. We read the +identity in an in-process ASGI middleware and stash it in a contextvar so the +engine-level wrapper (which runs in the same async task as the serving handler) +can recover it at record time. -No ``X-Session-ID`` on a request is fine: the recorder falls back to +No identity header on a request is fine: the recorder falls back to ``request_id`` as the task id so nothing is silently dropped. """ from contextvars import ContextVar @@ -15,24 +17,57 @@ from starlette.requests import Request # The task id for the in-flight request, or None when the client did not send -# ``X-Session-ID`` (the recorder then uses request_id as a fallback task id). +# a supported identity header (the recorder then uses request_id as a fallback +# task id). task_id_ctx: ContextVar[Optional[str]] = ContextVar("trinity_recording_task_id", default=None) -#: Canonical header name. Lower-cased per ASGI/httpx convention. +#: Preferred identity header for OpenAI-compatible clients. +AUTHORIZATION_HEADER = "authorization" +#: Compatibility identity header. Lower-cased per ASGI/httpx convention. SESSION_ID_HEADER = "x-session-id" -class SessionMiddleware(BaseHTTPMiddleware): - """Capture ``X-Session-ID`` into ``task_id_ctx`` for the request's lifetime. +def extract_bearer_token(authorization: Optional[str]) -> Optional[str]: + """Extract the bearer token from an Authorization header. + + Returns None when the header is missing or does not use the Bearer scheme. + """ + if not authorization: + return None + scheme, _, token = authorization.partition(" ") + if scheme.lower() != "bearer": + return None + token = token.strip() + return token or None + + +def get_recording_task_id(request: Request) -> Optional[str]: + """Return the recording identity for a request. + + Prefer OpenAI-compatible API keys because every supported agent platform + can pass them. ``X-Session-ID`` is kept as a backward-compatible fallback + for direct clients and debugging. + """ + return extract_bearer_token(request.headers.get(AUTHORIZATION_HEADER)) or request.headers.get( + SESSION_ID_HEADER + ) + + +class RecordingIdentityMiddleware(BaseHTTPMiddleware): + """Capture request identity into ``task_id_ctx`` for the request's lifetime. Runs in-process (ASGI) — no extra network hop, no serialization cost beyond a contextvar set/reset. """ async def dispatch(self, request: Request, call_next): - task_id = request.headers.get(SESSION_ID_HEADER) + task_id = get_recording_task_id(request) token = task_id_ctx.set(task_id) try: return await call_next(request) finally: task_id_ctx.reset(token) + + +# Backward-compatible export for existing imports. +SessionMiddleware = RecordingIdentityMiddleware diff --git a/trinity/common/models/vllm_patch/recording/models.py b/trinity/common/models/vllm_patch/recording/models.py index d7e60b49140..736c15f840d 100644 --- a/trinity/common/models/vllm_patch/recording/models.py +++ b/trinity/common/models/vllm_patch/recording/models.py @@ -12,6 +12,7 @@ ``eid.task``/``run``/``reward`` are left default here and assigned later by the proxy at ``/feedback`` time, matching ``explorer/proxy/service.record_experience`` semantics. + API key / X-Session-ID -> info["task_id"] (traceability only; not used as a key) prompt_token_ids -> tokens (prompt portion) + prompt_length response_token_ids-> tokens (response portion) @@ -84,7 +85,8 @@ def build_experience( Args: output: A ``RequestOutput`` with ``finished == True``. - task_id: From ``X-Session-ID``; stored in ``info`` for traceability only. + task_id: From the request API key, or legacy ``X-Session-ID``; stored + in ``info`` for traceability only. Not used as the storage key — ``eid.suffix`` is, so a missing ``X-Session-ID`` never drops a turn. rank: Data-parallel serving rank. diff --git a/trinity/common/models/vllm_patch/recording/server.py b/trinity/common/models/vllm_patch/recording/server.py index 565ef785677..806779d76f4 100644 --- a/trinity/common/models/vllm_patch/recording/server.py +++ b/trinity/common/models/vllm_patch/recording/server.py @@ -12,8 +12,9 @@ 1. ``patch_engine_for_recording`` — instance-level wrap of ``engine_client.generate`` to force top-k logprobs and record finished ``RequestOutput`` (covers chat/completion/responses, streaming and not). - 2. ``SessionMiddleware`` — in-process ASGI middleware reading - ``X-Session-ID`` into a contextvar. + 2. ``RecordingIdentityMiddleware`` — in-process ASGI middleware reading + ``Authorization: Bearer `` (or legacy ``X-Session-ID``) into a + contextvar. 3. ``query_router`` — ``/records/*`` endpoints for later analysis. Only for vllm versions >= 0.17.0. @@ -45,7 +46,9 @@ from trinity.common.models.vllm_patch import get_vllm_version from trinity.common.models.vllm_patch.recording.config import RecordingConfig -from trinity.common.models.vllm_patch.recording.context import SessionMiddleware +from trinity.common.models.vllm_patch.recording.context import ( + RecordingIdentityMiddleware, +) from trinity.common.models.vllm_patch.recording.query import query_router from trinity.common.models.vllm_patch.recording.recorder import ( Recorder, @@ -169,8 +172,9 @@ def _setup_recording( # the wrapped reference. Idempotent via the __patched_*__ guard. patch_engine_for_recording(engine_client, recorder, logger) - # (2) in-process middleware: X-Session-ID -> contextvar. Zero network hop. - app.add_middleware(SessionMiddleware) + # (2) in-process middleware: API key / X-Session-ID -> contextvar. + # Zero network hop. + app.add_middleware(RecordingIdentityMiddleware) # (3) query routes mounted on the main app; OpenAI /v1/* surface untouched. app.include_router(query_router) diff --git a/trinity/common/models/vllm_patch/recording/store.py b/trinity/common/models/vllm_patch/recording/store.py index 481059a3c3d..27ec0df0033 100644 --- a/trinity/common/models/vllm_patch/recording/store.py +++ b/trinity/common/models/vllm_patch/recording/store.py @@ -25,7 +25,7 @@ #: Attribute carrying the vLLM request id on each experience's ``info`` dict. _REQUEST_ID_INFO_KEY = "request_id" -#: Attribute carrying the task id (X-Session-ID) on each experience's ``info``. +#: Attribute carrying the recording identity on each experience's ``info``. _TASK_ID_INFO_KEY = "task_id" @@ -56,9 +56,10 @@ async def delete_task(self, task_id: str) -> None: class MemoryStore(RecordStore): """In-process store. - Groups experiences by session (``info["task_id"]``) when an X-Session-ID - was supplied, otherwise each turn is keyed by its own ``eid.suffix`` - (request_id) — so a missing session header never collapses distinct turns. + Groups experiences by recording identity (``info["task_id"]``) when an API + key or legacy X-Session-ID was supplied, otherwise each turn is keyed by its + own ``eid.suffix`` (request_id) — so a missing identity never collapses + distinct turns. ``get_turn`` resolves an individual turn by ``info["request_id"]``. Note: per-process under data-parallel serving — each API-server rank holds From 40b53e87e842f9ae72ab76d9d005a71dc2150c92 Mon Sep 17 00:00:00 2001 From: "panxuchen.pxc" Date: Thu, 25 Jun 2026 11:30:13 +0800 Subject: [PATCH 05/84] clean stale header --- .../models/vllm_patch/recording/__init__.py | 2 -- .../models/vllm_patch/recording/context.py | 26 ++++++------------- .../models/vllm_patch/recording/models.py | 9 +++---- .../models/vllm_patch/recording/recorder.py | 4 +-- .../models/vllm_patch/recording/server.py | 6 ++--- .../models/vllm_patch/recording/store.py | 5 ++-- 6 files changed, 18 insertions(+), 34 deletions(-) diff --git a/trinity/common/models/vllm_patch/recording/__init__.py b/trinity/common/models/vllm_patch/recording/__init__.py index 4bdf5ebc272..adb525a172c 100644 --- a/trinity/common/models/vllm_patch/recording/__init__.py +++ b/trinity/common/models/vllm_patch/recording/__init__.py @@ -4,7 +4,6 @@ from trinity.common.models.vllm_patch.recording.config import RecordingConfig from trinity.common.models.vllm_patch.recording.context import ( RecordingIdentityMiddleware, - SessionMiddleware, task_id_ctx, ) from trinity.common.models.vllm_patch.recording.models import build_experience @@ -29,7 +28,6 @@ "RecordingIdentityMiddleware", "SqlStore", "Recorder", - "SessionMiddleware", "build_experience", "patch_engine_for_recording", "query_router", diff --git a/trinity/common/models/vllm_patch/recording/context.py b/trinity/common/models/vllm_patch/recording/context.py index 195c5ad9134..7244c15c987 100644 --- a/trinity/common/models/vllm_patch/recording/context.py +++ b/trinity/common/models/vllm_patch/recording/context.py @@ -1,11 +1,10 @@ """Per-request recording identity propagation. -The OpenAI ``Authorization: Bearer `` header is the preferred identity -source for recording because it works with CLI agents that do not support -custom headers. ``X-Session-ID`` remains a compatibility fallback. We read the -identity in an in-process ASGI middleware and stash it in a contextvar so the -engine-level wrapper (which runs in the same async task as the serving handler) -can recover it at record time. +The OpenAI ``Authorization: Bearer `` header is the recording identity +source because it works with CLI agents that do not support custom headers. We +read the identity in an in-process ASGI middleware and stash it in a contextvar +so the engine-level wrapper (which runs in the same async task as the serving +handler) can recover it at record time. No identity header on a request is fine: the recorder falls back to ``request_id`` as the task id so nothing is silently dropped. @@ -23,8 +22,6 @@ #: Preferred identity header for OpenAI-compatible clients. AUTHORIZATION_HEADER = "authorization" -#: Compatibility identity header. Lower-cased per ASGI/httpx convention. -SESSION_ID_HEADER = "x-session-id" def extract_bearer_token(authorization: Optional[str]) -> Optional[str]: @@ -44,13 +41,10 @@ def extract_bearer_token(authorization: Optional[str]) -> Optional[str]: def get_recording_task_id(request: Request) -> Optional[str]: """Return the recording identity for a request. - Prefer OpenAI-compatible API keys because every supported agent platform - can pass them. ``X-Session-ID`` is kept as a backward-compatible fallback - for direct clients and debugging. + OpenAI-compatible API keys are used because every supported agent platform + can pass them. """ - return extract_bearer_token(request.headers.get(AUTHORIZATION_HEADER)) or request.headers.get( - SESSION_ID_HEADER - ) + return extract_bearer_token(request.headers.get(AUTHORIZATION_HEADER)) class RecordingIdentityMiddleware(BaseHTTPMiddleware): @@ -67,7 +61,3 @@ async def dispatch(self, request: Request, call_next): return await call_next(request) finally: task_id_ctx.reset(token) - - -# Backward-compatible export for existing imports. -SessionMiddleware = RecordingIdentityMiddleware diff --git a/trinity/common/models/vllm_patch/recording/models.py b/trinity/common/models/vllm_patch/recording/models.py index 736c15f840d..2da1b398224 100644 --- a/trinity/common/models/vllm_patch/recording/models.py +++ b/trinity/common/models/vllm_patch/recording/models.py @@ -12,8 +12,7 @@ ``eid.task``/``run``/``reward`` are left default here and assigned later by the proxy at ``/feedback`` time, matching ``explorer/proxy/service.record_experience`` semantics. - API key / - X-Session-ID -> info["task_id"] (traceability only; not used as a key) + API key -> info["task_id"] (traceability only; not used as a key) prompt_token_ids -> tokens (prompt portion) + prompt_length response_token_ids-> tokens (response portion) logprobs -> Experience.logprobs -- but ONLY the *chosen* token's @@ -85,10 +84,10 @@ def build_experience( Args: output: A ``RequestOutput`` with ``finished == True``. - task_id: From the request API key, or legacy ``X-Session-ID``; stored - in ``info`` for traceability only. + task_id: From the request API key; stored in ``info`` for traceability + only. Not used as the storage key — ``eid.suffix`` is, so a missing - ``X-Session-ID`` never drops a turn. + API key never drops a turn. rank: Data-parallel serving rank. timestamp: UTC ISO-8601 string (caller-stamped to keep this pure). endpoint: Which OpenAI endpoint served the turn (best-effort). diff --git a/trinity/common/models/vllm_patch/recording/recorder.py b/trinity/common/models/vllm_patch/recording/recorder.py index 1db9b82f339..106e104e8b5 100644 --- a/trinity/common/models/vllm_patch/recording/recorder.py +++ b/trinity/common/models/vllm_patch/recording/recorder.py @@ -86,8 +86,8 @@ async def _patched_generate(*args, **kwargs): if recorder.enabled and last is not None and getattr(last, "finished", False): # Recover task id from the request's async context (set by - # SessionMiddleware). None when the client omitted X-Session-ID; - # the recorder then falls back to request_id. + # RecordingIdentityMiddleware). None when the client omitted an + # API key; the recorder then falls back to request_id. task_id = task_id_ctx.get() # Offload heavy serialization off the response critical path. asyncio.create_task(recorder.record(last, task_id)) diff --git a/trinity/common/models/vllm_patch/recording/server.py b/trinity/common/models/vllm_patch/recording/server.py index 806779d76f4..5236a4b9f3d 100644 --- a/trinity/common/models/vllm_patch/recording/server.py +++ b/trinity/common/models/vllm_patch/recording/server.py @@ -13,8 +13,7 @@ ``engine_client.generate`` to force top-k logprobs and record finished ``RequestOutput`` (covers chat/completion/responses, streaming and not). 2. ``RecordingIdentityMiddleware`` — in-process ASGI middleware reading - ``Authorization: Bearer `` (or legacy ``X-Session-ID``) into a - contextvar. + ``Authorization: Bearer `` into a contextvar. 3. ``query_router`` — ``/records/*`` endpoints for later analysis. Only for vllm versions >= 0.17.0. @@ -172,8 +171,7 @@ def _setup_recording( # the wrapped reference. Idempotent via the __patched_*__ guard. patch_engine_for_recording(engine_client, recorder, logger) - # (2) in-process middleware: API key / X-Session-ID -> contextvar. - # Zero network hop. + # (2) in-process middleware: API key -> contextvar. Zero network hop. app.add_middleware(RecordingIdentityMiddleware) # (3) query routes mounted on the main app; OpenAI /v1/* surface untouched. diff --git a/trinity/common/models/vllm_patch/recording/store.py b/trinity/common/models/vllm_patch/recording/store.py index 27ec0df0033..27d19c6f309 100644 --- a/trinity/common/models/vllm_patch/recording/store.py +++ b/trinity/common/models/vllm_patch/recording/store.py @@ -57,9 +57,8 @@ class MemoryStore(RecordStore): """In-process store. Groups experiences by recording identity (``info["task_id"]``) when an API - key or legacy X-Session-ID was supplied, otherwise each turn is keyed by its - own ``eid.suffix`` (request_id) — so a missing identity never collapses - distinct turns. + key was supplied, otherwise each turn is keyed by its own ``eid.suffix`` + (request_id) — so a missing identity never collapses distinct turns. ``get_turn`` resolves an individual turn by ``info["request_id"]``. Note: per-process under data-parallel serving — each API-server rank holds From d7b60e066de0a6f0fd2b7b4a1d2db37c2803c61d Mon Sep 17 00:00:00 2001 From: pxc Date: Thu, 25 Jun 2026 13:58:54 +0800 Subject: [PATCH 06/84] explorer: in-vLLM recording path with task_id batch join (coexists with legacy) Refactor experience production so heavy data (tokens/logprobs/routed_experts) no longer rides runner->scheduler->coordinator as serialized bytes. The vLLM recorder now captures it in-process into a MemoryStore keyed by task_id, and the coordinator pulls it at finalize time via /records/consume_task. Runners ship only a small reward map. Both paths coexist behind `explorer.use_recorded_experience` (default off = legacy). Recording module (trinity/common/models/vllm_patch/recording/): - store: drop SqlStore; MemoryStore.update_reward_by_task_id stamps reward/run/task on a whole task-id group, pops and returns it (the in-memory replacement for the SQL HistoryRecorder join). - recorder: track in-flight record tasks; add flush() (await pending + queue.join) so a consume sees a quiesced store; honor skip_recording_ctx. - models: build_experience emits one Experience per completion (n>1) with info["sample_index"]; eid.suffix=request_id kept for traceability. - context: add skip_recording_ctx; task_id already flows via api_key (RecordingIdentityMiddleware) and now also via VLLMModel.chat (Ray entry). - query: POST /records/consume_task (flush -> update_reward_by_task_id -> serialize_many); drop the SqlStore 503 branch. - config/server: remove RecordingConfig entirely; the logprob width is a recorder-internal constant (we store only the chosen token, which vLLM force-includes at logprobs=1). No static config threaded through launch. task_id propagation (Ray entry, same contextvar as the HTTP middleware): - vllm_model: chat/generate accept task_id_key, set task_id_ctx around _generate_internal; logprobs sets skip_recording_ctx (auxiliary forward). - model: ModelWrapper.chat/chat_async forward task_id_key; SGLang.chat accepts-and-ignores it (recording is vLLM-only). Coordinator + runner + workflow: - rollout_coordinator: _resolve_rank_urls (ray.get_actor per engine) and a recording-mode finalize that fans out /records/consume_task per engine, deserializes, and feeds objects to the pipeline (no re-serialization). - experience_pipeline: process_experiences(exps) public object entry. - workflow_runner: recording mode returns a pickled reward map keyed by the per-sample task_id_key the workflow stamped; legacy path unchanged. - workflow: SimpleWorkflow/AsyncSimpleWorkflow run a per-sample n=1 loop in recording mode (distinct task_id_key per sample == reward unit for GRPO), legacy n=repeat_times single-call path unchanged. - config: ExplorerConfig.use_recorded_experience flag. SQL path removal (MemoryStore only): - delete proxy/recorder.py (HistoryRecorder) and proxy_test.py; proxy service/app drop /feedback, /commit, record_feedback, submit_experiences, ready_experiences (keep allocate_model + weight sync); allocator no longer fills record_db_url; drop InferenceModelConfig.record_db_url and the dead ExplorerConfig.db_url field; RecordingConfig deleted. Serve-mode external reward reporting is intentionally left unimplemented this version (proxy /feedback//commit removed); the affected serve integration tests (TestServeWithTrainer, ServeTest) are skipped with a pointer to the recording refactor plan. convert_messages_to_experience redirect (multi-turn) is deferred with TODOs at its call sites. Co-Authored-By: Claude Fable 5 --- tests/explorer/explorer_test.py | 5 + tests/explorer/proxy_test.py | 85 ----------- tests/trainer/trainer_test.py | 5 + .../buffer/pipelines/experience_pipeline.py | 10 ++ trinity/common/config.py | 25 ++-- trinity/common/models/allocator.py | 10 +- trinity/common/models/model.py | 23 ++- trinity/common/models/sglang_model.py | 11 +- trinity/common/models/vllm_model.py | 56 +++++-- trinity/common/models/vllm_patch/__init__.py | 16 +- .../models/vllm_patch/recording/__init__.py | 11 +- .../models/vllm_patch/recording/config.py | 33 ----- .../models/vllm_patch/recording/context.py | 5 + .../models/vllm_patch/recording/models.py | 133 +++++++++-------- .../models/vllm_patch/recording/query.py | 85 +++++++++-- .../models/vllm_patch/recording/recorder.py | 100 +++++++++---- .../models/vllm_patch/recording/server.py | 72 +++------ .../models/vllm_patch/recording/store.py | 115 +++++++-------- .../workflows/envs/alfworld/RAFT_utils.py | 4 + trinity/common/workflows/workflow.py | 113 ++++++++++---- trinity/explorer/proxy/app.py | 27 ---- trinity/explorer/proxy/recorder.py | 138 ------------------ trinity/explorer/proxy/service.py | 43 ++---- trinity/explorer/rollout_coordinator.py | 111 +++++++++++++- trinity/explorer/workflow_runner.py | 33 ++++- 25 files changed, 650 insertions(+), 619 deletions(-) delete mode 100644 tests/explorer/proxy_test.py delete mode 100644 trinity/common/models/vllm_patch/recording/config.py delete mode 100644 trinity/explorer/proxy/recorder.py diff --git a/tests/explorer/explorer_test.py b/tests/explorer/explorer_test.py index a8fce8bbae3..27203d29268 100644 --- a/tests/explorer/explorer_test.py +++ b/tests/explorer/explorer_test.py @@ -438,6 +438,11 @@ def run_agent(proxy_url, model_path: str, stream: bool): return response.choices[0].message.content +@unittest.skip( + "serve-mode experience collection moved to the in-vLLM MemoryStore + " + "/records/consume_task flow; the proxy /feedback//commit path and external " + "reward reporting are being redesigned (see recording refactor plan)." +) class ServeTest(RayUnittestBaseAsync): def setUp(self): self.config = get_template_config() diff --git a/tests/explorer/proxy_test.py b/tests/explorer/proxy_test.py deleted file mode 100644 index e18ca0f3fd5..00000000000 --- a/tests/explorer/proxy_test.py +++ /dev/null @@ -1,85 +0,0 @@ -import os -import unittest -import uuid -from typing import List - -import torch - -from trinity.common.experience import EID, Experience -from trinity.explorer.proxy.recorder import HistoryRecorder - - -def get_dummy_experience(num: int) -> List[Experience]: - return [ - Experience( - eid=EID(suffix=uuid.uuid4().hex[:6]), - tokens=torch.zeros(5), - prompt_length=2, - info={ - "model_version": 0, - }, - ) - for _ in range(num) - ] - - -db_path = os.path.join(os.path.dirname(__file__), "test_recorder.db") - - -class RecorderTest(unittest.IsolatedAsyncioTestCase): - def setUp(self) -> None: - if os.path.exists(db_path): - os.remove(db_path) - - def tearDown(self) -> None: - if os.path.exists(db_path): - os.remove(db_path) - - async def test_recorder(self): - recorder = HistoryRecorder( - db_url="sqlite:///" + db_path, - table_name="experience", - ) - self.assertIsInstance(recorder, HistoryRecorder) - - experiences_1 = get_dummy_experience(3) - await recorder.record_history(experiences_1) - - msg_ids_1 = [exp.eid.suffix for exp in experiences_1] - experiences_2 = get_dummy_experience(2) - await recorder.record_history(experiences_2) - updated_experiences = await recorder.update_reward( - reward=1.0, msg_ids=msg_ids_1, run_id=1, task_id="test_task" - ) - self.assertEqual(len(updated_experiences), 3) - for exp in updated_experiences: - self.assertEqual(exp.reward, 1.0) - self.assertEqual(exp.eid.run, 1) - self.assertEqual(str(exp.eid.task), "test_task") - - updated_experiences_empty = await recorder.update_reward( - reward=2.0, msg_ids=["non_existing_id"], run_id=1, task_id="test_task" - ) - self.assertEqual(len(updated_experiences_empty), 0) - - await recorder.record_history([]) - - updated_experiences_2 = await recorder.update_reward( - reward=3.0, - msg_ids=[exp.eid.suffix for exp in experiences_2], - run_id=2, - task_id="test_task_2", - ) - self.assertEqual(len(updated_experiences_2), 2) - for exp in updated_experiences_2: - self.assertEqual(exp.reward, 3.0) - self.assertEqual(exp.eid.run, 2) - self.assertEqual(str(exp.eid.task), "test_task_2") - - updated_experiences_3 = await recorder.update_reward( - reward=4.0, - msg_ids=[exp.eid.suffix for exp in experiences_2], - run_id=3, - task_id="test_task_3", - ) - self.assertEqual(len(updated_experiences_3), 0) # already consumed diff --git a/tests/trainer/trainer_test.py b/tests/trainer/trainer_test.py index 28b2cd80d1f..ccc16b0d92d 100644 --- a/tests/trainer/trainer_test.py +++ b/tests/trainer/trainer_test.py @@ -1098,6 +1098,11 @@ async def run_math_workflow(serve_url: str, task: dict): await proxy_client.feedback_async(sum(reward.values()), [response.id]) +@unittest.skip( + "serve-mode experience collection moved to the in-vLLM MemoryStore + " + "/records/consume_task flow; the proxy /feedback//commit path and external " + "reward reporting are being redesigned (see recording refactor plan)." +) class TestServeWithTrainer(RayUnittestBaseAsync): def setUp(self): if multiprocessing.get_start_method(allow_none=True) != "spawn": diff --git a/trinity/buffer/pipelines/experience_pipeline.py b/trinity/buffer/pipelines/experience_pipeline.py index 4d436d7e144..0839c5fb4c1 100644 --- a/trinity/buffer/pipelines/experience_pipeline.py +++ b/trinity/buffer/pipelines/experience_pipeline.py @@ -180,6 +180,16 @@ async def process_serialized_chunks(self, exp_chunks: list[bytes]) -> Dict: exps.extend(Experience.deserialize_many(exp_bytes)) return await self._process_experiences(exps) + async def process_experiences(self, exps: list[Experience]) -> Dict: + """Process already-assembled experiences (objects, not serialized bytes). + + Used by the rollout coordinator's recording path, which joins reward + onto experiences pulled from the in-vLLM MemoryStore and hands them + over directly — avoiding a serialize/deserialize round-trip for the + heavy tensor payload. + """ + return await self._process_experiences(exps) + async def _process_experiences(self, exps: list[Experience]) -> Dict: st = time.time() if self.input_store is not None: diff --git a/trinity/common/config.py b/trinity/common/config.py index 1ade372c4ac..c877046651c 100644 --- a/trinity/common/config.py +++ b/trinity/common/config.py @@ -593,17 +593,12 @@ class InferenceModelConfig: # Turn on in-vLLM generation recording for the OpenAI API serving path: the # engine wraps ``engine_client.generate`` and writes each finished turn as a - # Trinity ``Experience`` to the explorer proxy's shared SQL table - # (``proxy_history``), keyed by the request id (== OpenAI ``response.id``). - # When True, the Allocator forces ``enable_return_routed_experts`` and fills - # ``record_db_url`` from ``ExplorerConfig.db_url``. VLLMModel then mirrors - # the recording config onto the engine instance for the recorder to read. - # The capture width (top-k logprobs) reuses ``logprobs`` below (default 1). + # Trinity ``Experience`` to the in-process ``MemoryStore``, keyed by the + # recording identity (``task_id``). When True, the Allocator forces + # ``enable_return_routed_experts``. VLLMModel mirrors the recording config + # onto the engine instance for the recorder to read. The capture width + # (top-k logprobs) reuses ``logprobs`` below (default 1). enable_recording: bool = False - # SQL db url the recorder writes to (shared with the explorer proxy's - # HistoryRecorder). Populated by the Allocator from ExplorerConfig.db_url - # (with a cache_dir fallback) when enable_recording is on; None otherwise. - record_db_url: Optional[str] = None # Buffer size (MB) for batched NCCL weight sync. Controls peak GPU memory during sync. weight_sync_buffer_size: int = 1024 # MB @@ -788,8 +783,6 @@ class ExplorerConfig: service_status_check_interval: int = 60 # keep at least 1 model in running status min_running_model_num: int = 1 - # db url for proxy history recorder, if not set, use proxy_history.db in buffer cache dir - db_url: Optional[str] = None # Experimental feature over_rollout: OverRolloutConfig = field(default_factory=OverRolloutConfig) @@ -800,6 +793,14 @@ class ExplorerConfig: # Maximum number of train batches that RolloutCoordinator can hold simultaneously. max_inflight_batches: int = 2 + # Use the in-vLLM recording flow: runners report only a small reward map + # (keyed by task id) and the coordinator pulls heavy experiences from each + # vLLM rank's MemoryStore via /records/consume_task at finalize time. + # Requires rollout_model.enable_recording=True and enable_openai_api=True. + # When False (default), runners ship serialized experiences through the + # scheduler as before (legacy path). + use_recorded_experience: bool = False + @dataclass class MegatronParallelConfig: diff --git a/trinity/common/models/allocator.py b/trinity/common/models/allocator.py index 0b51531c1ba..bdc202cb24b 100644 --- a/trinity/common/models/allocator.py +++ b/trinity/common/models/allocator.py @@ -83,15 +83,11 @@ async def create_engine( config.engine_id = engine_id # In-vLLM recording: force routed_experts capture (the engine reads - # this at build time — see vllm_model.py) and carry the db_url on the - # config so VLLMModel can mirror the recording config onto the engine - # instance for the recorder to read. No env / runtime_env involved. + # this at build time — see vllm_model.py). VLLMModel mirrors the + # recording config onto the engine instance for the recorder to read. + # No env / runtime_env involved. if config.enable_recording: config.enable_return_routed_experts = True - if not config.record_db_url: - config.record_db_url = self.config.db_url or ( - f"sqlite:///{self.config.buffer.cache_dir}/proxy_history.db" - ) actor_bundle_lists = [] for node_id in range(config.nnodes): diff --git a/trinity/common/models/model.py b/trinity/common/models/model.py index bc3050906cb..67976f7e675 100644 --- a/trinity/common/models/model.py +++ b/trinity/common/models/model.py @@ -284,6 +284,11 @@ async def convert_messages_to_experience( tools: Optional[List[dict]] = None, temperature: Optional[float] = None, ) -> Experience: + # TODO(recording): when the in-vLLM recorder is active, this is + # redundant — it re-tokenizes messages and runs an extra logprobs + # forward (and fakes routed_experts), all of which build_experience + # already captured at generation time into the MemoryStore. Redirect to + # a store lookup by the call's task_id_key once it's threaded here. """Convert a list of messages into an experience in async. Args: @@ -520,16 +525,26 @@ async def generate_async(self, prompts: List[str], **kwargs) -> List[Experience] return [exp for exps in results for exp in exps] @_history_recorder - def chat(self, messages: List[dict], **kwargs) -> List[Experience]: + def chat( + self, messages: List[dict], task_id_key: Optional[str] = None, **kwargs + ) -> List[Experience]: """Generate a list of experiences from a list of messages.""" lora_request = self.get_lora_request() - return ray.get(self.model.chat.remote(messages, lora_request=lora_request, **kwargs)) + return ray.get( + self.model.chat.remote( + messages, lora_request=lora_request, task_id_key=task_id_key, **kwargs + ) + ) @_history_recorder - async def chat_async(self, messages: List[dict], **kwargs) -> List[Experience]: + async def chat_async( + self, messages: List[dict], task_id_key: Optional[str] = None, **kwargs + ) -> List[Experience]: """Generate a list of experiences from a list of messages in async.""" lora_request = await self.get_lora_request_async() - return await self.model.chat.remote(messages, lora_request=lora_request, **kwargs) + return await self.model.chat.remote( + messages, lora_request=lora_request, task_id_key=task_id_key, **kwargs + ) def logprobs(self, tokens: List[int], temperature: Optional[float] = None) -> Tensor: """Calculate the logprobs of the given tokens.""" diff --git a/trinity/common/models/sglang_model.py b/trinity/common/models/sglang_model.py index a69c88de2f6..49263d72da3 100644 --- a/trinity/common/models/sglang_model.py +++ b/trinity/common/models/sglang_model.py @@ -434,7 +434,16 @@ async def generate(self, prompt: str, lora_request=None, **kwargs) -> Sequence[E ) return experiences - async def chat(self, messages: List[dict], lora_request=None, **kwargs) -> Sequence[Experience]: + async def chat( + self, + messages: List[dict], + lora_request=None, + task_id_key: Optional[str] = None, + **kwargs, + ) -> Sequence[Experience]: + # ``task_id_key`` is the in-vLLM recorder's group key; SGLang does not + # support recording, so accept and ignore it here. + del task_id_key if self.tokenizer is None: await self._initialize_tokenizer() diff --git a/trinity/common/models/vllm_model.py b/trinity/common/models/vllm_model.py index 4b133c3ac64..3b32d8bc4d5 100644 --- a/trinity/common/models/vllm_model.py +++ b/trinity/common/models/vllm_model.py @@ -19,6 +19,10 @@ ) from trinity.common.models.model import BaseInferenceModel from trinity.common.models.vllm_patch import get_vllm_version +from trinity.common.models.vllm_patch.recording.context import ( + skip_recording_ctx, + task_id_ctx, +) # V0 engine is deprecated since vLLM v0.10.2, related code will be removed in the future. @@ -202,11 +206,22 @@ async def prepare(self) -> None: self.headless_executor.start_worker_monitor() self._prepared = True - async def chat(self, messages: List[Dict], lora_request=None, **kwargs) -> Sequence[Experience]: + async def chat( + self, + messages: List[Dict], + lora_request=None, + task_id_key: Optional[str] = None, + **kwargs, + ) -> Sequence[Experience]: """Chat with the model with a list of messages in async. Args: messages (List[dict]): The input history messages. + task_id_key (Optional[str]): Recording identity for the in-vLLM + recorder (the MemoryStore group key). Propagated to + ``generate`` via ``task_id_ctx`` so the recorder stamps it into + ``info["task_id"]`` without an HTTP hop. None leaves it to the + engine request id (fallback grouping). kwargs (dict): A dictionary of sampling parameters. Returns: @@ -232,7 +247,9 @@ async def chat(self, messages: List[Dict], lora_request=None, **kwargs) -> Seque "prompt": prompt, "multi_modal_data": multi_modal_data or {}, } - return await self.generate(prompt=prompt, lora_request=lora_request, **kwargs) + return await self.generate( + prompt=prompt, lora_request=lora_request, task_id_key=task_id_key, **kwargs + ) def _extract_routed_experts(self, output: Any, output_index: int) -> Optional[torch.Tensor]: if not self.config.enable_return_routed_experts: @@ -256,12 +273,18 @@ def _extract_routed_experts(self, output: Any, output_index: int) -> Optional[to return torch.cat(routed_experts_parts, dim=0) async def generate( - self, prompt: Union[str, Dict], lora_request=None, **kwargs + self, + prompt: Union[str, Dict], + lora_request=None, + task_id_key: Optional[str] = None, + **kwargs, ) -> Sequence[Experience]: """Generate a response from the provided prompt in async. Args: prompt (str): The input prompt. + task_id_key (Optional[str]): Recording identity propagated to the + in-vLLM recorder via ``task_id_ctx`` (see ``chat``). kwargs (dict): A dictionary of sampling parameters. Returns: @@ -282,7 +305,16 @@ async def generate( } # is_valid is True: returned_seq is token_ids multi_modal_inputs = None - output = await self._generate_internal(prompt=prompt, lora_request=lora_request, **kwargs) + # Propagate the recording identity to the engine-level recorder (same + # async task, same process) so the recorded experience is grouped under + # this task id in the MemoryStore. + task_id_token = task_id_ctx.set(task_id_key) + try: + output = await self._generate_internal( + prompt=prompt, lora_request=lora_request, **kwargs + ) + finally: + task_id_ctx.reset(task_id_token) if is_mm_prompt: if self.mm_render is None: self.mm_render = vLLMMultiModalRender( @@ -353,11 +385,17 @@ async def logprobs( # type: ignore [override] # avoid using prefix cache when calculating logprobs, only for vLLM >= 0.12.0 if self.logprobs_no_prefix_cache: kwargs["skip_reading_prefix_cache"] = True - output = await self._generate_internal( - prompt={"prompt_token_ids": token_ids}, - lora_request=lora_request, - **kwargs, - ) + # This is an auxiliary 1-token forward, not a real turn — keep it out + # of the recording store so it doesn't pollute task-id groups. + skip_token = skip_recording_ctx.set(True) + try: + output = await self._generate_internal( + prompt={"prompt_token_ids": token_ids}, + lora_request=lora_request, + **kwargs, + ) + finally: + skip_recording_ctx.reset(skip_token) return torch.tensor( [list(logprob_dict.values())[0].logprob for logprob_dict in output.prompt_logprobs[1:]], dtype=torch.float32, diff --git a/trinity/common/models/vllm_patch/__init__.py b/trinity/common/models/vllm_patch/__init__.py index 692ae4e1cc7..cf54d3dec2d 100644 --- a/trinity/common/models/vllm_patch/__init__.py +++ b/trinity/common/models/vllm_patch/__init__.py @@ -154,8 +154,10 @@ def get_api_server( ): vllm_version = get_vllm_version() # Recording is driven by the config field (not env, not an engine attr): - # get_api_server already receives InferenceModelConfig, so it builds the - # static RecordingConfig here and threads it into the recording runner. + # when on, ``get_api_server`` selects the recording-enabled runner, which + # wires the in-process MemoryStore + engine wrap itself — no static config + # needs threading (the logprob width is a recorder-internal constant, and + # the checkpoint version is read live off the engine). recording = bool(config.enable_recording) and vllm_version >= VLLM_VERSION_0230 if config.enable_recording and not recording: logger.warning( @@ -176,15 +178,7 @@ def get_api_server( enable_log_requests=config.enable_log_requests, chat_template=config.chat_template, ) - if recording: - from trinity.common.models.vllm_patch.recording import RecordingConfig - - kwargs["recording_config"] = RecordingConfig( - db_url=config.record_db_url, - table="proxy_history", - topk=config.logprobs if config.logprobs else 1, - ) # The dynamic checkpoint version is read live off the engine instance # (``async_llm.trinity_model_version``, mirrored by VLLMModel), so it is - # not part of the static config passed here. + # not part of any static config passed here. return asyncio.create_task(run_api_server_in_ray_actor(async_llm, **kwargs)) diff --git a/trinity/common/models/vllm_patch/recording/__init__.py b/trinity/common/models/vllm_patch/recording/__init__.py index adb525a172c..a10f2f800f9 100644 --- a/trinity/common/models/vllm_patch/recording/__init__.py +++ b/trinity/common/models/vllm_patch/recording/__init__.py @@ -1,9 +1,9 @@ """Generation-recording patch for the vLLM OpenAI server. Designed for vllm >= 0.23.0. """ -from trinity.common.models.vllm_patch.recording.config import RecordingConfig from trinity.common.models.vllm_patch.recording.context import ( RecordingIdentityMiddleware, + skip_recording_ctx, task_id_ctx, ) from trinity.common.models.vllm_patch.recording.models import build_experience @@ -15,22 +15,17 @@ from trinity.common.models.vllm_patch.recording.server import ( run_api_server_with_recording, ) -from trinity.common.models.vllm_patch.recording.store import ( - MemoryStore, - RecordStore, - SqlStore, -) +from trinity.common.models.vllm_patch.recording.store import MemoryStore, RecordStore __all__ = [ "MemoryStore", "RecordStore", - "RecordingConfig", "RecordingIdentityMiddleware", - "SqlStore", "Recorder", "build_experience", "patch_engine_for_recording", "query_router", "run_api_server_with_recording", + "skip_recording_ctx", "task_id_ctx", ] diff --git a/trinity/common/models/vllm_patch/recording/config.py b/trinity/common/models/vllm_patch/recording/config.py deleted file mode 100644 index 30bdcd08acc..00000000000 --- a/trinity/common/models/vllm_patch/recording/config.py +++ /dev/null @@ -1,33 +0,0 @@ -# -*- coding: utf-8 -*- -"""Recording configuration carried explicitly through the launch chain. - -Instead of env vars or attributes bolted onto the engine instance, the static -recording config (db url / table / top-k) is passed as an explicit parameter -into ``run_api_server_with_recording``. ``get_api_server`` builds one from -``InferenceModelConfig`` when ``enable_recording`` is on. - -The *dynamic* checkpoint version is NOT here: it changes at runtime -(``sync_model_weights``), so the recorder reads it live off the engine -instance attribute ``trinity_model_version`` (mirrored by VLLMModel). -""" -from dataclasses import dataclass -from typing import Optional - - -@dataclass -class RecordingConfig: - """Static configuration for the in-vLLM recorder. - - Attributes: - db_url: SQL db url the recorder writes to (shared with the explorer - proxy's HistoryRecorder). When None, the recorder falls back to an - in-process MemoryStore (no cross-process visibility). - table: SQL table name (default ``proxy_history``, matching the proxy). - topk: How many top-k logprobs the engine computes per generated token. - Only the chosen token's logprob is stored, and vLLM force-includes - the sampled token, so 1 suffices by default. - """ - - db_url: Optional[str] = None - table: str = "proxy_history" - topk: int = 1 diff --git a/trinity/common/models/vllm_patch/recording/context.py b/trinity/common/models/vllm_patch/recording/context.py index 7244c15c987..9aa6cf1cc60 100644 --- a/trinity/common/models/vllm_patch/recording/context.py +++ b/trinity/common/models/vllm_patch/recording/context.py @@ -20,6 +20,11 @@ # task id). task_id_ctx: ContextVar[Optional[str]] = ContextVar("trinity_recording_task_id", default=None) +# Set around auxiliary engine.generate calls (logprobs recomputation, +# convert_messages_to_experience) so the recorder skips them — those 1-token +# forwards are not real turns and would pollute the store. +skip_recording_ctx: ContextVar[bool] = ContextVar("trinity_recording_skip", default=False) + #: Preferred identity header for OpenAI-compatible clients. AUTHORIZATION_HEADER = "authorization" diff --git a/trinity/common/models/vllm_patch/recording/models.py b/trinity/common/models/vllm_patch/recording/models.py index 2da1b398224..30cf0535105 100644 --- a/trinity/common/models/vllm_patch/recording/models.py +++ b/trinity/common/models/vllm_patch/recording/models.py @@ -1,18 +1,25 @@ # -*- coding: utf-8 -*- -"""Build a Trinity ``Experience`` from a finished vLLM ``RequestOutput``. +"""Build Trinity ``Experience`` objects from a finished vLLM ``RequestOutput``. We record into Trinity's native ``Experience`` struct (see ``trinity.common.experience``) rather than a bespoke record, so captured data drops straight into Trinity's RL/buffer pipeline without a conversion step. +A single ``RequestOutput`` may carry multiple completions (``n > 1``); we emit +one ``Experience`` per completion so no sample is lost. + Field mapping (captured ``RequestOutput`` fields -> ``Experience``): - request_id -> eid.suffix (``EID(suffix=...)``; this is the msg_id the - proxy/openai client sees as ``response.id`` — the key the - proxy's ``HistoryRecorder.update_reward`` looks up by). + request_id -> eid.suffix (``EID(suffix=...)``; the vLLM engine request + id == the OpenAI ``response.id``. Kept for traceability; ``eid.task``/``run``/``reward`` are left default here and - assigned later by the proxy at ``/feedback`` time, matching - ``explorer/proxy/service.record_experience`` semantics. - API key -> info["task_id"] (traceability only; not used as a key) + assigned by ``MemoryStore.update_reward_by_task_id`` at + consume time.) + API key / task id -> info["task_id"] (the recording identity; **the group + key** the MemoryStore batches experiences by, so a whole + task's samples/turns are reward-updated and consumed + together. Falls back to ``eid.suffix`` when absent.) + sample index -> info["sample_index"] (position within the n-completion + set; orders samples/turns inside a task-id group) prompt_token_ids -> tokens (prompt portion) + prompt_length response_token_ids-> tokens (response portion) logprobs -> Experience.logprobs -- but ONLY the *chosen* token's @@ -24,11 +31,11 @@ model_version -> info["model_version"] (which checkpoint policy served the turn; read in-actor by the recorder's provider) -Plus bookkeeping (request_id / task_id / rank / timestamp / endpoint / -model_version) stashed in ``Experience.info`` so it round-trips with the -experience through serialize/deserialize. +Plus bookkeeping (request_id / task_id / sample_index / rank / timestamp / +endpoint / model_version) stashed in ``Experience.info`` so it round-trips +with the experience through serialize/deserialize. """ -from typing import Any, Optional +from typing import Any, List, Optional from trinity.common.experience import EID, Experience @@ -79,15 +86,18 @@ def build_experience( timestamp: str, endpoint: str = "unknown", model_version: Optional[int] = None, -) -> Optional[Experience]: - """Build a Trinity ``Experience`` from a finished ``RequestOutput``. +) -> List[Experience]: + """Build Trinity ``Experience`` objects from a finished ``RequestOutput``. + + One experience per completion (``output.outputs``), so ``n > 1`` sampling + is captured in full. Each experience shares ``eid.suffix = request_id`` and + ``info["task_id"] = task_id`` (the group key); ``info["sample_index"]`` + distinguishes samples within the group. Args: output: A ``RequestOutput`` with ``finished == True``. - task_id: From the request API key; stored in ``info`` for traceability - only. - Not used as the storage key — ``eid.suffix`` is, so a missing - API key never drops a turn. + task_id: The recording identity (API key / Ray-injected task id); + stored in ``info["task_id"]`` and used as the MemoryStore group key. rank: Data-parallel serving rank. timestamp: UTC ISO-8601 string (caller-stamped to keep this pure). endpoint: Which OpenAI endpoint served the turn (best-effort). @@ -95,48 +105,55 @@ def build_experience( into ``info`` for RL attribution (read in-actor by the recorder). Returns: - A populated ``Experience``, or None if the turn is degenerate (no - prompt or no response tokens) and cannot form a valid experience. + One ``Experience`` per non-degenerate completion. Empty list if the + request had no prompt or no completion with response tokens. """ request_id = output.request_id - # Key by the request id (= the OpenAI response ``id`` / proxy msg_id) so the - # proxy's HistoryRecorder.update_reward can find this row at feedback time. - # task/run/reward are intentionally left default — the proxy assigns them. + # eid.suffix = request_id for traceability; task/run/reward are left + # default and assigned by MemoryStore.update_reward_by_task_id at consume. prompt_token_ids = list(output.prompt_token_ids or []) - - completion = output.outputs[0] if output.outputs else None - if completion is None: - return None - response_token_ids = list(completion.token_ids or []) - - # A valid single-turn experience needs both a prompt and a response; - # Experience.__init__ asserts len(tokens) > prompt_length otherwise. - if not prompt_token_ids or not response_token_ids: - return None - - tokens = prompt_token_ids + response_token_ids - prompt_length = len(prompt_token_ids) - - chosen_logprobs = _extract_chosen_logprobs(completion.logprobs, response_token_ids) - routed_experts = completion.routed_experts - - info = { - "request_id": request_id, - "task_id": task_id, - "rank": rank, - "timestamp": timestamp, - "endpoint": endpoint, - "model_version": model_version, - } - - return Experience( - eid=EID(suffix=request_id), - tokens=tokens, - logprobs=chosen_logprobs, - prompt_length=prompt_length, - routed_experts=routed_experts, - prompt_text=output.prompt, - response_text=completion.text, - info=info, - ) + if not prompt_token_ids: + return [] + + completions = list(output.outputs or []) + if not completions: + return [] + + experiences: List[Experience] = [] + for sample_index, completion in enumerate(completions): + response_token_ids = list(completion.token_ids or []) + # A valid single-turn experience needs both a prompt and a response; + # Experience.__init__ asserts len(tokens) > prompt_length otherwise. + if not response_token_ids: + continue + + tokens = prompt_token_ids + response_token_ids + prompt_length = len(prompt_token_ids) + + chosen_logprobs = _extract_chosen_logprobs(completion.logprobs, response_token_ids) + routed_experts = completion.routed_experts + + info = { + "request_id": request_id, + "task_id": task_id, + "sample_index": sample_index, + "rank": rank, + "timestamp": timestamp, + "endpoint": endpoint, + "model_version": model_version, + } + + experiences.append( + Experience( + eid=EID(suffix=request_id), + tokens=tokens, + logprobs=chosen_logprobs, + prompt_length=prompt_length, + routed_experts=routed_experts, + prompt_text=output.prompt, + response_text=completion.text, + info=info, + ) + ) + return experiences diff --git a/trinity/common/models/vllm_patch/recording/query.py b/trinity/common/models/vllm_patch/recording/query.py index 7f53783b294..3629adf2275 100644 --- a/trinity/common/models/vllm_patch/recording/query.py +++ b/trinity/common/models/vllm_patch/recording/query.py @@ -1,4 +1,4 @@ -"""HTTP query endpoints over recorded experiences. +"""HTTP query/consume endpoints over recorded experiences. Mounted on the main vLLM app (we own ``app`` in the bootstrap) via ``app.include_router(query_router)``. Routes are prefixed under ``/records`` to @@ -9,33 +9,55 @@ tensor payload). A single-turn endpoint returns the full experience as Trinity serialized bytes (``Experience.serialize()``) so it can be fed straight back into Trinity pipelines. + +``POST /records/consume_task`` is the consume path: it drains the recorder, +reward-stamps every experience in the matching task-id groups, pops them, and +returns the heavy experiences as ``Experience.serialize_many`` bytes — the +coordinator fans this out per rank at finalize time. """ +from typing import List + from fastapi import APIRouter, HTTPException, Request, Response +from pydantic import BaseModel from trinity.common.experience import Experience +from trinity.common.models.vllm_patch.recording.recorder import Recorder from trinity.common.models.vllm_patch.recording.store import RecordStore #: Mounted on app.state by the bootstrap wiring in ``server.py``. _STORE_ATTR = "trinity_record_store" +_RECORDER_ATTR = "trinity_recorder" query_router = APIRouter(prefix="/records", tags=["trinity-recording"]) -def _store(request: Request, *, read: bool = False) -> RecordStore: +class _TaskRewardUpdate(BaseModel): + """One task-id group to reward-stamp and consume.""" + + task_id: str + reward: float + run: int = 0 + task: str = "" + + +class _ConsumeTaskRequest(BaseModel): + updates: List[_TaskRewardUpdate] + + +def _store(request: Request) -> RecordStore: store = getattr(request.app.state, _STORE_ATTR, None) if store is None: raise HTTPException(status_code=503, detail="recording store not configured") - if read and not getattr(store, "supports_reads", True): - # SqlStore shares the explorer proxy's table; reads/consumes are served - # by the proxy's /feedback + /commit, not by these debug endpoints. - raise HTTPException( - status_code=503, - detail="this record store backend is read-only from the vLLM " - "side; query via the explorer proxy instead", - ) return store +def _recorder(request: Request) -> Recorder: + rec = getattr(request.app.state, _RECORDER_ATTR, None) + if rec is None: + raise HTTPException(status_code=503, detail="recorder not configured") + return rec + + async def _get_exp(store: RecordStore, task_id: str, request_id: str) -> Experience: exp = await store.get_turn(task_id, request_id) if exp is None: @@ -46,7 +68,7 @@ async def _get_exp(store: RecordStore, task_id: str, request_id: str) -> Experie @query_router.get("/tasks") async def list_tasks(request: Request) -> dict: """List all known task ids.""" - store = _store(request, read=True) + store = _store(request) return {"task_ids": await store.list_tasks()} @@ -57,7 +79,7 @@ async def get_task(task_id: str, request: Request) -> dict: Tensor payloads (tokens/logprobs/routed_experts) are omitted here to keep listing responses small; fetch the per-turn blob endpoint for full data. """ - store = _store(request, read=True) + store = _store(request) experiences = await store.get_task(task_id) return {"task_id": task_id, "turns": [e.to_dict() for e in experiences]} @@ -65,7 +87,7 @@ async def get_task(task_id: str, request: Request) -> dict: @query_router.get("/tasks/{task_id}/turns/{request_id}") async def get_turn(task_id: str, request_id: str, request: Request) -> dict: """Return lightweight metadata for a single experience.""" - store = _store(request, read=True) + store = _store(request) exp = await _get_exp(store, task_id, request_id) return exp.to_dict() @@ -77,7 +99,7 @@ async def get_turn_blob(task_id: str, request_id: str, request: Request) -> Resp Equivalent to ``Experience.serialize()``; deserializable via ``Experience.deserialize(...)`` or ``deserialize_many``. """ - store = _store(request, read=True) + store = _store(request) exp = await _get_exp(store, task_id, request_id) return Response( content=Experience.serialize(exp), @@ -89,6 +111,39 @@ async def get_turn_blob(task_id: str, request_id: str, request: Request) -> Resp @query_router.delete("/tasks/{task_id}") async def delete_task(task_id: str, request: Request) -> dict: """Delete all experiences for a task.""" - store = _store(request, read=True) + store = _store(request) await store.delete_task(task_id) return {"task_id": task_id, "deleted": True} + + +@query_router.post("/consume_task") +async def consume_task(req: _ConsumeTaskRequest, request: Request) -> Response: + """Reward-stamp and pop the experiences of the given task-id groups. + + Drains the recorder first (so every finished turn already made it into the + store), then for each update sets ``reward``/``run``/``task`` on the whole + matching group and pops it. Returns the union as Trinity serialized bytes + (``Experience.serialize_many``), ready for the coordinator pipeline. + + A task_id absent from this rank yields no experiences (it lived on another + rank); the coordinator fans this call out to every rank and merges. + """ + store = _store(request) + recorder = _recorder(request) + # Ensure in-flight record tasks have been appended before we pop. + await recorder.flush() + + exps: List[Experience] = [] + for update in req.updates: + exps.extend( + await store.update_reward_by_task_id( + task_id=update.task_id, + reward=update.reward, + run=update.run, + task=update.task, + ) + ) + return Response( + content=Experience.serialize_many(exps), + media_type="application/octet-stream", + ) diff --git a/trinity/common/models/vllm_patch/recording/recorder.py b/trinity/common/models/vllm_patch/recording/recorder.py index 106e104e8b5..942af6fe4ca 100644 --- a/trinity/common/models/vllm_patch/recording/recorder.py +++ b/trinity/common/models/vllm_patch/recording/recorder.py @@ -23,7 +23,10 @@ from typing import Optional from trinity.common.experience import Experience -from trinity.common.models.vllm_patch.recording.context import task_id_ctx +from trinity.common.models.vllm_patch.recording.context import ( + skip_recording_ctx, + task_id_ctx, +) from trinity.common.models.vllm_patch.recording.models import build_experience from trinity.common.models.vllm_patch.recording.store import RecordStore @@ -34,6 +37,13 @@ #: engine creation); read live here so each experience is attributed to the #: right policy without a launch-time parameter. _MODEL_VERSION_ATTR = "trinity_model_version" +#: Force at least this many top-k logprobs per generated token so recording +#: captures the chosen token's logprob even when the caller didn't request +#: logprobs. We store ONLY the sampled token's logprob, and vLLM force-includes +#: the sampled token at ``logprobs=1``, so 1 is the only useful value — no need +#: to thread a knob through the launcher. The engine's ``max_logprobs`` cap +#: (default 20, set at engine build) already covers it. +_RECORDER_LOGPROB_WIDTH = 1 def patch_engine_for_recording( @@ -73,9 +83,12 @@ async def _patched_generate(*args, **kwargs): sampling_params = args[1] if recorder.enabled and sampling_params is not None: - desired = recorder.topk + # Ensure logprobs are computed for recording (callers may omit + # them, e.g. on the HTTP path). See _RECORDER_LOGPROB_WIDTH. cur = sampling_params.logprobs - sampling_params.logprobs = max(cur, desired) if cur is not None else desired + sampling_params.logprobs = ( + max(cur, _RECORDER_LOGPROB_WIDTH) if cur is not None else _RECORDER_LOGPROB_WIDTH + ) last = None # ``current`` is the original *bound* method captured pre-wrap, so it @@ -86,11 +99,13 @@ async def _patched_generate(*args, **kwargs): if recorder.enabled and last is not None and getattr(last, "finished", False): # Recover task id from the request's async context (set by - # RecordingIdentityMiddleware). None when the client omitted an - # API key; the recorder then falls back to request_id. + # RecordingIdentityMiddleware on the HTTP path, or by + # VLLMModel.chat on the Ray-direct path). None when neither was + # supplied; the store then falls back to request_id grouping. task_id = task_id_ctx.get() - # Offload heavy serialization off the response critical path. - asyncio.create_task(recorder.record(last, task_id)) + # Offload heavy serialization off the response critical path. The + # task is tracked so ``flush`` can await it before a consume. + recorder.schedule_record(last, task_id) setattr(_patched_generate, _PATCHED_FLAG, True) engine_client.generate = _patched_generate @@ -101,21 +116,23 @@ class Recorder: """Drains finished turns into a ``RecordStore`` from a background task. Putting records into an ``asyncio.Queue`` and flushing from a single worker - keeps the response path cheap (record == one ``queue.put``) and serializes - expensive payloads (ndarray -> .npy, json) off the serving hot loop. + keeps the response path cheap (record == enqueue) and serializes expensive + payloads (ndarray -> .npy, json) off the serving hot loop. + + ``schedule_record`` spawns a task per finished ``RequestOutput`` and tracks + it in ``_pending``; ``flush`` awaits all of them plus ``queue.join`` so a + caller (``/records/consume_task``) sees a quiesced store before popping. """ def __init__( self, store: RecordStore, *, - topk: int, enabled: bool, rank: int = 0, engine_client=None, ) -> None: self.store = store - self.topk = topk self.enabled = enabled self.rank = rank # The engine_client is the same AsyncLLM instance VLLMModel updates in @@ -124,6 +141,11 @@ def __init__( self.engine_client = engine_client self._queue: "asyncio.Queue[Optional[Experience]]" = asyncio.Queue() self._flusher: Optional[asyncio.Task] = None + # In-flight ``_record`` tasks spawned by ``schedule_record``. Tracked + # so ``flush`` can await them — a record task that hasn't put yet would + # otherwise race a consume (the experience would be missing from the + # store even though the request already finished). + self._pending: "set[asyncio.Task]" = set() def start(self) -> None: """Start the background flusher. Idempotent.""" @@ -132,49 +154,65 @@ def start(self) -> None: self._flusher = asyncio.create_task(self._flush_loop()) async def stop(self) -> None: - """Cancel the flusher and drain remaining queued turns.""" + """Drain in-flight + queued turns, then stop the flusher.""" if self._flusher is None: return + # Drain everything put in flight so we don't lose turns at shutdown. + await self.flush() self._flusher.cancel() self._flusher = None - # Drain anything already queued so we don't lose in-flight turns. - while not self._queue.empty(): - exp = self._queue.get_nowait() - if exp is not None: - await self._safe_append(exp) - async def record(self, output, task_id: Optional[str]) -> None: - """Enqueue a finished ``RequestOutput`` for recording as an Experience. + def schedule_record(self, output, task_id: Optional[str]) -> None: + """Spawn (and track) a record task for a finished ``RequestOutput``.""" + task = asyncio.create_task(self._record(output, task_id)) + self._pending.add(task) + task.add_done_callback(self._pending.discard) + + async def flush(self) -> None: + """Wait until every in-flight record has been appended to the store. - Args: - output: A finished ``RequestOutput``. - task_id: From ``task_id_ctx``; stored in ``info`` for traceability. + Awaits all pending ``_record`` tasks (so every finished turn has been + enqueued), then ``queue.join``s (so the flusher has appended them all). + Call this before consuming the store to avoid reading a partial state. """ + if self._pending: + await asyncio.gather(*self._pending, return_exceptions=True) + if self._flusher is not None: + await self._queue.join() + + async def _record(self, output, task_id: Optional[str]) -> None: + """Build experiences for a finished turn and enqueue each for append.""" + # Auxiliary forwards (logprobs recomputation, convert_messages) set + # this to avoid polluting the store with 1-token degenerate turns. + if skip_recording_ctx.get(): + return # Stamp now (real runtime, not a workflow sandbox): permitted here. from datetime import datetime, timezone timestamp = datetime.now(timezone.utc).isoformat() # Read the live checkpoint version the actor mirrors onto the engine. model_version = getattr(self.engine_client, _MODEL_VERSION_ATTR, None) - exp = build_experience( + exps = build_experience( output, task_id, rank=self.rank, timestamp=timestamp, model_version=model_version, ) - if exp is None: - # Degenerate turn (no prompt/response) — nothing to record. - return - await self._queue.put(exp) + for exp in exps: + await self._queue.put(exp) async def _flush_loop(self) -> None: while True: exp = await self._queue.get() - if exp is None: - # Sentinel for graceful shutdown. - return - await self._safe_append(exp) + try: + if exp is None: + # Sentinel for graceful shutdown. + return + await self._safe_append(exp) + finally: + # Paired with put() so queue.join() in flush() can complete. + self._queue.task_done() async def _safe_append(self, exp: Experience) -> None: try: diff --git a/trinity/common/models/vllm_patch/recording/server.py b/trinity/common/models/vllm_patch/recording/server.py index 5236a4b9f3d..a9c304668a7 100644 --- a/trinity/common/models/vllm_patch/recording/server.py +++ b/trinity/common/models/vllm_patch/recording/server.py @@ -44,7 +44,6 @@ from vllm.version import __version__ as VLLM_VERSION from trinity.common.models.vllm_patch import get_vllm_version -from trinity.common.models.vllm_patch.recording.config import RecordingConfig from trinity.common.models.vllm_patch.recording.context import ( RecordingIdentityMiddleware, ) @@ -53,11 +52,7 @@ Recorder, patch_engine_for_recording, ) -from trinity.common.models.vllm_patch.recording.store import ( - MemoryStore, - RecordStore, - SqlStore, -) +from trinity.common.models.vllm_patch.recording.store import MemoryStore, RecordStore #: Attribute on app.state holding the active RecordStore. _STORE_STATE_ATTR = "trinity_record_store" @@ -117,39 +112,29 @@ def _setup_recording( engine_client, app, logger, - recording_config: Optional[RecordingConfig] = None, -) -> Optional[Recorder]: +) -> Recorder: """Wire generation recording onto the in-construction server. - Returns the started Recorder (for lifecycle management), or None if - recording is disabled (``recording_config`` is None). + Returns the started Recorder (for lifecycle management). This is only + called when recording is on, so there is no disable switch here. + + No static config is threaded in: the chosen-token logprob width is a + constant inside the recorder (``_RECORDER_LOGPROB_WIDTH`` — we store only + the sampled token's logprob, so 1 suffices). The *dynamic* checkpoint + version is read live off ``engine_client.trinity_model_version`` (mirrored + by VLLMModel at engine creation and in ``sync_model_weights``). - The static config (db_url/table/topk) arrives explicitly via - ``recording_config`` (built by ``get_api_server`` from - ``InferenceModelConfig``). The *dynamic* checkpoint version is read live - off ``engine_client.trinity_model_version`` (mirrored by VLLMModel at - engine creation and in ``sync_model_weights``). + The store backend is always the in-process ``MemoryStore``; the coordinator + drains it at finalize time via ``/records/consume_task`` (fanned out per + rank), so heavy experience bytes never touch SQL or Ray serialization. Args: args: Parsed vLLM CLI args. engine_client: AsyncLLM instance (we own it pre-init_app_state). app: FastAPI app from ``build_app`` (we own it pre-serve_http). logger: Logger. - recording_config: Static recording config; None disables recording. """ - if recording_config is None: - return None - - if recording_config.db_url: - store: RecordStore = SqlStore( - db_url=recording_config.db_url, table_name=recording_config.table - ) - else: - logger.warning( - "recording enabled but recording_config.db_url is None; falling " - "back to in-process MemoryStore (no cross-process visibility)" - ) - store = MemoryStore() + store: RecordStore = MemoryStore() # Rank is constant per process; capture once (RequestOutput does not expose # parallel_config, so we read it from engine_client here, mirroring @@ -161,7 +146,6 @@ def _setup_recording( recorder = Recorder( store=store, - topk=recording_config.topk, enabled=True, rank=rank, engine_client=engine_client, @@ -181,8 +165,7 @@ def _setup_recording( setattr(app.state, _RECORDER_STATE_ATTR, recorder) logger.info( - "Generation recording enabled: topk=%d store=%s rank=%d", - recording_config.topk, + "Generation recording enabled: store=%s rank=%d", type(store).__name__, rank, ) @@ -195,7 +178,6 @@ async def run_server_worker_in_ray( args, engine_client, logger, - recording_config: Optional[RecordingConfig] = None, ) -> None: """Modified from vllm.entrypoints.openai.api_server.run_server_worker. @@ -212,7 +194,7 @@ async def run_server_worker_in_ray( app = build_app(args) # --- recording wiring: engine wrap must precede init_app_state ----------- - recorder = _setup_recording(args, engine_client, app, logger, recording_config=recording_config) + recorder = _setup_recording(args, engine_client, app, logger) # ------------------------------------------------------------------------ await init_app_state(engine_client, app.state, args) @@ -226,8 +208,7 @@ async def run_server_worker_in_ray( listen_address, ) - if recorder is not None: - recorder.start() + recorder.start() shutdown_task = await serve_http( app, @@ -252,8 +233,7 @@ async def run_server_worker_in_ray( try: await shutdown_task finally: - if recorder is not None: - await recorder.stop() + await recorder.stop() sock.close() @@ -261,14 +241,11 @@ async def run_server_in_ray( args, engine_client, logger, - recording_config: Optional[RecordingConfig] = None, ): # Modified from vllm.entrypoints.openai.api_server.run_server listen_address, sock = setup_server_in_ray(args, logger) logger.info("vLLM API server listening on %s", listen_address) - await run_server_worker_in_ray( - listen_address, sock, args, engine_client, logger, recording_config - ) + await run_server_worker_in_ray(listen_address, sock, args, engine_client, logger) async def run_api_server_with_recording( @@ -282,17 +259,14 @@ async def run_api_server_with_recording( tool_call_parser: Optional[str] = None, reasoning_parser: Optional[str] = None, enable_log_requests: bool = False, - recording_config: Optional[RecordingConfig] = None, ): """Drop-in recording-enabled variant of ``api_patch_v17.run_api_server_in_ray_actor_v17``. - Same signature plus an optional ``recording_config`` so launchers can - switch by import path. Requires vllm >= 0.17.0. Recording is on iff - ``recording_config`` is provided (built by ``get_api_server`` from - ``InferenceModelConfig`` when ``enable_recording`` is on). The dynamic + Requires vllm >= 0.17.0. No static recording config is threaded in: the + logprob capture width is a recorder-internal constant. The dynamic checkpoint version is read off ``async_llm.trinity_model_version`` - (mirrored by VLLMModel), so it is not part of the static config here. + (mirrored by VLLMModel). """ vllm_version = get_vllm_version() if vllm_version < parse_version("0.17.0"): @@ -338,4 +312,4 @@ async def run_api_server_with_recording( args = parser.parse_args(cli_args) args.structured_outputs_config.reasoning_parser = reasoning_parser logger.info(f"Starting vLLM OpenAI API server with args: {args}") - await run_server_in_ray(args, async_llm, logger, recording_config) + await run_server_in_ray(args, async_llm, logger) diff --git a/trinity/common/models/vllm_patch/recording/store.py b/trinity/common/models/vllm_patch/recording/store.py index 27d19c6f309..f745aa77459 100644 --- a/trinity/common/models/vllm_patch/recording/store.py +++ b/trinity/common/models/vllm_patch/recording/store.py @@ -1,21 +1,26 @@ """Pluggable storage backends for recorded experiences. -A ``RecordStore`` persists Trinity ``Experience`` objects. Backends: - -* ``MemoryStore`` — in-process, request/session-keyed; for standalone runs and - the ``/records`` debug endpoints. -* ``SqlStore`` — delegates to ``trinity.explorer.proxy.recorder.HistoryRecorder`` - so the in-vLLM recorder writes to the *same* SQL table the explorer proxy - reads (``proxy_history``). This is the online-RL path: experiences written - here by the vLLM process are later picked up by the proxy's - ``update_reward``/``submit_experiences`` via the shared ``msg_id`` key. - -Keying: experiences are identified by ``eid.suffix`` (the vLLM ``request_id``, -== the OpenAI ``response.id`` == the proxy ``msg_id``). ``eid.task``/``run``/ -``reward`` are assigned by the proxy at feedback time, not here. +A ``RecordStore`` persists Trinity ``Experience`` objects in the vLLM API +server process. The only backend is ``MemoryStore`` — in-process, keyed by the +recording identity (``info["task_id"]`` = the API key / Ray-injected task id), +falling back to ``eid.suffix`` (the vLLM ``request_id``) when no identity was +supplied. + +The consume side is ``update_reward_by_task_id``: it sets ``reward``/``run``/ +``task`` on every experience in a task-id group, pops the group, and returns +it. This is the in-memory replacement for the old SQL ``HistoryRecorder``- +mediated join — the coordinator calls it (via ``/records/consume_task``) at +finalize time, so heavy experience bytes cross the network exactly once (store +→ coordinator pipeline) and never through Ray. + +Keying: experiences are identified by ``eid.suffix`` (the vLLM ``request_id``) +for traceability, but **grouped** by ``info["task_id"]`` so a whole task's +worth of turns/samples can be reward-updated and consumed together. Concurrency: ``append_turn`` is called from a single background flusher task; -the async signatures keep the door open for I/O-bound backends. +``update_reward_by_task_id`` is called from the ``/records/consume_task`` HTTP +handler. Both run in the same asyncio loop, so the dict is single-writer-safe +across these two without a lock. """ import abc from collections import defaultdict @@ -36,6 +41,24 @@ class RecordStore(abc.ABC): async def append_turn(self, exp: Experience) -> None: """Persist one completed experience.""" + @abc.abstractmethod + async def update_reward_by_task_id( + self, task_id: str, reward: float, run: int, task: str + ) -> list[Experience]: + """Set reward/run/task on every experience in the group, pop and return it. + + Args: + task_id: The recording identity (group key). When the recorded + experience had no identity, this is its ``eid.suffix``. + reward: Reward to stamp on every experience in the group. + run: Run id to stamp on ``eid.run``. + task: Task id to stamp on ``eid.task``. + + Returns: + The (now reward-stamped) experiences of the group, in insertion + order. Empty list if the group was absent. + """ + @abc.abstractmethod async def get_task(self, task_id: str) -> list[Experience]: """Return all experiences for a task, in insertion order.""" @@ -57,12 +80,15 @@ class MemoryStore(RecordStore): """In-process store. Groups experiences by recording identity (``info["task_id"]``) when an API - key was supplied, otherwise each turn is keyed by its own ``eid.suffix`` - (request_id) — so a missing identity never collapses distinct turns. - ``get_turn`` resolves an individual turn by ``info["request_id"]``. + key / task id was supplied, otherwise each turn is keyed by its own + ``eid.suffix`` (request_id) — so a missing identity never collapses + distinct turns. ``get_turn`` resolves an individual turn by + ``info["request_id"]``. Note: per-process under data-parallel serving — each API-server rank holds - only the experiences it served. For cross-rank aggregation, use ``SqlStore``. + only the experiences it served. The coordinator fans out + ``/records/consume_task`` to every rank and merges, so cross-rank + aggregation happens at consume time, not in storage. """ def __init__(self) -> None: @@ -77,6 +103,16 @@ def _group_key(exp: Experience) -> str: async def append_turn(self, exp: Experience) -> None: self._records[self._group_key(exp)].append(exp) + async def update_reward_by_task_id( + self, task_id: str, reward: float, run: int, task: str + ) -> list[Experience]: + exps = self._records.pop(task_id, []) + for exp in exps: + exp.reward = reward + exp.eid.run = run + exp.eid.task = task + return exps + async def get_task(self, task_id: str) -> list[Experience]: return list(self._records.get(task_id, [])) @@ -91,46 +127,3 @@ async def list_tasks(self) -> list[str]: async def delete_task(self, task_id: str) -> None: self._records.pop(task_id, None) - - -class SqlStore(RecordStore): - """SQL-backed store sharing the explorer proxy's ``proxy_history`` table. - - Writes go through ``HistoryRecorder.record_history`` (which ``prepare()``s - the engine on first use and maps ``eid.suffix`` -> ``msg_id``). The explorer - proxy's own ``HistoryRecorder`` instance reads/updates the same rows for - ``/feedback`` and ``/commit``, so the in-vLLM recorder and the proxy share - one table by ``db_url`` + ``table_name``. - - Reads (``get_task``/``get_turn``/``list_tasks``/``delete_task``) are NOT - implemented here: in the online-RL setup the proxy owns the read/consume - side. The ``/records`` query endpoints surface this as 503 when this backend - is active. - """ - - #: Marks that this backend does not serve the ``/records`` read endpoints. - supports_reads = False - - def __init__(self, db_url: str, table_name: str) -> None: - # Imported lazily so the vLLM process only pulls in the SQL/explorer - # stack when this backend is actually selected. - from trinity.explorer.proxy.recorder import HistoryRecorder - - self._recorder = HistoryRecorder(db_url=db_url, table_name=table_name) - - async def append_turn(self, exp: Experience) -> None: - # record_history() calls prepare() on first use; serializes the - # experience into the blob column and writes meta keyed by msg_id. - await self._recorder.record_history([exp]) - - async def get_task(self, task_id: str) -> list[Experience]: - raise NotImplementedError("SqlStore reads are served by the explorer proxy; use /feedback") - - async def get_turn(self, task_id: str, request_id: str) -> Optional[Experience]: - raise NotImplementedError("SqlStore reads are served by the explorer proxy; use /feedback") - - async def list_tasks(self) -> list[str]: - raise NotImplementedError("SqlStore reads are served by the explorer proxy; use /feedback") - - async def delete_task(self, task_id: str) -> None: - raise NotImplementedError("SqlStore reads are served by the explorer proxy; use /feedback") diff --git a/trinity/common/workflows/envs/alfworld/RAFT_utils.py b/trinity/common/workflows/envs/alfworld/RAFT_utils.py index 5e57ba597a0..77c48af3dc5 100644 --- a/trinity/common/workflows/envs/alfworld/RAFT_utils.py +++ b/trinity/common/workflows/envs/alfworld/RAFT_utils.py @@ -113,6 +113,10 @@ async def process_messages_to_experience_async(model, messages, info=None) -> Ex info = {} try: + # TODO(recording): when use_recorded_experience is on, replace this + # client-side conversion with a MemoryStore lookup by the session's + # task_id_key (concatenate turns via info["sample_index"]); see + # workflow.process_messages_to_experience. converted_experience = await model.convert_messages_to_experience_async(messages) metrics = {} diff --git a/trinity/common/workflows/workflow.py b/trinity/common/workflows/workflow.py index 25853322fd9..43db800b5c0 100644 --- a/trinity/common/workflows/workflow.py +++ b/trinity/common/workflows/workflow.py @@ -89,6 +89,14 @@ class Workflow: can_repeat: bool = False # whether the workflow can be repeated multiple times. If true, `set_repeat_times()` must be implemented. is_async: bool = False # whether the workflow runs in async mode. If true, `run_async()` must be implemented, else `run()` must be implemented. + # When True, the workflow cooperates with the in-vLLM recorder: each chat + # call carries a per-sample ``task_id_key`` (the MemoryStore group key) so + # the runner can later report ``{task_id_key: reward}`` and the coordinator + # can join reward inside the store. Set by the WorkflowRunner from + # ``config.explorer.use_recorded_experience``. See ``SimpleWorkflow`` for + # the per-sample (n=1) loop this triggers. + use_recorded_experience: bool = False + def __init__( self, *, @@ -219,6 +227,14 @@ def _build_experience_from_converted( def process_messages_to_experience( self, messages, reward, info={}, truncate_status=None ) -> Experience: + # TODO(recording): when use_recorded_experience is on, this client-side + # conversion is redundant — the vLLM recorder's build_experience already + # captured the authoritative heavy data (real logprobs without an extra + # forward, real routed_experts) into the MemoryStore, keyed by the + # task_id the chat call carried. Replace this with an in-process lookup + # by task_id (store.get_task / consume), then concatenate the session's + # turns (info["sample_index"] orders them) into one experience here. + # Requires threading the per-call task_id_key down to this call site. converted_experience = self.model.convert_messages_to_experience(messages) return self._build_experience_from_converted( converted_experience, @@ -230,6 +246,8 @@ def process_messages_to_experience( async def process_messages_to_experience_async( self, messages, reward, info={}, truncate_status=None ) -> Experience: + # TODO(recording): see process_messages_to_experience — replace with a + # MemoryStore lookup by task_id once task_id_key is threaded here. converted_experience = await self.model.convert_messages_to_experience_async(messages) return self._build_experience_from_converted( converted_experience, @@ -289,37 +307,68 @@ def format_messages(self): messages.append({"role": "assistant", "content": self.reply_prefix}) return messages + # -- recording-path helpers (shared by SimpleWorkflow / AsyncSimpleWorkflow) - + def _recorded_task_id_key(self, run_index: int) -> str: + """Per-sample recording identity (the MemoryStore group key).""" + return f"{self.task.batch_id}/{self.task.task_id}/{self.run_id_base + run_index}" -class SimpleWorkflow(BaseSimpleWorkflow): - """A workflow for simple single-round task.""" - - can_reset: bool = True - can_repeat: bool = True - - def run(self) -> List[Experience]: - # TODO: Optimize the generate function - messages = self.format_messages() - - self.logger.debug("start chat") - responses = self.model.chat(messages, **self.rollout_args) + def _attach_rewards(self, responses, *, base: int) -> List[Experience]: for i, response in enumerate(responses): + run = base + i reward_dict = self.reward_fn( # type: ignore [misc] response=response.response_text, # type: ignore [arg-type] truth=self.truth, ) - if response.metrics is None: response.metrics = {} response.metrics.update(reward_dict) reward = sum(reward_dict.values()) response.reward = reward - response.eid.run = i + self.run_id_base - + response.eid.run = run self.logger.debug( - f"self.task_desc: {self.task_desc}, messages: {messages}, response: {response.response_text}, reward: {reward}" + f"self.task_desc: {self.task_desc}, response: {response.response_text}, reward: {reward}" ) return responses + @staticmethod + def _stamp_task_id_key(exps: List[Experience], task_id_key: str) -> None: + for exp in exps: + if exp.info is None: + exp.info = {} + exp.info["task_id_key"] = task_id_key + + +class SimpleWorkflow(BaseSimpleWorkflow): + """A workflow for simple single-round task.""" + + can_reset: bool = True + can_repeat: bool = True + + def run(self) -> List[Experience]: + # TODO: Optimize the generate function + messages = self.format_messages() + + self.logger.debug("start chat") + if self.use_recorded_experience: + return self._run_recorded(messages) + responses = self.model.chat(messages, **self.rollout_args) + return self._attach_rewards(responses, base=self.run_id_base) + + def _run_recorded(self, messages) -> List[Experience]: + # One chat call per sample (n=1) so each gets a distinct task_id_key + # (the recording group key == reward unit). The runner later reports + # {task_id_key: reward} and the coordinator joins reward in-store. + rollout_args = dict(self.rollout_args) + rollout_args["n"] = 1 + exps: List[Experience] = [] + for i in range(self.repeat_times): + task_id_key = self._recorded_task_id_key(i) + responses = self.model.chat(messages, task_id_key=task_id_key, **rollout_args) + rewarded = self._attach_rewards(responses, base=self.run_id_base + i) + self._stamp_task_id_key(rewarded, task_id_key) + exps.extend(rewarded) + return exps + class AsyncSimpleWorkflow(BaseSimpleWorkflow): is_async: bool = True @@ -329,24 +378,24 @@ async def run_async(self) -> List[Experience]: messages = self.format_messages() self.logger.info("start chat") + if self.use_recorded_experience: + return await self._run_recorded_async(messages) responses = await self.model.chat_async(messages, **self.rollout_args) - for i, response in enumerate(responses): - reward_dict = self.reward_fn( # type: ignore [misc] - response=response.response_text, # type: ignore [arg-type] - truth=self.truth, + return self._attach_rewards(responses, base=self.run_id_base) + + async def _run_recorded_async(self, messages) -> List[Experience]: + rollout_args = dict(self.rollout_args) + rollout_args["n"] = 1 + exps: List[Experience] = [] + for i in range(self.repeat_times): + task_id_key = self._recorded_task_id_key(i) + responses = await self.model.chat_async( + messages, task_id_key=task_id_key, **rollout_args ) - - if response.metrics is None: - response.metrics = {} - response.metrics.update(reward_dict) - reward = sum(reward_dict.values()) - response.reward = reward - response.eid.run = i + self.run_id_base - - self.logger.debug( - f"self.task_desc: {self.task_desc}, messages: {messages}, response: {response.response_text}, reward: {reward}" - ) - return responses + rewarded = self._attach_rewards(responses, base=self.run_id_base + i) + self._stamp_task_id_key(rewarded, task_id_key) + exps.extend(rewarded) + return exps class MathWorkflow(SimpleWorkflow): diff --git a/trinity/explorer/proxy/app.py b/trinity/explorer/proxy/app.py index cbc97eb57f4..2a3bb1101a1 100644 --- a/trinity/explorer/proxy/app.py +++ b/trinity/explorer/proxy/app.py @@ -233,33 +233,6 @@ async def metrics(request: Request): return JSONResponse(content=metrics) -@app.post("/feedback") -async def feedback(request: Request): - """Receive feedback for the current session.""" - body = await request.json() - reward = body.get("reward") - msg_ids = body.get("msg_ids") - task_id = body.get("task_id") - run_id = body.get("run_id", 0) - if msg_ids is None or reward is None: - return JSONResponse(status_code=400, content={"error": "msg_ids and reward are required"}) - if not isinstance(msg_ids, list) or not isinstance(reward, (int, float)): - return JSONResponse( - status_code=400, content={"error": "msg_ids must be a list and reward must be a number"} - ) - await request.app.state.service.record_feedback( - reward=reward, msg_ids=msg_ids, task_id=task_id, run_id=run_id - ) - return JSONResponse(content={"status": "success"}) - - -@app.post("/commit") -async def commit(request: Request): - """Commit the current experiences.""" - await request.app.state.service.submit_experiences() - return JSONResponse(content={"status": "success"}) - - async def serve_http(app: FastAPI, host: str, port: int) -> None: config = uvicorn.Config(app, host=host, port=port) server = uvicorn.Server(config) diff --git a/trinity/explorer/proxy/recorder.py b/trinity/explorer/proxy/recorder.py deleted file mode 100644 index d670978652f..00000000000 --- a/trinity/explorer/proxy/recorder.py +++ /dev/null @@ -1,138 +0,0 @@ -from typing import Dict, List, Set - -from sqlalchemy import select, update -from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker - -from trinity.buffer.schema.sql_schema import init_async_engine -from trinity.buffer.utils import async_run_with_retry_session -from trinity.common.experience import Experience -from trinity.utils.log import get_logger - - -class HistoryRecorder: - """Record chat history into the database using async SQL.""" - - def __init__(self, db_url: str, table_name: str): - self.logger = get_logger() - self._db_url = db_url - self._table_name = table_name - self._initialized = False - - async def prepare(self) -> None: - if self._initialized: - return - engine, self.meta_cls, self.blob_cls = await init_async_engine( - db_url=self._db_url, - table_name=self._table_name, - schema_type="experience", - ) - self.session = async_sessionmaker(engine, expire_on_commit=False) - self._initialized = True - self.logger.info(f"Init async SQL storage at {self._db_url}") - - async def record_history(self, experiences: List[Experience]) -> None: - """Save experiences to the database.""" - await self.prepare() - - async def operation(session: AsyncSession): - for exp in experiences: - meta_row = self.meta_cls.from_experience(exp) - session.add(meta_row) - await session.flush() - blob_row = self.blob_cls(id=meta_row.id, experience_bytes=exp.serialize()) - session.add(blob_row) - - await async_run_with_retry_session(self.session, operation) - - async def update_reward( - self, reward: float, msg_ids: list, run_id: int, task_id: str - ) -> List[Experience]: - """Update reward for given response IDs and return the updated experiences. - - Only experiences that have not been consumed (consumed == 0) will be returned. - """ - await self.prepare() - - meta_cls = self.meta_cls - blob_cls = self.blob_cls - - async def operation(session: AsyncSession): - stmt = ( - select(meta_cls) - .where(meta_cls.msg_id.in_(msg_ids), meta_cls.consumed == 0) - .with_for_update() - ) - result = await session.execute(stmt) - records = result.scalars().all() - - if not records: - return [] - - ids = [record.id for record in records] - - update_stmt = ( - update(meta_cls) - .where(meta_cls.id.in_(ids)) - .values( - reward=reward, - run_id=run_id, - task_id=task_id, - consumed=meta_cls.consumed + 1, - ) - ) - await session.execute(update_stmt) - - blob_stmt = select(blob_cls).where(blob_cls.id.in_(ids)) - blob_result = await session.execute(blob_stmt) - blobs = blob_result.scalars().all() - blob_map = {b.id: b.experience_bytes for b in blobs} - - # Re-fetch meta rows to get updated values - refresh_stmt = select(meta_cls).where(meta_cls.id.in_(ids)) - refresh_result = await session.execute(refresh_stmt) - updated_records = refresh_result.scalars().all() - - updated_experiences = [] - for record in updated_records: - blob_bytes = blob_map.get(record.id) - if blob_bytes is not None: - updated_experiences.append(record.to_experience(blob_bytes)) - return updated_experiences - - return await async_run_with_retry_session(self.session, operation) - - -class MemoryHistoryRecorder: - """ - In-memory version of HistoryRecorder for high-performance reward update and history recording. - All data is stored in memory, and can be flushed to persistent storage as needed. - """ - - def __init__(self): - self.logger = get_logger() - # msg_id -> Experience - self._exp_map: Dict[str, Experience] = {} - # Set of msg_id that are not consumed - self._unconsumed: Set[str] = set() - - async def record_history(self, experiences: List[Experience]) -> None: - """Save experiences in memory.""" - for exp in experiences: - self._exp_map[exp.eid.suffix] = exp - if getattr(exp, "consumed", 0) == 0: - self._unconsumed.add(exp.eid.suffix) - - async def update_reward( - self, reward: float, msg_ids: list, run_id: int, task_id: str - ) -> List[Experience]: - """Update reward for given response IDs and return the updated experiences.""" - updated = [] - for msg_id in msg_ids: - if msg_id in self._unconsumed and msg_id in self._exp_map: - exp = self._exp_map.pop(msg_id) - exp.reward = reward - exp.eid.run = run_id - exp.eid.task = task_id - self._unconsumed.remove(msg_id) - updated.append(exp) - return updated diff --git a/trinity/explorer/proxy/service.py b/trinity/explorer/proxy/service.py index 1fe77d1564a..4a247442410 100644 --- a/trinity/explorer/proxy/service.py +++ b/trinity/explorer/proxy/service.py @@ -4,15 +4,21 @@ from typing import Dict, List, Tuple from trinity.common.constants import RunningStatus, SyncMethod -from trinity.common.experience import Experience from trinity.common.models.model import ModelWrapper from trinity.explorer.explorer import Explorer -from trinity.explorer.proxy.recorder import HistoryRecorder from trinity.utils.log import get_logger class ExplorerService: - """Manages the lifecycle and operations of the Explorer API service.""" + """Manages the lifecycle and operations of the Explorer API service. + + The proxy is a request router + model-weight sync coordinator for serve + mode. Experience collection used to live here (SQL-mediated + ``/feedback``/``/commit``); it has been removed in favor of the in-vLLM + ``MemoryStore`` + ``/records/consume_task`` flow driven by the rollout + coordinator. Serve-mode external reward reporting is therefore pending + (see the recording refactor plan). + """ def __init__(self, explorer: Explorer, listen_address: str = "localhost", port: int = 8010): self.logger = get_logger(__name__) @@ -29,15 +35,6 @@ def __init__(self, explorer: Explorer, listen_address: str = "localhost", port: self.model_version_map: Dict[int, int] = {} # model index -> model version self.sync_task_map: Dict[asyncio.Future, int] = {} # sync task -> model index self.latest_model_version = 0 - self.session_level_experience_queue: Dict[int, deque[Experience]] = {} - self.commit_lock = asyncio.Lock() - self.ready_experiences = deque() - self.recorder = HistoryRecorder( - db_url=explorer.config.explorer.db_url - or f"sqlite:///{explorer.config.buffer.cache_dir}/proxy_history.db", - table_name="proxy_history", - ) - self.ready_experience_count = 0 async def serve(self) -> None: from trinity.explorer.proxy.app import run_app @@ -129,30 +126,8 @@ def collect_metrics(self) -> Dict: for i, model in enumerate(self.models): metrics[f"rollout/model_{i}/total_request_count"] = model.request_count metrics[f"rollout/model_{i}/model_version"] = model.model_version - metrics["rollout/ready_experience_count"] = self.ready_experience_count return metrics - async def submit_experiences(self) -> None: - async with self.commit_lock: - experiences = list(self.ready_experiences) - self.ready_experiences.clear() - metrics = await self.explorer.rollout_coordinator.process_experiences.remote( - [Experience.serialize_many(experiences)] - ) - metrics.update(self.collect_metrics()) - self.explorer.explore_step_num += 1 - self.explorer.monitor.log(metrics, self.explorer.explore_step_num) - - async def record_feedback(self, reward: float, msg_ids: List[str], task_id: str, run_id: int): - exps = await self.recorder.update_reward( - reward=reward, - msg_ids=msg_ids, - task_id=task_id, - run_id=run_id, - ) - self.ready_experience_count += len(exps) - self.ready_experiences.extend(exps) - async def shutdown(self): if not self.running: self.logger.warning("Server is not running.") diff --git a/trinity/explorer/rollout_coordinator.py b/trinity/explorer/rollout_coordinator.py index b5c98032ab2..b2fae011dc0 100644 --- a/trinity/explorer/rollout_coordinator.py +++ b/trinity/explorer/rollout_coordinator.py @@ -1,15 +1,19 @@ """Rollout coordinator for async batch submission and finalize.""" import asyncio +import pickle +from collections import defaultdict from dataclasses import dataclass, field from enum import Enum -from typing import Any, Dict, Literal, Optional, Union +from typing import Any, Dict, List, Literal, Optional, Union +import httpx import ray from ray.actor import ActorHandle from trinity.buffer.pipelines.experience_pipeline import ExperiencePipeline from trinity.common.config import Config +from trinity.common.experience import Experience from trinity.common.workflows import Task from trinity.explorer.scheduler import Scheduler from trinity.utils.log import get_logger @@ -18,6 +22,10 @@ BatchId = Union[int, str] BatchType = Literal["train", "eval"] +#: Default per-rank consume HTTP timeout (seconds). The consume returns heavy +#: experience bytes, so allow generous headroom over the inference timeout. +_CONSUME_TIMEOUT = 300.0 + class BatchLifecycleState(str, Enum): """Lifecycle states for one submitted batch.""" @@ -64,6 +72,42 @@ def __init__( self.pending_batches: Dict[BatchId, BatchState] = {} self.running = False self.detailed_stats = getattr(getattr(config, "monitor", None), "detailed_stats", False) + # Lazily-resolved map of rollout engine_id -> API server URL, for the + # recording path's per-rank /records/consume_task fan-out. + self._rank_urls: Optional[Dict[int, str]] = None + + def _use_recorded_experience(self) -> bool: + """Whether the recording-consume path is active for train batches.""" + return bool(self.config.explorer.use_recorded_experience) + + def _resolve_rank_urls(self) -> Dict[int, str]: + """Resolve each rollout engine's API server URL via named Ray actors. + + Mirrors ``Allocator.get_actor_name`` + ``ray.get_actor``: rollout model + actors are named ``f"{explorer.name}_rollout_model_{engine_id}_0"`` + (node_id 0 holds the API server). Cached after first resolution. + """ + if self._rank_urls is not None: + return self._rank_urls + rollout_cfg = self.config.explorer.rollout_model + name = self.config.explorer.name + namespace = rollout_cfg.ray_namespace + urls: Dict[int, str] = {} + for engine_id in range(rollout_cfg.engine_num): + actor_name = f"{name}_rollout_model_{engine_id}_0" + try: + actor = ray.get_actor(actor_name, namespace=namespace) + except ValueError: + self.logger.warning( + "rollout actor %s not found in namespace %s; skipping rank %d", + actor_name, + namespace, + engine_id, + ) + continue + urls[engine_id] = ray.get(actor.get_api_server_url.remote()) + self._rank_urls = urls + return urls async def prepare(self) -> None: """Initialize the owned pipeline and scheduler.""" @@ -275,7 +319,10 @@ async def _finalize_train_batch( batch_state.state = BatchLifecycleState.FINALIZING try: - pipeline_metrics = await self.process_experiences(payload_chunks) + if self._use_recorded_experience(): + pipeline_metrics = await self._consume_recorded_experiences(payload_chunks) + else: + pipeline_metrics = await self.process_experiences(payload_chunks) if not is_complete: await self._cleanup_train_batch_runtime(batch_state) except Exception: @@ -284,6 +331,66 @@ async def _finalize_train_batch( return self._finish_batch(batch_state, pipeline_metrics=pipeline_metrics) + async def _consume_recorded_experiences(self, payload_chunks: List[bytes]) -> dict: + """Recording path: pull heavy experiences from each vLLM rank's store. + + ``payload_chunks`` are small pickle reward maps produced by the runners + (``{"engine_id": int, "updates": [{"task_id", "reward", "run", "task"}]}``). + Group updates by engine, fan out ``POST /records/consume_task`` to each + rank (which drains its recorder, reward-stamps the matching task-id + groups, pops them, and returns ``serialize_many`` bytes), deserialize, + and feed the assembled experiences straight into the pipeline — no Ray + serialization of heavy tensors, and reward is fused inside the store. + """ + if self.experience_pipeline is None: + raise RuntimeError("Experience pipeline is not initialized.") + per_engine: Dict[int, List[dict]] = defaultdict(list) + for chunk in payload_chunks: + if not chunk: + continue + data = pickle.loads(chunk) + per_engine[int(data["engine_id"])].extend(data["updates"]) + + if not per_engine: + return {} + + rank_urls = self._resolve_rank_urls() + async with httpx.AsyncClient(timeout=_CONSUME_TIMEOUT) as client: + requests = [ + self._post_consume_task(client, rank_urls[engine_id], updates) + for engine_id, updates in per_engine.items() + if engine_id in rank_urls + ] + responses = await asyncio.gather(*requests) + + exps: List[Experience] = [] + for resp_bytes in responses: + if resp_bytes: + exps.extend(Experience.deserialize_many(resp_bytes)) + return await self.experience_pipeline.process_experiences(exps) + + async def _post_consume_task( + self, client: httpx.AsyncClient, rank_url: str, updates: List[dict] + ) -> bytes: + """POST a batch of task-id reward updates to one rank; return heavy bytes.""" + try: + resp = await client.post( + f"{rank_url}/records/consume_task", + json={"updates": updates}, + ) + except (httpx.TimeoutException, httpx.RequestError) as exc: + self.logger.error("consume_task to %s failed: %s", rank_url, exc) + return b"" + if resp.status_code != 200: + self.logger.error( + "consume_task to %s returned %d: %s", + rank_url, + resp.status_code, + resp.text[:200], + ) + return b"" + return resp.content + def _finish_batch( self, batch_state: BatchState, diff --git a/trinity/explorer/workflow_runner.py b/trinity/explorer/workflow_runner.py index dcdcd2fb5ea..b52f529bcc0 100644 --- a/trinity/explorer/workflow_runner.py +++ b/trinity/explorer/workflow_runner.py @@ -3,6 +3,7 @@ import asyncio import os +import pickle import time import traceback from dataclasses import dataclass @@ -117,8 +118,12 @@ def _create_workflow_instance(self, task: Task) -> Workflow: ) else: self.workflow_instance.reset(task) + self.workflow_instance.use_recorded_experience = self._use_recorded_experience() return self.workflow_instance + def _use_recorded_experience(self) -> bool: + return bool(self.config.explorer.use_recorded_experience) + async def _run_workflow(self, workflow_instance: Workflow) -> List[Experience]: if workflow_instance.asynchronous: exps = await workflow_instance.run_async() @@ -127,7 +132,7 @@ async def _run_workflow(self, workflow_instance: Workflow) -> List[Experience]: return exps def _create_isolated_workflow_instance(self, task: Task) -> Workflow: - return task.to_workflow( + wf = task.to_workflow( ( self.model_wrapper.clone_with_isolated_history() if self.config.explorer.rollout_model.enable_history @@ -135,6 +140,8 @@ def _create_isolated_workflow_instance(self, task: Task) -> Workflow: ), self.auxiliary_model_wrappers, ) + wf.use_recorded_experience = self._use_recorded_experience() + return wf def _build_execution_result( self, @@ -417,8 +424,30 @@ async def run_task( status = execution_result.status if task.is_eval: - # If the task is an evaluation task, we do not record the experiences to the buffer + # Eval tasks are not written to the training buffer. Under the + # recording path their turns are still in the vLLM MemoryStore; + # the coordinator's eval finalize must drain-and-discard them + # (TODO: wire eval consume-and-discard so eval turns don't leak + # in the store). For now, return no payload. return status, b"" + elif self._use_recorded_experience(): + # Recording path: ship only the small reward map keyed by the + # per-sample task_id_key the workflow stamped on each exp. The + # heavy experiences live in the vLLM MemoryStore and are pulled + # by the coordinator at finalize time. + updates = [ + { + "task_id": exp.info.get("task_id_key") or exp.eid.suffix, + "reward": exp.reward, + "run": exp.eid.run, + "task": str(task.task_id), + } + for exp in exps + ] + reward_payload = pickle.dumps( + {"engine_id": self.rollout_model_id, "updates": updates} + ) + return status, reward_payload else: exp_payload = Experience.serialize_many(exps) return status, exp_payload From 4ae939220642bd43f43921ffb78cd03544de814e Mon Sep 17 00:00:00 2001 From: pxc Date: Thu, 25 Jun 2026 14:38:12 +0800 Subject: [PATCH 07/84] simplify code --- trinity/common/config.py | 22 +++--- trinity/common/models/model.py | 10 +-- trinity/common/models/sglang_model.py | 6 +- trinity/common/models/vllm_model.py | 26 +++---- .../models/vllm_patch/recording/__init__.py | 4 +- .../models/vllm_patch/recording/context.py | 25 ++++--- .../models/vllm_patch/recording/models.py | 28 ++++---- .../models/vllm_patch/recording/query.py | 62 ++++++++--------- .../models/vllm_patch/recording/recorder.py | 22 +++--- .../models/vllm_patch/recording/store.py | 69 ++++++++++--------- .../workflows/envs/alfworld/RAFT_utils.py | 4 +- trinity/common/workflows/workflow.py | 51 +++++++------- trinity/explorer/rollout_coordinator.py | 12 ++-- trinity/explorer/workflow_runner.py | 14 ++-- 14 files changed, 180 insertions(+), 175 deletions(-) diff --git a/trinity/common/config.py b/trinity/common/config.py index c877046651c..18fb999f20c 100644 --- a/trinity/common/config.py +++ b/trinity/common/config.py @@ -594,10 +594,16 @@ class InferenceModelConfig: # Turn on in-vLLM generation recording for the OpenAI API serving path: the # engine wraps ``engine_client.generate`` and writes each finished turn as a # Trinity ``Experience`` to the in-process ``MemoryStore``, keyed by the - # recording identity (``task_id``). When True, the Allocator forces - # ``enable_return_routed_experts``. VLLMModel mirrors the recording config - # onto the engine instance for the recorder to read. The capture width - # (top-k logprobs) reuses ``logprobs`` below (default 1). + # recording identity (``record_key``). This is the single switch for the + # recording flow — when on, the explorer also consumes from the store: the + # WorkflowRunner ships only a small reward map keyed by ``record_key`` and + # the RolloutCoordinator pulls heavy experiences via ``/records/consume_task`` + # at finalize time. When off (default), runners ship serialized experiences + # through the scheduler as before (legacy path). When True, the Allocator + # forces ``enable_return_routed_experts``. VLLMModel mirrors the recording + # config onto the engine instance for the recorder to read. The capture + # width (top-k logprobs) reuses ``logprobs`` below (default 1). Requires + # ``enable_openai_api=True`` (the recording runner is the API server). enable_recording: bool = False # Buffer size (MB) for batched NCCL weight sync. Controls peak GPU memory during sync. @@ -793,14 +799,6 @@ class ExplorerConfig: # Maximum number of train batches that RolloutCoordinator can hold simultaneously. max_inflight_batches: int = 2 - # Use the in-vLLM recording flow: runners report only a small reward map - # (keyed by task id) and the coordinator pulls heavy experiences from each - # vLLM rank's MemoryStore via /records/consume_task at finalize time. - # Requires rollout_model.enable_recording=True and enable_openai_api=True. - # When False (default), runners ship serialized experiences through the - # scheduler as before (legacy path). - use_recorded_experience: bool = False - @dataclass class MegatronParallelConfig: diff --git a/trinity/common/models/model.py b/trinity/common/models/model.py index 67976f7e675..70e56a367d0 100644 --- a/trinity/common/models/model.py +++ b/trinity/common/models/model.py @@ -288,7 +288,7 @@ async def convert_messages_to_experience( # redundant — it re-tokenizes messages and runs an extra logprobs # forward (and fakes routed_experts), all of which build_experience # already captured at generation time into the MemoryStore. Redirect to - # a store lookup by the call's task_id_key once it's threaded here. + # a store lookup by the call's record_key once it's threaded here. """Convert a list of messages into an experience in async. Args: @@ -526,24 +526,24 @@ async def generate_async(self, prompts: List[str], **kwargs) -> List[Experience] @_history_recorder def chat( - self, messages: List[dict], task_id_key: Optional[str] = None, **kwargs + self, messages: List[dict], record_key: Optional[str] = None, **kwargs ) -> List[Experience]: """Generate a list of experiences from a list of messages.""" lora_request = self.get_lora_request() return ray.get( self.model.chat.remote( - messages, lora_request=lora_request, task_id_key=task_id_key, **kwargs + messages, lora_request=lora_request, record_key=record_key, **kwargs ) ) @_history_recorder async def chat_async( - self, messages: List[dict], task_id_key: Optional[str] = None, **kwargs + self, messages: List[dict], record_key: Optional[str] = None, **kwargs ) -> List[Experience]: """Generate a list of experiences from a list of messages in async.""" lora_request = await self.get_lora_request_async() return await self.model.chat.remote( - messages, lora_request=lora_request, task_id_key=task_id_key, **kwargs + messages, lora_request=lora_request, record_key=record_key, **kwargs ) def logprobs(self, tokens: List[int], temperature: Optional[float] = None) -> Tensor: diff --git a/trinity/common/models/sglang_model.py b/trinity/common/models/sglang_model.py index 49263d72da3..b2602daba5a 100644 --- a/trinity/common/models/sglang_model.py +++ b/trinity/common/models/sglang_model.py @@ -438,12 +438,12 @@ async def chat( self, messages: List[dict], lora_request=None, - task_id_key: Optional[str] = None, + record_key: Optional[str] = None, **kwargs, ) -> Sequence[Experience]: - # ``task_id_key`` is the in-vLLM recorder's group key; SGLang does not + # ``record_key`` is the in-vLLM recorder's group key; SGLang does not # support recording, so accept and ignore it here. - del task_id_key + del record_key if self.tokenizer is None: await self._initialize_tokenizer() diff --git a/trinity/common/models/vllm_model.py b/trinity/common/models/vllm_model.py index 3b32d8bc4d5..6011e421dde 100644 --- a/trinity/common/models/vllm_model.py +++ b/trinity/common/models/vllm_model.py @@ -20,8 +20,8 @@ from trinity.common.models.model import BaseInferenceModel from trinity.common.models.vllm_patch import get_vllm_version from trinity.common.models.vllm_patch.recording.context import ( + record_key_ctx, skip_recording_ctx, - task_id_ctx, ) @@ -210,18 +210,18 @@ async def chat( self, messages: List[Dict], lora_request=None, - task_id_key: Optional[str] = None, + record_key: Optional[str] = None, **kwargs, ) -> Sequence[Experience]: """Chat with the model with a list of messages in async. Args: messages (List[dict]): The input history messages. - task_id_key (Optional[str]): Recording identity for the in-vLLM + record_key (Optional[str]): Recording identity for the in-vLLM recorder (the MemoryStore group key). Propagated to - ``generate`` via ``task_id_ctx`` so the recorder stamps it into - ``info["task_id"]`` without an HTTP hop. None leaves it to the - engine request id (fallback grouping). + ``generate`` via ``record_key_ctx`` so the recorder stamps it + into ``info["record_key"]`` without an HTTP hop. None leaves it + to the engine request id (fallback grouping). kwargs (dict): A dictionary of sampling parameters. Returns: @@ -248,7 +248,7 @@ async def chat( "multi_modal_data": multi_modal_data or {}, } return await self.generate( - prompt=prompt, lora_request=lora_request, task_id_key=task_id_key, **kwargs + prompt=prompt, lora_request=lora_request, record_key=record_key, **kwargs ) def _extract_routed_experts(self, output: Any, output_index: int) -> Optional[torch.Tensor]: @@ -276,15 +276,15 @@ async def generate( self, prompt: Union[str, Dict], lora_request=None, - task_id_key: Optional[str] = None, + record_key: Optional[str] = None, **kwargs, ) -> Sequence[Experience]: """Generate a response from the provided prompt in async. Args: prompt (str): The input prompt. - task_id_key (Optional[str]): Recording identity propagated to the - in-vLLM recorder via ``task_id_ctx`` (see ``chat``). + record_key (Optional[str]): Recording identity propagated to the + in-vLLM recorder via ``record_key_ctx`` (see ``chat``). kwargs (dict): A dictionary of sampling parameters. Returns: @@ -307,14 +307,14 @@ async def generate( # Propagate the recording identity to the engine-level recorder (same # async task, same process) so the recorded experience is grouped under - # this task id in the MemoryStore. - task_id_token = task_id_ctx.set(task_id_key) + # this record key in the MemoryStore. + record_key_token = record_key_ctx.set(record_key) try: output = await self._generate_internal( prompt=prompt, lora_request=lora_request, **kwargs ) finally: - task_id_ctx.reset(task_id_token) + record_key_ctx.reset(record_key_token) if is_mm_prompt: if self.mm_render is None: self.mm_render = vLLMMultiModalRender( diff --git a/trinity/common/models/vllm_patch/recording/__init__.py b/trinity/common/models/vllm_patch/recording/__init__.py index a10f2f800f9..a8ec841bef6 100644 --- a/trinity/common/models/vllm_patch/recording/__init__.py +++ b/trinity/common/models/vllm_patch/recording/__init__.py @@ -3,8 +3,8 @@ """ from trinity.common.models.vllm_patch.recording.context import ( RecordingIdentityMiddleware, + record_key_ctx, skip_recording_ctx, - task_id_ctx, ) from trinity.common.models.vllm_patch.recording.models import build_experience from trinity.common.models.vllm_patch.recording.query import query_router @@ -25,7 +25,7 @@ "build_experience", "patch_engine_for_recording", "query_router", + "record_key_ctx", "run_api_server_with_recording", "skip_recording_ctx", - "task_id_ctx", ] diff --git a/trinity/common/models/vllm_patch/recording/context.py b/trinity/common/models/vllm_patch/recording/context.py index 9aa6cf1cc60..8e43d98714f 100644 --- a/trinity/common/models/vllm_patch/recording/context.py +++ b/trinity/common/models/vllm_patch/recording/context.py @@ -6,8 +6,13 @@ so the engine-level wrapper (which runs in the same async task as the serving handler) can recover it at record time. +The identity value is the **record key** — the MemoryStore group key under +which a whole reward unit's worth of turns (one sample / one episode) is +recorded, and the key the runner later reports ``{record_key: reward}`` against +so the coordinator can join the reward inside the store. + No identity header on a request is fine: the recorder falls back to -``request_id`` as the task id so nothing is silently dropped. +``request_id`` as the record key so nothing is silently dropped. """ from contextvars import ContextVar from typing import Optional @@ -15,10 +20,10 @@ from starlette.middleware.base import BaseHTTPMiddleware from starlette.requests import Request -# The task id for the in-flight request, or None when the client did not send -# a supported identity header (the recorder then uses request_id as a fallback -# task id). -task_id_ctx: ContextVar[Optional[str]] = ContextVar("trinity_recording_task_id", default=None) +# The record key for the in-flight request (the MemoryStore group key), or None +# when the client did not send a supported identity header (the recorder then +# uses request_id as a fallback record key). +record_key_ctx: ContextVar[Optional[str]] = ContextVar("trinity_recording_record_key", default=None) # Set around auxiliary engine.generate calls (logprobs recomputation, # convert_messages_to_experience) so the recorder skips them — those 1-token @@ -43,7 +48,7 @@ def extract_bearer_token(authorization: Optional[str]) -> Optional[str]: return token or None -def get_recording_task_id(request: Request) -> Optional[str]: +def get_recording_record_key(request: Request) -> Optional[str]: """Return the recording identity for a request. OpenAI-compatible API keys are used because every supported agent platform @@ -53,16 +58,16 @@ def get_recording_task_id(request: Request) -> Optional[str]: class RecordingIdentityMiddleware(BaseHTTPMiddleware): - """Capture request identity into ``task_id_ctx`` for the request's lifetime. + """Capture request identity into ``record_key_ctx`` for the request's lifetime. Runs in-process (ASGI) — no extra network hop, no serialization cost beyond a contextvar set/reset. """ async def dispatch(self, request: Request, call_next): - task_id = get_recording_task_id(request) - token = task_id_ctx.set(task_id) + record_key = get_recording_record_key(request) + token = record_key_ctx.set(record_key) try: return await call_next(request) finally: - task_id_ctx.reset(token) + record_key_ctx.reset(token) diff --git a/trinity/common/models/vllm_patch/recording/models.py b/trinity/common/models/vllm_patch/recording/models.py index 30cf0535105..99f558d8f25 100644 --- a/trinity/common/models/vllm_patch/recording/models.py +++ b/trinity/common/models/vllm_patch/recording/models.py @@ -12,14 +12,15 @@ request_id -> eid.suffix (``EID(suffix=...)``; the vLLM engine request id == the OpenAI ``response.id``. Kept for traceability; ``eid.task``/``run``/``reward`` are left default here and - assigned by ``MemoryStore.update_reward_by_task_id`` at + assigned by ``MemoryStore.update_reward_by_record_key`` at consume time.) - API key / task id -> info["task_id"] (the recording identity; **the group - key** the MemoryStore batches experiences by, so a whole - task's samples/turns are reward-updated and consumed - together. Falls back to ``eid.suffix`` when absent.) + API key / record key -> info["record_key"] (the recording identity; **the + group key** the MemoryStore batches experiences by, so a + whole reward unit's samples/turns are reward-updated and + consumed together. Falls back to ``eid.suffix`` when + absent.) sample index -> info["sample_index"] (position within the n-completion - set; orders samples/turns inside a task-id group) + set; orders samples/turns inside a record-key group) prompt_token_ids -> tokens (prompt portion) + prompt_length response_token_ids-> tokens (response portion) logprobs -> Experience.logprobs -- but ONLY the *chosen* token's @@ -31,7 +32,7 @@ model_version -> info["model_version"] (which checkpoint policy served the turn; read in-actor by the recorder's provider) -Plus bookkeeping (request_id / task_id / sample_index / rank / timestamp / +Plus bookkeeping (request_id / record_key / sample_index / rank / timestamp / endpoint / model_version) stashed in ``Experience.info`` so it round-trips with the experience through serialize/deserialize. """ @@ -80,7 +81,7 @@ def _extract_chosen_logprobs( def build_experience( output: Any, - task_id: Optional[str], + record_key: Optional[str], *, rank: int, timestamp: str, @@ -91,13 +92,14 @@ def build_experience( One experience per completion (``output.outputs``), so ``n > 1`` sampling is captured in full. Each experience shares ``eid.suffix = request_id`` and - ``info["task_id"] = task_id`` (the group key); ``info["sample_index"]`` + ``info["record_key"] = record_key`` (the group key); ``info["sample_index"]`` distinguishes samples within the group. Args: output: A ``RequestOutput`` with ``finished == True``. - task_id: The recording identity (API key / Ray-injected task id); - stored in ``info["task_id"]`` and used as the MemoryStore group key. + record_key: The recording identity (API key / Ray-injected record key); + stored in ``info["record_key"]`` and used as the MemoryStore group + key. rank: Data-parallel serving rank. timestamp: UTC ISO-8601 string (caller-stamped to keep this pure). endpoint: Which OpenAI endpoint served the turn (best-effort). @@ -110,7 +112,7 @@ def build_experience( """ request_id = output.request_id # eid.suffix = request_id for traceability; task/run/reward are left - # default and assigned by MemoryStore.update_reward_by_task_id at consume. + # default and assigned by MemoryStore.update_reward_by_record_key at consume. prompt_token_ids = list(output.prompt_token_ids or []) if not prompt_token_ids: @@ -136,7 +138,7 @@ def build_experience( info = { "request_id": request_id, - "task_id": task_id, + "record_key": record_key, "sample_index": sample_index, "rank": rank, "timestamp": timestamp, diff --git a/trinity/common/models/vllm_patch/recording/query.py b/trinity/common/models/vllm_patch/recording/query.py index 3629adf2275..9e98772ed69 100644 --- a/trinity/common/models/vllm_patch/recording/query.py +++ b/trinity/common/models/vllm_patch/recording/query.py @@ -11,8 +11,8 @@ into Trinity pipelines. ``POST /records/consume_task`` is the consume path: it drains the recorder, -reward-stamps every experience in the matching task-id groups, pops them, and -returns the heavy experiences as ``Experience.serialize_many`` bytes — the +reward-stamps every experience in the matching record-key groups, pops them, +and returns the heavy experiences as ``Experience.serialize_many`` bytes — the coordinator fans this out per rank at finalize time. """ from typing import List @@ -31,17 +31,17 @@ query_router = APIRouter(prefix="/records", tags=["trinity-recording"]) -class _TaskRewardUpdate(BaseModel): - """One task-id group to reward-stamp and consume.""" +class _RecordKeyUpdate(BaseModel): + """One record-key group to reward-stamp and consume.""" - task_id: str + record_key: str reward: float run: int = 0 task: str = "" class _ConsumeTaskRequest(BaseModel): - updates: List[_TaskRewardUpdate] + updates: List[_RecordKeyUpdate] def _store(request: Request) -> RecordStore: @@ -58,8 +58,8 @@ def _recorder(request: Request) -> Recorder: return rec -async def _get_exp(store: RecordStore, task_id: str, request_id: str) -> Experience: - exp = await store.get_turn(task_id, request_id) +async def _get_exp(store: RecordStore, record_key: str, request_id: str) -> Experience: + exp = await store.get_turn(record_key, request_id) if exp is None: raise HTTPException(status_code=404, detail="experience not found") return exp @@ -67,40 +67,40 @@ async def _get_exp(store: RecordStore, task_id: str, request_id: str) -> Experie @query_router.get("/tasks") async def list_tasks(request: Request) -> dict: - """List all known task ids.""" + """List all known record keys.""" store = _store(request) - return {"task_ids": await store.list_tasks()} + return {"record_keys": await store.list_tasks()} -@query_router.get("/tasks/{task_id}") -async def get_task(task_id: str, request: Request) -> dict: - """Return lightweight metadata for all experiences of a task. +@query_router.get("/tasks/{record_key}") +async def get_task(record_key: str, request: Request) -> dict: + """Return lightweight metadata for all experiences of a record key. Tensor payloads (tokens/logprobs/routed_experts) are omitted here to keep listing responses small; fetch the per-turn blob endpoint for full data. """ store = _store(request) - experiences = await store.get_task(task_id) - return {"task_id": task_id, "turns": [e.to_dict() for e in experiences]} + experiences = await store.get_task(record_key) + return {"record_key": record_key, "turns": [e.to_dict() for e in experiences]} -@query_router.get("/tasks/{task_id}/turns/{request_id}") -async def get_turn(task_id: str, request_id: str, request: Request) -> dict: +@query_router.get("/tasks/{record_key}/turns/{request_id}") +async def get_turn(record_key: str, request_id: str, request: Request) -> dict: """Return lightweight metadata for a single experience.""" store = _store(request) - exp = await _get_exp(store, task_id, request_id) + exp = await _get_exp(store, record_key, request_id) return exp.to_dict() -@query_router.get("/tasks/{task_id}/turns/{request_id}/blob") -async def get_turn_blob(task_id: str, request_id: str, request: Request) -> Response: +@query_router.get("/tasks/{record_key}/turns/{request_id}/blob") +async def get_turn_blob(record_key: str, request_id: str, request: Request) -> Response: """Return the full experience as Trinity serialized bytes. Equivalent to ``Experience.serialize()``; deserializable via ``Experience.deserialize(...)`` or ``deserialize_many``. """ store = _store(request) - exp = await _get_exp(store, task_id, request_id) + exp = await _get_exp(store, record_key, request_id) return Response( content=Experience.serialize(exp), media_type="application/octet-stream", @@ -108,25 +108,25 @@ async def get_turn_blob(task_id: str, request_id: str, request: Request) -> Resp ) -@query_router.delete("/tasks/{task_id}") -async def delete_task(task_id: str, request: Request) -> dict: - """Delete all experiences for a task.""" +@query_router.delete("/tasks/{record_key}") +async def delete_task(record_key: str, request: Request) -> dict: + """Delete all experiences for a record key.""" store = _store(request) - await store.delete_task(task_id) - return {"task_id": task_id, "deleted": True} + await store.delete_task(record_key) + return {"record_key": record_key, "deleted": True} @query_router.post("/consume_task") async def consume_task(req: _ConsumeTaskRequest, request: Request) -> Response: - """Reward-stamp and pop the experiences of the given task-id groups. + """Reward-stamp and pop the experiences of the given record-key groups. Drains the recorder first (so every finished turn already made it into the store), then for each update sets ``reward``/``run``/``task`` on the whole matching group and pops it. Returns the union as Trinity serialized bytes (``Experience.serialize_many``), ready for the coordinator pipeline. - A task_id absent from this rank yields no experiences (it lived on another - rank); the coordinator fans this call out to every rank and merges. + A record_key absent from this rank yields no experiences (it lived on + another rank); the coordinator fans this call out to every rank and merges. """ store = _store(request) recorder = _recorder(request) @@ -136,8 +136,8 @@ async def consume_task(req: _ConsumeTaskRequest, request: Request) -> Response: exps: List[Experience] = [] for update in req.updates: exps.extend( - await store.update_reward_by_task_id( - task_id=update.task_id, + await store.update_reward_by_record_key( + record_key=update.record_key, reward=update.reward, run=update.run, task=update.task, diff --git a/trinity/common/models/vllm_patch/recording/recorder.py b/trinity/common/models/vllm_patch/recording/recorder.py index 942af6fe4ca..e7b9f95c667 100644 --- a/trinity/common/models/vllm_patch/recording/recorder.py +++ b/trinity/common/models/vllm_patch/recording/recorder.py @@ -24,8 +24,8 @@ from trinity.common.experience import Experience from trinity.common.models.vllm_patch.recording.context import ( + record_key_ctx, skip_recording_ctx, - task_id_ctx, ) from trinity.common.models.vllm_patch.recording.models import build_experience from trinity.common.models.vllm_patch.recording.store import RecordStore @@ -98,14 +98,14 @@ async def _patched_generate(*args, **kwargs): yield out if recorder.enabled and last is not None and getattr(last, "finished", False): - # Recover task id from the request's async context (set by - # RecordingIdentityMiddleware on the HTTP path, or by - # VLLMModel.chat on the Ray-direct path). None when neither was - # supplied; the store then falls back to request_id grouping. - task_id = task_id_ctx.get() + # Recover the record key from the request's async context (set by + # RecordingIdentityMiddleware on the HTTP path, or by VLLMModel.chat + # on the Ray-direct path). None when neither was supplied; the store + # then falls back to request_id grouping. + record_key = record_key_ctx.get() # Offload heavy serialization off the response critical path. The # task is tracked so ``flush`` can await it before a consume. - recorder.schedule_record(last, task_id) + recorder.schedule_record(last, record_key) setattr(_patched_generate, _PATCHED_FLAG, True) engine_client.generate = _patched_generate @@ -162,9 +162,9 @@ async def stop(self) -> None: self._flusher.cancel() self._flusher = None - def schedule_record(self, output, task_id: Optional[str]) -> None: + def schedule_record(self, output, record_key: Optional[str]) -> None: """Spawn (and track) a record task for a finished ``RequestOutput``.""" - task = asyncio.create_task(self._record(output, task_id)) + task = asyncio.create_task(self._record(output, record_key)) self._pending.add(task) task.add_done_callback(self._pending.discard) @@ -180,7 +180,7 @@ async def flush(self) -> None: if self._flusher is not None: await self._queue.join() - async def _record(self, output, task_id: Optional[str]) -> None: + async def _record(self, output, record_key: Optional[str]) -> None: """Build experiences for a finished turn and enqueue each for append.""" # Auxiliary forwards (logprobs recomputation, convert_messages) set # this to avoid polluting the store with 1-token degenerate turns. @@ -194,7 +194,7 @@ async def _record(self, output, task_id: Optional[str]) -> None: model_version = getattr(self.engine_client, _MODEL_VERSION_ATTR, None) exps = build_experience( output, - task_id, + record_key, rank=self.rank, timestamp=timestamp, model_version=model_version, diff --git a/trinity/common/models/vllm_patch/recording/store.py b/trinity/common/models/vllm_patch/recording/store.py index f745aa77459..e08a4d449be 100644 --- a/trinity/common/models/vllm_patch/recording/store.py +++ b/trinity/common/models/vllm_patch/recording/store.py @@ -2,25 +2,25 @@ A ``RecordStore`` persists Trinity ``Experience`` objects in the vLLM API server process. The only backend is ``MemoryStore`` — in-process, keyed by the -recording identity (``info["task_id"]`` = the API key / Ray-injected task id), -falling back to ``eid.suffix`` (the vLLM ``request_id``) when no identity was -supplied. +recording identity (``info["record_key"]`` = the API key / Ray-injected record +key), falling back to ``eid.suffix`` (the vLLM ``request_id``) when no identity +was supplied. -The consume side is ``update_reward_by_task_id``: it sets ``reward``/``run``/ -``task`` on every experience in a task-id group, pops the group, and returns +The consume side is ``update_reward_by_record_key``: it sets ``reward``/``run``/ +``task`` on every experience in a record-key group, pops the group, and returns it. This is the in-memory replacement for the old SQL ``HistoryRecorder``- mediated join — the coordinator calls it (via ``/records/consume_task``) at finalize time, so heavy experience bytes cross the network exactly once (store → coordinator pipeline) and never through Ray. Keying: experiences are identified by ``eid.suffix`` (the vLLM ``request_id``) -for traceability, but **grouped** by ``info["task_id"]`` so a whole task's -worth of turns/samples can be reward-updated and consumed together. +for traceability, but **grouped** by ``info["record_key"]`` so a whole reward +unit's worth of turns/samples can be reward-updated and consumed together. Concurrency: ``append_turn`` is called from a single background flusher task; -``update_reward_by_task_id`` is called from the ``/records/consume_task`` HTTP -handler. Both run in the same asyncio loop, so the dict is single-writer-safe -across these two without a lock. +``update_reward_by_record_key`` is called from the ``/records/consume_task`` +HTTP handler. Both run in the same asyncio loop, so the dict is single-writer- +safe across these two without a lock. """ import abc from collections import defaultdict @@ -30,8 +30,9 @@ #: Attribute carrying the vLLM request id on each experience's ``info`` dict. _REQUEST_ID_INFO_KEY = "request_id" -#: Attribute carrying the recording identity on each experience's ``info``. -_TASK_ID_INFO_KEY = "task_id" +#: Attribute carrying the recording identity (the MemoryStore group key) on +#: each experience's ``info``. +_RECORD_KEY_INFO_KEY = "record_key" class RecordStore(abc.ABC): @@ -42,13 +43,13 @@ async def append_turn(self, exp: Experience) -> None: """Persist one completed experience.""" @abc.abstractmethod - async def update_reward_by_task_id( - self, task_id: str, reward: float, run: int, task: str + async def update_reward_by_record_key( + self, record_key: str, reward: float, run: int, task: str ) -> list[Experience]: """Set reward/run/task on every experience in the group, pop and return it. Args: - task_id: The recording identity (group key). When the recorded + record_key: The recording identity (group key). When the recorded experience had no identity, this is its ``eid.suffix``. reward: Reward to stamp on every experience in the group. run: Run id to stamp on ``eid.run``. @@ -60,27 +61,27 @@ async def update_reward_by_task_id( """ @abc.abstractmethod - async def get_task(self, task_id: str) -> list[Experience]: - """Return all experiences for a task, in insertion order.""" + async def get_task(self, record_key: str) -> list[Experience]: + """Return all experiences for a record key, in insertion order.""" @abc.abstractmethod - async def get_turn(self, task_id: str, request_id: str) -> Optional[Experience]: + async def get_turn(self, record_key: str, request_id: str) -> Optional[Experience]: """Return a single experience, or None if not found.""" @abc.abstractmethod async def list_tasks(self) -> list[str]: - """Return all known task ids.""" + """Return all known record keys.""" @abc.abstractmethod - async def delete_task(self, task_id: str) -> None: - """Drop all experiences for a task.""" + async def delete_task(self, record_key: str) -> None: + """Drop all experiences for a record key.""" class MemoryStore(RecordStore): """In-process store. - Groups experiences by recording identity (``info["task_id"]``) when an API - key / task id was supplied, otherwise each turn is keyed by its own + Groups experiences by recording identity (``info["record_key"]``) when an + API key / record key was supplied, otherwise each turn is keyed by its own ``eid.suffix`` (request_id) — so a missing identity never collapses distinct turns. ``get_turn`` resolves an individual turn by ``info["request_id"]``. @@ -97,27 +98,27 @@ def __init__(self) -> None: @staticmethod def _group_key(exp: Experience) -> str: - session = exp.info.get(_TASK_ID_INFO_KEY) - return session if session else exp.eid.suffix + record_key = exp.info.get(_RECORD_KEY_INFO_KEY) + return record_key if record_key else exp.eid.suffix async def append_turn(self, exp: Experience) -> None: self._records[self._group_key(exp)].append(exp) - async def update_reward_by_task_id( - self, task_id: str, reward: float, run: int, task: str + async def update_reward_by_record_key( + self, record_key: str, reward: float, run: int, task: str ) -> list[Experience]: - exps = self._records.pop(task_id, []) + exps = self._records.pop(record_key, []) for exp in exps: exp.reward = reward exp.eid.run = run exp.eid.task = task return exps - async def get_task(self, task_id: str) -> list[Experience]: - return list(self._records.get(task_id, [])) + async def get_task(self, record_key: str) -> list[Experience]: + return list(self._records.get(record_key, [])) - async def get_turn(self, task_id: str, request_id: str) -> Optional[Experience]: - for exp in self._records.get(task_id, []): + async def get_turn(self, record_key: str, request_id: str) -> Optional[Experience]: + for exp in self._records.get(record_key, []): if exp.info.get(_REQUEST_ID_INFO_KEY) == request_id: return exp return None @@ -125,5 +126,5 @@ async def get_turn(self, task_id: str, request_id: str) -> Optional[Experience]: async def list_tasks(self) -> list[str]: return list(self._records.keys()) - async def delete_task(self, task_id: str) -> None: - self._records.pop(task_id, None) + async def delete_task(self, record_key: str) -> None: + self._records.pop(record_key, None) diff --git a/trinity/common/workflows/envs/alfworld/RAFT_utils.py b/trinity/common/workflows/envs/alfworld/RAFT_utils.py index 77c48af3dc5..8fe5c9f6b51 100644 --- a/trinity/common/workflows/envs/alfworld/RAFT_utils.py +++ b/trinity/common/workflows/envs/alfworld/RAFT_utils.py @@ -113,9 +113,9 @@ async def process_messages_to_experience_async(model, messages, info=None) -> Ex info = {} try: - # TODO(recording): when use_recorded_experience is on, replace this + # TODO(recording): when enable_recording is on, replace this # client-side conversion with a MemoryStore lookup by the session's - # task_id_key (concatenate turns via info["sample_index"]); see + # record_key (concatenate turns via info["sample_index"]); see # workflow.process_messages_to_experience. converted_experience = await model.convert_messages_to_experience_async(messages) diff --git a/trinity/common/workflows/workflow.py b/trinity/common/workflows/workflow.py index 43db800b5c0..e7670c79011 100644 --- a/trinity/common/workflows/workflow.py +++ b/trinity/common/workflows/workflow.py @@ -90,12 +90,12 @@ class Workflow: is_async: bool = False # whether the workflow runs in async mode. If true, `run_async()` must be implemented, else `run()` must be implemented. # When True, the workflow cooperates with the in-vLLM recorder: each chat - # call carries a per-sample ``task_id_key`` (the MemoryStore group key) so - # the runner can later report ``{task_id_key: reward}`` and the coordinator + # call carries a per-sample ``record_key`` (the MemoryStore group key) so + # the runner can later report ``{record_key: reward}`` and the coordinator # can join reward inside the store. Set by the WorkflowRunner from - # ``config.explorer.use_recorded_experience``. See ``SimpleWorkflow`` for - # the per-sample (n=1) loop this triggers. - use_recorded_experience: bool = False + # ``config.explorer.rollout_model.enable_recording``. See ``SimpleWorkflow`` + # for the per-sample (n=1) loop this triggers. + enable_recording: bool = False def __init__( self, @@ -227,14 +227,15 @@ def _build_experience_from_converted( def process_messages_to_experience( self, messages, reward, info={}, truncate_status=None ) -> Experience: - # TODO(recording): when use_recorded_experience is on, this client-side + # TODO(recording): when enable_recording is on, this client-side # conversion is redundant — the vLLM recorder's build_experience already # captured the authoritative heavy data (real logprobs without an extra # forward, real routed_experts) into the MemoryStore, keyed by the - # task_id the chat call carried. Replace this with an in-process lookup - # by task_id (store.get_task / consume), then concatenate the session's - # turns (info["sample_index"] orders them) into one experience here. - # Requires threading the per-call task_id_key down to this call site. + # record_key the chat call carried. Replace this with an in-process + # lookup by record_key (store.get_task / consume), then concatenate the + # session's turns (info["sample_index"] orders them) into one + # experience here. Requires threading the per-call record_key down to + # this call site. converted_experience = self.model.convert_messages_to_experience(messages) return self._build_experience_from_converted( converted_experience, @@ -247,7 +248,7 @@ async def process_messages_to_experience_async( self, messages, reward, info={}, truncate_status=None ) -> Experience: # TODO(recording): see process_messages_to_experience — replace with a - # MemoryStore lookup by task_id once task_id_key is threaded here. + # MemoryStore lookup by record_key once it is threaded here. converted_experience = await self.model.convert_messages_to_experience_async(messages) return self._build_experience_from_converted( converted_experience, @@ -308,7 +309,7 @@ def format_messages(self): return messages # -- recording-path helpers (shared by SimpleWorkflow / AsyncSimpleWorkflow) - - def _recorded_task_id_key(self, run_index: int) -> str: + def _record_key(self, run_index: int) -> str: """Per-sample recording identity (the MemoryStore group key).""" return f"{self.task.batch_id}/{self.task.task_id}/{self.run_id_base + run_index}" @@ -331,11 +332,11 @@ def _attach_rewards(self, responses, *, base: int) -> List[Experience]: return responses @staticmethod - def _stamp_task_id_key(exps: List[Experience], task_id_key: str) -> None: + def _stamp_record_key(exps: List[Experience], record_key: str) -> None: for exp in exps: if exp.info is None: exp.info = {} - exp.info["task_id_key"] = task_id_key + exp.info["record_key"] = record_key class SimpleWorkflow(BaseSimpleWorkflow): @@ -349,23 +350,23 @@ def run(self) -> List[Experience]: messages = self.format_messages() self.logger.debug("start chat") - if self.use_recorded_experience: + if self.enable_recording: return self._run_recorded(messages) responses = self.model.chat(messages, **self.rollout_args) return self._attach_rewards(responses, base=self.run_id_base) def _run_recorded(self, messages) -> List[Experience]: - # One chat call per sample (n=1) so each gets a distinct task_id_key + # One chat call per sample (n=1) so each gets a distinct record_key # (the recording group key == reward unit). The runner later reports - # {task_id_key: reward} and the coordinator joins reward in-store. + # {record_key: reward} and the coordinator joins reward in-store. rollout_args = dict(self.rollout_args) rollout_args["n"] = 1 exps: List[Experience] = [] for i in range(self.repeat_times): - task_id_key = self._recorded_task_id_key(i) - responses = self.model.chat(messages, task_id_key=task_id_key, **rollout_args) + record_key = self._record_key(i) + responses = self.model.chat(messages, record_key=record_key, **rollout_args) rewarded = self._attach_rewards(responses, base=self.run_id_base + i) - self._stamp_task_id_key(rewarded, task_id_key) + self._stamp_record_key(rewarded, record_key) exps.extend(rewarded) return exps @@ -378,7 +379,7 @@ async def run_async(self) -> List[Experience]: messages = self.format_messages() self.logger.info("start chat") - if self.use_recorded_experience: + if self.enable_recording: return await self._run_recorded_async(messages) responses = await self.model.chat_async(messages, **self.rollout_args) return self._attach_rewards(responses, base=self.run_id_base) @@ -388,12 +389,10 @@ async def _run_recorded_async(self, messages) -> List[Experience]: rollout_args["n"] = 1 exps: List[Experience] = [] for i in range(self.repeat_times): - task_id_key = self._recorded_task_id_key(i) - responses = await self.model.chat_async( - messages, task_id_key=task_id_key, **rollout_args - ) + record_key = self._record_key(i) + responses = await self.model.chat_async(messages, record_key=record_key, **rollout_args) rewarded = self._attach_rewards(responses, base=self.run_id_base + i) - self._stamp_task_id_key(rewarded, task_id_key) + self._stamp_record_key(rewarded, record_key) exps.extend(rewarded) return exps diff --git a/trinity/explorer/rollout_coordinator.py b/trinity/explorer/rollout_coordinator.py index b2fae011dc0..feed993a337 100644 --- a/trinity/explorer/rollout_coordinator.py +++ b/trinity/explorer/rollout_coordinator.py @@ -76,9 +76,9 @@ def __init__( # recording path's per-rank /records/consume_task fan-out. self._rank_urls: Optional[Dict[int, str]] = None - def _use_recorded_experience(self) -> bool: + def _enable_recording(self) -> bool: """Whether the recording-consume path is active for train batches.""" - return bool(self.config.explorer.use_recorded_experience) + return bool(self.config.explorer.rollout_model.enable_recording) def _resolve_rank_urls(self) -> Dict[int, str]: """Resolve each rollout engine's API server URL via named Ray actors. @@ -319,7 +319,7 @@ async def _finalize_train_batch( batch_state.state = BatchLifecycleState.FINALIZING try: - if self._use_recorded_experience(): + if self._enable_recording(): pipeline_metrics = await self._consume_recorded_experiences(payload_chunks) else: pipeline_metrics = await self.process_experiences(payload_chunks) @@ -335,9 +335,9 @@ async def _consume_recorded_experiences(self, payload_chunks: List[bytes]) -> di """Recording path: pull heavy experiences from each vLLM rank's store. ``payload_chunks`` are small pickle reward maps produced by the runners - (``{"engine_id": int, "updates": [{"task_id", "reward", "run", "task"}]}``). + (``{"engine_id": int, "updates": [{"record_key", "reward", "run", "task"}]}``). Group updates by engine, fan out ``POST /records/consume_task`` to each - rank (which drains its recorder, reward-stamps the matching task-id + rank (which drains its recorder, reward-stamps the matching record-key groups, pops them, and returns ``serialize_many`` bytes), deserialize, and feed the assembled experiences straight into the pipeline — no Ray serialization of heavy tensors, and reward is fused inside the store. @@ -372,7 +372,7 @@ async def _consume_recorded_experiences(self, payload_chunks: List[bytes]) -> di async def _post_consume_task( self, client: httpx.AsyncClient, rank_url: str, updates: List[dict] ) -> bytes: - """POST a batch of task-id reward updates to one rank; return heavy bytes.""" + """POST a batch of record-key reward updates to one rank; return heavy bytes.""" try: resp = await client.post( f"{rank_url}/records/consume_task", diff --git a/trinity/explorer/workflow_runner.py b/trinity/explorer/workflow_runner.py index b52f529bcc0..e1003366e74 100644 --- a/trinity/explorer/workflow_runner.py +++ b/trinity/explorer/workflow_runner.py @@ -118,11 +118,11 @@ def _create_workflow_instance(self, task: Task) -> Workflow: ) else: self.workflow_instance.reset(task) - self.workflow_instance.use_recorded_experience = self._use_recorded_experience() + self.workflow_instance.enable_recording = self._enable_recording() return self.workflow_instance - def _use_recorded_experience(self) -> bool: - return bool(self.config.explorer.use_recorded_experience) + def _enable_recording(self) -> bool: + return bool(self.config.explorer.rollout_model.enable_recording) async def _run_workflow(self, workflow_instance: Workflow) -> List[Experience]: if workflow_instance.asynchronous: @@ -140,7 +140,7 @@ def _create_isolated_workflow_instance(self, task: Task) -> Workflow: ), self.auxiliary_model_wrappers, ) - wf.use_recorded_experience = self._use_recorded_experience() + wf.enable_recording = self._enable_recording() return wf def _build_execution_result( @@ -430,14 +430,14 @@ async def run_task( # (TODO: wire eval consume-and-discard so eval turns don't leak # in the store). For now, return no payload. return status, b"" - elif self._use_recorded_experience(): + elif self._enable_recording(): # Recording path: ship only the small reward map keyed by the - # per-sample task_id_key the workflow stamped on each exp. The + # per-sample record_key the workflow stamped on each exp. The # heavy experiences live in the vLLM MemoryStore and are pulled # by the coordinator at finalize time. updates = [ { - "task_id": exp.info.get("task_id_key") or exp.eid.suffix, + "record_key": exp.info.get("record_key") or exp.eid.suffix, "reward": exp.reward, "run": exp.eid.run, "task": str(task.task_id), From 4f0f0a7eeb980a2b8a8df76f801c260707085007 Mon Sep 17 00:00:00 2001 From: "panxuchen.pxc" Date: Thu, 25 Jun 2026 17:05:54 +0800 Subject: [PATCH 08/84] unify vllm experience recording --- trinity/common/models/recording/__init__.py | 22 ++ trinity/common/models/recording/context.py | 47 ++++ .../{vllm_patch => }/recording/query.py | 58 +---- trinity/common/models/recording/recorder.py | 103 +++++++++ trinity/common/models/recording/store.py | 80 +++++++ trinity/common/models/vllm_model.py | 102 ++++----- .../models/vllm_patch/recording/__init__.py | 10 +- .../models/vllm_patch/recording/context.py | 73 ------ .../models/vllm_patch/recording/models.py | 87 ++++++-- .../models/vllm_patch/recording/recorder.py | 210 +++++++----------- .../models/vllm_patch/recording/server.py | 76 +++---- .../models/vllm_patch/recording/store.py | 130 ----------- 12 files changed, 493 insertions(+), 505 deletions(-) create mode 100644 trinity/common/models/recording/__init__.py create mode 100644 trinity/common/models/recording/context.py rename trinity/common/models/{vllm_patch => }/recording/query.py (54%) create mode 100644 trinity/common/models/recording/recorder.py create mode 100644 trinity/common/models/recording/store.py delete mode 100644 trinity/common/models/vllm_patch/recording/context.py delete mode 100644 trinity/common/models/vllm_patch/recording/store.py diff --git a/trinity/common/models/recording/__init__.py b/trinity/common/models/recording/__init__.py new file mode 100644 index 00000000000..a9c0b3f105e --- /dev/null +++ b/trinity/common/models/recording/__init__.py @@ -0,0 +1,22 @@ +"""Engine-agnostic generation recording utilities.""" + +from trinity.common.models.recording.context import ( + RecordingIdentityMiddleware, + extract_bearer_token, + get_recording_record_key, + record_key_ctx, + skip_recording_ctx, +) +from trinity.common.models.recording.recorder import Recorder +from trinity.common.models.recording.store import MemoryStore, RecordStore + +__all__ = [ + "MemoryStore", + "Recorder", + "RecordingIdentityMiddleware", + "RecordStore", + "extract_bearer_token", + "get_recording_record_key", + "record_key_ctx", + "skip_recording_ctx", +] diff --git a/trinity/common/models/recording/context.py b/trinity/common/models/recording/context.py new file mode 100644 index 00000000000..c31af625f58 --- /dev/null +++ b/trinity/common/models/recording/context.py @@ -0,0 +1,47 @@ +"""Per-request recording identity propagation shared by model engines.""" + +from contextvars import ContextVar +from typing import Any, Optional + +try: + from starlette.middleware.base import BaseHTTPMiddleware +except ModuleNotFoundError: + BaseHTTPMiddleware = object # type: ignore + +# The record key for the in-flight request (the MemoryStore group key), or None +# when the client did not send a supported identity header. +record_key_ctx: ContextVar[Optional[str]] = ContextVar("trinity_recording_record_key", default=None) + +# Set around auxiliary generate calls (logprobs recomputation, message +# conversion) so recorders skip them. +skip_recording_ctx: ContextVar[bool] = ContextVar("trinity_recording_skip", default=False) + +AUTHORIZATION_HEADER = "authorization" + + +def extract_bearer_token(authorization: Optional[str]) -> Optional[str]: + """Extract the bearer token from an Authorization header.""" + if not authorization: + return None + scheme, _, token = authorization.partition(" ") + if scheme.lower() != "bearer": + return None + token = token.strip() + return token or None + + +def get_recording_record_key(request: Any) -> Optional[str]: + """Return the recording identity for an HTTP request.""" + return extract_bearer_token(request.headers.get(AUTHORIZATION_HEADER)) + + +class RecordingIdentityMiddleware(BaseHTTPMiddleware): + """Capture request identity into ``record_key_ctx`` for the request lifetime.""" + + async def dispatch(self, request: Any, call_next): + record_key = get_recording_record_key(request) + token = record_key_ctx.set(record_key) + try: + return await call_next(request) + finally: + record_key_ctx.reset(token) diff --git a/trinity/common/models/vllm_patch/recording/query.py b/trinity/common/models/recording/query.py similarity index 54% rename from trinity/common/models/vllm_patch/recording/query.py rename to trinity/common/models/recording/query.py index 9e98772ed69..7d6e2de0cf5 100644 --- a/trinity/common/models/vllm_patch/recording/query.py +++ b/trinity/common/models/recording/query.py @@ -1,39 +1,21 @@ -"""HTTP query/consume endpoints over recorded experiences. - -Mounted on the main vLLM app (we own ``app`` in the bootstrap) via -``app.include_router(query_router)``. Routes are prefixed under ``/records`` to -avoid colliding with the OpenAI ``/v1/*`` surface — the OpenAI protocol is left -untouched. - -Listing endpoints return ``Experience.to_dict()`` (lightweight metadata, no -tensor payload). A single-turn endpoint returns the full experience as Trinity -serialized bytes (``Experience.serialize()``) so it can be fed straight back -into Trinity pipelines. - -``POST /records/consume_task`` is the consume path: it drains the recorder, -reward-stamps every experience in the matching record-key groups, pops them, -and returns the heavy experiences as ``Experience.serialize_many`` bytes — the -coordinator fans this out per rank at finalize time. -""" +"""HTTP query/consume endpoints over recorded experiences.""" + from typing import List from fastapi import APIRouter, HTTPException, Request, Response from pydantic import BaseModel from trinity.common.experience import Experience -from trinity.common.models.vllm_patch.recording.recorder import Recorder -from trinity.common.models.vllm_patch.recording.store import RecordStore +from trinity.common.models.recording.recorder import Recorder +from trinity.common.models.recording.store import RecordStore -#: Mounted on app.state by the bootstrap wiring in ``server.py``. -_STORE_ATTR = "trinity_record_store" -_RECORDER_ATTR = "trinity_recorder" +STORE_STATE_ATTR = "trinity_record_store" +RECORDER_STATE_ATTR = "trinity_recorder" query_router = APIRouter(prefix="/records", tags=["trinity-recording"]) class _RecordKeyUpdate(BaseModel): - """One record-key group to reward-stamp and consume.""" - record_key: str reward: float run: int = 0 @@ -45,14 +27,14 @@ class _ConsumeTaskRequest(BaseModel): def _store(request: Request) -> RecordStore: - store = getattr(request.app.state, _STORE_ATTR, None) + store = getattr(request.app.state, STORE_STATE_ATTR, None) if store is None: raise HTTPException(status_code=503, detail="recording store not configured") return store def _recorder(request: Request) -> Recorder: - rec = getattr(request.app.state, _RECORDER_ATTR, None) + rec = getattr(request.app.state, RECORDER_STATE_ATTR, None) if rec is None: raise HTTPException(status_code=503, detail="recorder not configured") return rec @@ -67,18 +49,12 @@ async def _get_exp(store: RecordStore, record_key: str, request_id: str) -> Expe @query_router.get("/tasks") async def list_tasks(request: Request) -> dict: - """List all known record keys.""" store = _store(request) return {"record_keys": await store.list_tasks()} @query_router.get("/tasks/{record_key}") async def get_task(record_key: str, request: Request) -> dict: - """Return lightweight metadata for all experiences of a record key. - - Tensor payloads (tokens/logprobs/routed_experts) are omitted here to keep - listing responses small; fetch the per-turn blob endpoint for full data. - """ store = _store(request) experiences = await store.get_task(record_key) return {"record_key": record_key, "turns": [e.to_dict() for e in experiences]} @@ -86,7 +62,6 @@ async def get_task(record_key: str, request: Request) -> dict: @query_router.get("/tasks/{record_key}/turns/{request_id}") async def get_turn(record_key: str, request_id: str, request: Request) -> dict: - """Return lightweight metadata for a single experience.""" store = _store(request) exp = await _get_exp(store, record_key, request_id) return exp.to_dict() @@ -94,11 +69,6 @@ async def get_turn(record_key: str, request_id: str, request: Request) -> dict: @query_router.get("/tasks/{record_key}/turns/{request_id}/blob") async def get_turn_blob(record_key: str, request_id: str, request: Request) -> Response: - """Return the full experience as Trinity serialized bytes. - - Equivalent to ``Experience.serialize()``; deserializable via - ``Experience.deserialize(...)`` or ``deserialize_many``. - """ store = _store(request) exp = await _get_exp(store, record_key, request_id) return Response( @@ -110,7 +80,6 @@ async def get_turn_blob(record_key: str, request_id: str, request: Request) -> R @query_router.delete("/tasks/{record_key}") async def delete_task(record_key: str, request: Request) -> dict: - """Delete all experiences for a record key.""" store = _store(request) await store.delete_task(record_key) return {"record_key": record_key, "deleted": True} @@ -118,19 +87,8 @@ async def delete_task(record_key: str, request: Request) -> dict: @query_router.post("/consume_task") async def consume_task(req: _ConsumeTaskRequest, request: Request) -> Response: - """Reward-stamp and pop the experiences of the given record-key groups. - - Drains the recorder first (so every finished turn already made it into the - store), then for each update sets ``reward``/``run``/``task`` on the whole - matching group and pops it. Returns the union as Trinity serialized bytes - (``Experience.serialize_many``), ready for the coordinator pipeline. - - A record_key absent from this rank yields no experiences (it lived on - another rank); the coordinator fans this call out to every rank and merges. - """ store = _store(request) recorder = _recorder(request) - # Ensure in-flight record tasks have been appended before we pop. await recorder.flush() exps: List[Experience] = [] diff --git a/trinity/common/models/recording/recorder.py b/trinity/common/models/recording/recorder.py new file mode 100644 index 00000000000..1515f11a7b2 --- /dev/null +++ b/trinity/common/models/recording/recorder.py @@ -0,0 +1,103 @@ +"""Engine-agnostic background recorder for generated experiences.""" + +import asyncio +import logging +from collections.abc import Callable, Sequence +from datetime import datetime, timezone +from typing import Any, Optional + +from trinity.common.experience import Experience +from trinity.common.models.recording.context import skip_recording_ctx +from trinity.common.models.recording.store import RecordStore + +MODEL_VERSION_ATTR = "trinity_model_version" + +BuildExperiencesFn = Callable[..., Sequence[Experience]] + + +class Recorder: + """Drains finished turns into a ``RecordStore`` from a background task. + + Engine-specific code supplies ``build_experiences``, which converts a + finished engine output object into Trinity ``Experience`` instances. + """ + + def __init__( + self, + store: RecordStore, + *, + build_experiences: BuildExperiencesFn, + enabled: bool, + rank: int = 0, + engine_client: Any = None, + ) -> None: + self.store = store + self.enabled = enabled + self.rank = rank + self.engine_client = engine_client + self._build_experiences = build_experiences + self._queue: "asyncio.Queue[Optional[Experience]]" = asyncio.Queue() + self._flusher: Optional[asyncio.Task] = None + self._pending: "set[asyncio.Task]" = set() + + def start(self) -> None: + """Start the background flusher. Idempotent.""" + if self._flusher is not None or not self.enabled: + return + self._flusher = asyncio.create_task(self._flush_loop()) + + async def stop(self) -> None: + """Drain in-flight + queued turns, then stop the flusher.""" + if self._flusher is None: + return + await self.flush() + self._flusher.cancel() + self._flusher = None + + def schedule_record(self, output: Any, record_key: Optional[str], **builder_kwargs) -> None: + """Spawn and track a record task for a finished engine output.""" + task = asyncio.create_task(self._record(output, record_key, **builder_kwargs)) + self._pending.add(task) + task.add_done_callback(self._pending.discard) + + async def flush(self) -> None: + """Wait until every in-flight record has been appended to the store.""" + if self._pending: + await asyncio.gather(*self._pending, return_exceptions=True) + if self._flusher is not None: + await self._queue.join() + + async def _record(self, output: Any, record_key: Optional[str], **builder_kwargs) -> None: + if skip_recording_ctx.get(): + return + timestamp = datetime.now(timezone.utc).isoformat() + model_version = getattr(self.engine_client, MODEL_VERSION_ATTR, None) + exps = self._build_experiences( + output, + record_key, + rank=self.rank, + timestamp=timestamp, + model_version=model_version, + **builder_kwargs, + ) + for exp in exps: + await self._queue.put(exp) + + async def _flush_loop(self) -> None: + while True: + exp = await self._queue.get() + try: + if exp is None: + return + await self._safe_append(exp) + finally: + self._queue.task_done() + + async def _safe_append(self, exp: Experience) -> None: + try: + await self.store.append_turn(exp) + except Exception: + logging.getLogger(__name__).exception( + "recording store.append_turn failed for request %s", + exp.info.get("request_id"), + ) diff --git a/trinity/common/models/recording/store.py b/trinity/common/models/recording/store.py new file mode 100644 index 00000000000..ee7a764e367 --- /dev/null +++ b/trinity/common/models/recording/store.py @@ -0,0 +1,80 @@ +"""Pluggable storage backends for recorded experiences.""" + +import abc +from collections import defaultdict +from typing import Optional + +from trinity.common.experience import Experience + +REQUEST_ID_INFO_KEY = "request_id" +RECORD_KEY_INFO_KEY = "record_key" + + +class RecordStore(abc.ABC): + """Abstract persistence interface for recorded experiences.""" + + @abc.abstractmethod + async def append_turn(self, exp: Experience) -> None: + """Persist one completed experience.""" + + @abc.abstractmethod + async def update_reward_by_record_key( + self, record_key: str, reward: float, run: int, task: str + ) -> list[Experience]: + """Set reward/run/task on every experience in the group, pop and return it.""" + + @abc.abstractmethod + async def get_task(self, record_key: str) -> list[Experience]: + """Return all experiences for a record key, in insertion order.""" + + @abc.abstractmethod + async def get_turn(self, record_key: str, request_id: str) -> Optional[Experience]: + """Return a single experience, or None if not found.""" + + @abc.abstractmethod + async def list_tasks(self) -> list[str]: + """Return all known record keys.""" + + @abc.abstractmethod + async def delete_task(self, record_key: str) -> None: + """Drop all experiences for a record key.""" + + +class MemoryStore(RecordStore): + """In-process store grouped by recording identity.""" + + def __init__(self) -> None: + self._records: dict[str, list[Experience]] = defaultdict(list) + + @staticmethod + def _group_key(exp: Experience) -> str: + record_key = exp.info.get(RECORD_KEY_INFO_KEY) + return record_key if record_key else exp.eid.suffix + + async def append_turn(self, exp: Experience) -> None: + self._records[self._group_key(exp)].append(exp) + + async def update_reward_by_record_key( + self, record_key: str, reward: float, run: int, task: str + ) -> list[Experience]: + exps = self._records.pop(record_key, []) + for exp in exps: + exp.reward = reward + exp.eid.run = run + exp.eid.task = task + return exps + + async def get_task(self, record_key: str) -> list[Experience]: + return list(self._records.get(record_key, [])) + + async def get_turn(self, record_key: str, request_id: str) -> Optional[Experience]: + for exp in self._records.get(record_key, []): + if exp.info.get(REQUEST_ID_INFO_KEY) == request_id: + return exp + return None + + async def list_tasks(self) -> list[str]: + return list(self._records.keys()) + + async def delete_task(self, record_key: str) -> None: + self._records.pop(record_key, None) diff --git a/trinity/common/models/vllm_model.py b/trinity/common/models/vllm_model.py index 6011e421dde..116423dca56 100644 --- a/trinity/common/models/vllm_model.py +++ b/trinity/common/models/vllm_model.py @@ -13,16 +13,11 @@ from trinity.common.config import InferenceModelConfig from trinity.common.constants import SyncMethod from trinity.common.experience import Experience -from trinity.common.models.mm_utils import ( - combine_output_token_ids, - vLLMMultiModalRender, -) +from trinity.common.models.mm_utils import vLLMMultiModalRender from trinity.common.models.model import BaseInferenceModel +from trinity.common.models.recording.context import record_key_ctx, skip_recording_ctx from trinity.common.models.vllm_patch import get_vllm_version -from trinity.common.models.vllm_patch.recording.context import ( - record_key_ctx, - skip_recording_ctx, -) +from trinity.common.models.vllm_patch.recording.models import build_experience # V0 engine is deprecated since vLLM v0.10.2, related code will be removed in the future. @@ -92,6 +87,7 @@ def __init__( self.api_server_host = None self.api_server_port = None self.api_server = None + self.recorder = None self._prepared = False self.async_llm = None self.headless_executor = None @@ -195,6 +191,19 @@ async def prepare(self) -> None: # attribute experiences to the right policy without an extra # launch-time parameter. Updated in sync_model_weights. self.async_llm.trinity_model_version = self.model_version + if self.config.enable_recording: + from trinity.common.models.vllm_patch.recording.recorder import ( + TRINITY_MM_RENDER_ATTR, + create_vllm_recorder, + ) + + if self.mm_render is None: + self.mm_render = vLLMMultiModalRender( + model_path=self.config.model_path, # type: ignore + ) + setattr(self.async_llm, TRINITY_MM_RENDER_ATTR, self.mm_render) + self.recorder = create_vllm_recorder(self.async_llm, self.logger) + self.recorder.start() await self._collective_rpc("apply_patches") await self.run_api_server() else: @@ -251,27 +260,6 @@ async def chat( prompt=prompt, lora_request=lora_request, record_key=record_key, **kwargs ) - def _extract_routed_experts(self, output: Any, output_index: int) -> Optional[torch.Tensor]: - if not self.config.enable_return_routed_experts: - return None - - routed_experts_parts = [] - prompt_routed_experts = getattr(output, "prompt_routed_experts", None) - if prompt_routed_experts is not None: - routed_experts_parts.append(torch.as_tensor(prompt_routed_experts, dtype=torch.uint8)) - - completion_routed_experts = getattr(output.outputs[output_index], "routed_experts", None) - if completion_routed_experts is not None: - routed_experts_parts.append( - torch.as_tensor(completion_routed_experts, dtype=torch.uint8) - ) - - if not routed_experts_parts: - return None - if len(routed_experts_parts) == 1: - return routed_experts_parts[0] - return torch.cat(routed_experts_parts, dim=0) - async def generate( self, prompt: Union[str, Dict], @@ -324,36 +312,19 @@ async def generate( input_ids=output.prompt_token_ids, multi_modal_data=prompt.get("multi_modal_data", {}), ) - experiences = [ - Experience( - tokens=torch.cat( - ( - torch.tensor(output.prompt_token_ids, dtype=torch.int32), - torch.tensor(output.outputs[i].token_ids, dtype=torch.int32), - ) - ), - logprobs=torch.cat( - ( - torch.tensor( - [ - list(logprob_dict.values())[0].logprob - for logprob_dict in output.outputs[i].logprobs - ], - dtype=torch.float32, - ), - ) - ), - prompt_length=len(output.prompt_token_ids), - prompt_text=self.tokenizer.decode(output.prompt_token_ids), - response_text=output.outputs[i].text, - multi_modal_inputs=combine_output_token_ids( - output.outputs[i].token_ids, multi_modal_inputs - ), - routed_experts=self._extract_routed_experts(output, i), - ) - for i in range(len(output.outputs)) - ] - return experiences + if self.tokenizer is None: + await self._initialize_tokenizer() + return build_experience( + output, + record_key=None, + rank=0, + timestamp="", + multi_modal_inputs=multi_modal_inputs, + prompt_text=self.tokenizer.decode(output.prompt_token_ids), + include_recording_info=False, + include_routed_experts=self.config.enable_return_routed_experts, + include_prompt_routed_experts=True, + ) async def logprobs( # type: ignore [override] self, @@ -543,14 +514,18 @@ async def _generate_internal(self, prompt: Any, lora_request=None, **kwargs) -> **generate_kwargs, ) - # Consume the stream until the request is finished. + # Consume the stream to completion so engine-level recording runs only + # after the full generation stream has ended. + finished_output = None async for request_output in stream: if request_output.finished: # Bypass the original full prompt. # request_output.prompt = request.prompt - return request_output + finished_output = request_output - raise RuntimeError("[vLLM] The request is not finished. This should not happen.") + if finished_output is None: + raise RuntimeError("[vLLM] The request is not finished. This should not happen.") + return finished_output async def shutdown(self): """Shutdown the vLLM v1 engine. This kills child processes forked @@ -565,6 +540,9 @@ async def shutdown(self): except asyncio.CancelledError: pass self.api_server = None + if self.recorder is not None: + await self.recorder.stop() + self.recorder = None if self.headless_executor is not None: self.logger.info("Shutting down headless executor") self.headless_executor.shutdown() diff --git a/trinity/common/models/vllm_patch/recording/__init__.py b/trinity/common/models/vllm_patch/recording/__init__.py index a8ec841bef6..32eef1fc0dd 100644 --- a/trinity/common/models/vllm_patch/recording/__init__.py +++ b/trinity/common/models/vllm_patch/recording/__init__.py @@ -1,21 +1,22 @@ """Generation-recording patch for the vLLM OpenAI server. Designed for vllm >= 0.23.0. """ -from trinity.common.models.vllm_patch.recording.context import ( +from trinity.common.models.recording.context import ( RecordingIdentityMiddleware, record_key_ctx, skip_recording_ctx, ) +from trinity.common.models.recording.query import query_router +from trinity.common.models.recording.recorder import Recorder +from trinity.common.models.recording.store import MemoryStore, RecordStore from trinity.common.models.vllm_patch.recording.models import build_experience -from trinity.common.models.vllm_patch.recording.query import query_router from trinity.common.models.vllm_patch.recording.recorder import ( - Recorder, + create_vllm_recorder, patch_engine_for_recording, ) from trinity.common.models.vllm_patch.recording.server import ( run_api_server_with_recording, ) -from trinity.common.models.vllm_patch.recording.store import MemoryStore, RecordStore __all__ = [ "MemoryStore", @@ -23,6 +24,7 @@ "RecordingIdentityMiddleware", "Recorder", "build_experience", + "create_vllm_recorder", "patch_engine_for_recording", "query_router", "record_key_ctx", diff --git a/trinity/common/models/vllm_patch/recording/context.py b/trinity/common/models/vllm_patch/recording/context.py deleted file mode 100644 index 8e43d98714f..00000000000 --- a/trinity/common/models/vllm_patch/recording/context.py +++ /dev/null @@ -1,73 +0,0 @@ -"""Per-request recording identity propagation. - -The OpenAI ``Authorization: Bearer `` header is the recording identity -source because it works with CLI agents that do not support custom headers. We -read the identity in an in-process ASGI middleware and stash it in a contextvar -so the engine-level wrapper (which runs in the same async task as the serving -handler) can recover it at record time. - -The identity value is the **record key** — the MemoryStore group key under -which a whole reward unit's worth of turns (one sample / one episode) is -recorded, and the key the runner later reports ``{record_key: reward}`` against -so the coordinator can join the reward inside the store. - -No identity header on a request is fine: the recorder falls back to -``request_id`` as the record key so nothing is silently dropped. -""" -from contextvars import ContextVar -from typing import Optional - -from starlette.middleware.base import BaseHTTPMiddleware -from starlette.requests import Request - -# The record key for the in-flight request (the MemoryStore group key), or None -# when the client did not send a supported identity header (the recorder then -# uses request_id as a fallback record key). -record_key_ctx: ContextVar[Optional[str]] = ContextVar("trinity_recording_record_key", default=None) - -# Set around auxiliary engine.generate calls (logprobs recomputation, -# convert_messages_to_experience) so the recorder skips them — those 1-token -# forwards are not real turns and would pollute the store. -skip_recording_ctx: ContextVar[bool] = ContextVar("trinity_recording_skip", default=False) - -#: Preferred identity header for OpenAI-compatible clients. -AUTHORIZATION_HEADER = "authorization" - - -def extract_bearer_token(authorization: Optional[str]) -> Optional[str]: - """Extract the bearer token from an Authorization header. - - Returns None when the header is missing or does not use the Bearer scheme. - """ - if not authorization: - return None - scheme, _, token = authorization.partition(" ") - if scheme.lower() != "bearer": - return None - token = token.strip() - return token or None - - -def get_recording_record_key(request: Request) -> Optional[str]: - """Return the recording identity for a request. - - OpenAI-compatible API keys are used because every supported agent platform - can pass them. - """ - return extract_bearer_token(request.headers.get(AUTHORIZATION_HEADER)) - - -class RecordingIdentityMiddleware(BaseHTTPMiddleware): - """Capture request identity into ``record_key_ctx`` for the request's lifetime. - - Runs in-process (ASGI) — no extra network hop, no serialization cost beyond - a contextvar set/reset. - """ - - async def dispatch(self, request: Request, call_next): - record_key = get_recording_record_key(request) - token = record_key_ctx.set(record_key) - try: - return await call_next(request) - finally: - record_key_ctx.reset(token) diff --git a/trinity/common/models/vllm_patch/recording/models.py b/trinity/common/models/vllm_patch/recording/models.py index 99f558d8f25..b6817c40d28 100644 --- a/trinity/common/models/vllm_patch/recording/models.py +++ b/trinity/common/models/vllm_patch/recording/models.py @@ -38,7 +38,10 @@ """ from typing import Any, List, Optional +import torch + from trinity.common.experience import EID, Experience +from trinity.common.models.mm_utils import combine_output_token_ids def _extract_chosen_logprobs( @@ -79,6 +82,33 @@ def _extract_chosen_logprobs( return [float(sample_logprobs[i][tid].logprob) for i, tid in enumerate(response_token_ids)] +def _extract_routed_experts( + output: Any, + completion: Any, + *, + include_routed_experts: bool, + include_prompt_routed_experts: bool, +): + if not include_routed_experts: + return None + + routed_experts_parts = [] + if include_prompt_routed_experts: + prompt_routed_experts = getattr(output, "prompt_routed_experts", None) + if prompt_routed_experts is not None: + routed_experts_parts.append(torch.as_tensor(prompt_routed_experts, dtype=torch.uint8)) + + completion_routed_experts = getattr(completion, "routed_experts", None) + if completion_routed_experts is not None: + routed_experts_parts.append(torch.as_tensor(completion_routed_experts, dtype=torch.uint8)) + + if not routed_experts_parts: + return None + if len(routed_experts_parts) == 1: + return routed_experts_parts[0] + return torch.cat(routed_experts_parts, dim=0) + + def build_experience( output: Any, record_key: Optional[str], @@ -87,6 +117,11 @@ def build_experience( timestamp: str, endpoint: str = "unknown", model_version: Optional[int] = None, + multi_modal_inputs: Optional[dict] = None, + prompt_text: Optional[str] = None, + include_recording_info: bool = True, + include_routed_experts: bool = True, + include_prompt_routed_experts: bool = False, ) -> List[Experience]: """Build Trinity ``Experience`` objects from a finished ``RequestOutput``. @@ -105,6 +140,19 @@ def build_experience( endpoint: Which OpenAI endpoint served the turn (best-effort). model_version: Checkpoint version the serving policy was at; stamped into ``info`` for RL attribution (read in-actor by the recorder). + multi_modal_inputs: Optional training-time multimodal tensors aligned + with the prompt tokens. Response token type ids are appended per + completion before storing on the ``Experience``. + prompt_text: Optional prompt text override. Direct model calls can pass + tokenizer-decoded prompt text when ``RequestOutput.prompt`` is not + suitable for training records. + include_recording_info: Whether to attach recording metadata and + ``EID(suffix=request_id)``. Direct generate return values can turn + this off to preserve normal Experience construction semantics. + include_routed_experts: Whether routed experts should be copied. + include_prompt_routed_experts: Whether to prepend prompt routed experts + to completion routed experts. Direct generate uses this to match + its full-token training representation. Returns: One ``Experience`` per non-degenerate completion. Empty list if the @@ -134,27 +182,40 @@ def build_experience( prompt_length = len(prompt_token_ids) chosen_logprobs = _extract_chosen_logprobs(completion.logprobs, response_token_ids) - routed_experts = completion.routed_experts - - info = { - "request_id": request_id, - "record_key": record_key, - "sample_index": sample_index, - "rank": rank, - "timestamp": timestamp, - "endpoint": endpoint, - "model_version": model_version, - } + routed_experts = _extract_routed_experts( + output, + completion, + include_routed_experts=include_routed_experts, + include_prompt_routed_experts=include_prompt_routed_experts, + ) + + info = None + eid = None + if include_recording_info: + eid = EID(suffix=request_id) + info = { + "request_id": request_id, + "record_key": record_key, + "sample_index": sample_index, + "rank": rank, + "timestamp": timestamp, + "endpoint": endpoint, + "model_version": model_version, + } experiences.append( Experience( - eid=EID(suffix=request_id), + eid=eid, tokens=tokens, logprobs=chosen_logprobs, prompt_length=prompt_length, routed_experts=routed_experts, - prompt_text=output.prompt, + prompt_text=prompt_text if prompt_text is not None else output.prompt, response_text=completion.text, + multi_modal_inputs=combine_output_token_ids( + response_token_ids, + multi_modal_inputs, + ), info=info, ) ) diff --git a/trinity/common/models/vllm_patch/recording/recorder.py b/trinity/common/models/vllm_patch/recording/recorder.py index e7b9f95c667..585feca4b02 100644 --- a/trinity/common/models/vllm_patch/recording/recorder.py +++ b/trinity/common/models/vllm_patch/recording/recorder.py @@ -17,18 +17,14 @@ computation — the client response is unchanged unless the client itself requested logprobs. Recording stays transparent. """ -import asyncio import functools import logging from typing import Optional -from trinity.common.experience import Experience -from trinity.common.models.vllm_patch.recording.context import ( - record_key_ctx, - skip_recording_ctx, -) +from trinity.common.models.recording.context import record_key_ctx +from trinity.common.models.recording.recorder import Recorder +from trinity.common.models.recording.store import MemoryStore, RecordStore from trinity.common.models.vllm_patch.recording.models import build_experience -from trinity.common.models.vllm_patch.recording.store import RecordStore #: Guard attribute marking the wrapped generate, mirroring api_patch_v17 style. _PATCHED_FLAG = "__patched_engine_recording__" @@ -44,6 +40,73 @@ #: to thread a knob through the launcher. The engine's ``max_logprobs`` cap #: (default 20, set at engine build) already covers it. _RECORDER_LOGPROB_WIDTH = 1 +TRINITY_RECORDER_ATTR = "trinity_recorder" +TRINITY_RECORD_STORE_ATTR = "trinity_record_store" +TRINITY_MM_RENDER_ATTR = "trinity_mm_render" + + +def _get_api_process_rank(engine_client) -> int: + try: + return int(engine_client.vllm_config.parallel_config._api_process_rank) + except Exception: + return 0 + + +def create_vllm_recorder( + engine_client, + logger: logging.Logger, + *, + store: Optional[RecordStore] = None, + enabled: bool = True, +) -> Recorder: + """Create and install a vLLM-backed recorder on ``engine_client``.""" + existing = getattr(engine_client, TRINITY_RECORDER_ATTR, None) + if existing is not None: + return existing + + recorder = Recorder( + store=store or MemoryStore(), + build_experiences=build_experience, + enabled=enabled, + rank=_get_api_process_rank(engine_client), + engine_client=engine_client, + ) + patch_engine_for_recording(engine_client, recorder, logger) + setattr(engine_client, TRINITY_RECORDER_ATTR, recorder) + setattr(engine_client, TRINITY_RECORD_STORE_ATTR, recorder.store) + return recorder + + +def _get_prompt_arg(args, kwargs): + if "prompt" in kwargs: + return kwargs["prompt"] + if args: + return args[0] + return None + + +def _build_multi_modal_inputs(engine_client, prompt, output, logger: logging.Logger): + if not isinstance(prompt, dict): + return None + multi_modal_data = prompt.get("multi_modal_data") + if not multi_modal_data: + return None + mm_render = getattr(engine_client, TRINITY_MM_RENDER_ATTR, None) + if mm_render is None: + logger.warning( + "Recording saw a multimodal vLLM prompt but no %s is attached to engine_client; " + "recorded Experience will not include multi_modal_inputs.", + TRINITY_MM_RENDER_ATTR, + ) + return None + try: + return mm_render.build_mm_input_for_training( + input_ids=output.prompt_token_ids, + multi_modal_data=multi_modal_data, + ) + except Exception: + logger.exception("Failed to build multi_modal_inputs for recorded vLLM Experience") + return None def patch_engine_for_recording( @@ -81,6 +144,7 @@ async def _patched_generate(*args, **kwargs): sampling_params = kwargs.get("sampling_params") if sampling_params is None and len(args) >= 2: sampling_params = args[1] + prompt = _get_prompt_arg(args, kwargs) if recorder.enabled and sampling_params is not None: # Ensure logprobs are computed for recording (callers may omit @@ -100,126 +164,22 @@ async def _patched_generate(*args, **kwargs): if recorder.enabled and last is not None and getattr(last, "finished", False): # Recover the record key from the request's async context (set by # RecordingIdentityMiddleware on the HTTP path, or by VLLMModel.chat - # on the Ray-direct path). None when neither was supplied; the store - # then falls back to request_id grouping. + # on the Ray-direct path). A missing key means the caller did not + # opt into grouping this turn, so skip recording entirely. record_key = record_key_ctx.get() - # Offload heavy serialization off the response critical path. The - # task is tracked so ``flush`` can await it before a consume. - recorder.schedule_record(last, record_key) + if record_key is not None: + multi_modal_inputs = _build_multi_modal_inputs( + engine_client, + prompt, + last, + logger, + ) + recorder.schedule_record( + last, + record_key, + multi_modal_inputs=multi_modal_inputs, + ) setattr(_patched_generate, _PATCHED_FLAG, True) engine_client.generate = _patched_generate logger.info("Patched vLLM engine_client.generate for generation recording") - - -class Recorder: - """Drains finished turns into a ``RecordStore`` from a background task. - - Putting records into an ``asyncio.Queue`` and flushing from a single worker - keeps the response path cheap (record == enqueue) and serializes expensive - payloads (ndarray -> .npy, json) off the serving hot loop. - - ``schedule_record`` spawns a task per finished ``RequestOutput`` and tracks - it in ``_pending``; ``flush`` awaits all of them plus ``queue.join`` so a - caller (``/records/consume_task``) sees a quiesced store before popping. - """ - - def __init__( - self, - store: RecordStore, - *, - enabled: bool, - rank: int = 0, - engine_client=None, - ) -> None: - self.store = store - self.enabled = enabled - self.rank = rank - # The engine_client is the same AsyncLLM instance VLLMModel updates in - # sync_model_weights (``.trinity_model_version``), so we read the live - # checkpoint version off it at record time. - self.engine_client = engine_client - self._queue: "asyncio.Queue[Optional[Experience]]" = asyncio.Queue() - self._flusher: Optional[asyncio.Task] = None - # In-flight ``_record`` tasks spawned by ``schedule_record``. Tracked - # so ``flush`` can await them — a record task that hasn't put yet would - # otherwise race a consume (the experience would be missing from the - # store even though the request already finished). - self._pending: "set[asyncio.Task]" = set() - - def start(self) -> None: - """Start the background flusher. Idempotent.""" - if self._flusher is not None or not self.enabled: - return - self._flusher = asyncio.create_task(self._flush_loop()) - - async def stop(self) -> None: - """Drain in-flight + queued turns, then stop the flusher.""" - if self._flusher is None: - return - # Drain everything put in flight so we don't lose turns at shutdown. - await self.flush() - self._flusher.cancel() - self._flusher = None - - def schedule_record(self, output, record_key: Optional[str]) -> None: - """Spawn (and track) a record task for a finished ``RequestOutput``.""" - task = asyncio.create_task(self._record(output, record_key)) - self._pending.add(task) - task.add_done_callback(self._pending.discard) - - async def flush(self) -> None: - """Wait until every in-flight record has been appended to the store. - - Awaits all pending ``_record`` tasks (so every finished turn has been - enqueued), then ``queue.join``s (so the flusher has appended them all). - Call this before consuming the store to avoid reading a partial state. - """ - if self._pending: - await asyncio.gather(*self._pending, return_exceptions=True) - if self._flusher is not None: - await self._queue.join() - - async def _record(self, output, record_key: Optional[str]) -> None: - """Build experiences for a finished turn and enqueue each for append.""" - # Auxiliary forwards (logprobs recomputation, convert_messages) set - # this to avoid polluting the store with 1-token degenerate turns. - if skip_recording_ctx.get(): - return - # Stamp now (real runtime, not a workflow sandbox): permitted here. - from datetime import datetime, timezone - - timestamp = datetime.now(timezone.utc).isoformat() - # Read the live checkpoint version the actor mirrors onto the engine. - model_version = getattr(self.engine_client, _MODEL_VERSION_ATTR, None) - exps = build_experience( - output, - record_key, - rank=self.rank, - timestamp=timestamp, - model_version=model_version, - ) - for exp in exps: - await self._queue.put(exp) - - async def _flush_loop(self) -> None: - while True: - exp = await self._queue.get() - try: - if exp is None: - # Sentinel for graceful shutdown. - return - await self._safe_append(exp) - finally: - # Paired with put() so queue.join() in flush() can complete. - self._queue.task_done() - - async def _safe_append(self, exp: Experience) -> None: - try: - await self.store.append_turn(exp) - except Exception: - # Never let a storage hiccup crash the flusher loop. - logging.getLogger(__name__).exception( - "recording store.append_turn failed for request %s", - exp.info.get("request_id"), - ) diff --git a/trinity/common/models/vllm_patch/recording/server.py b/trinity/common/models/vllm_patch/recording/server.py index a9c304668a7..a33a79a2afb 100644 --- a/trinity/common/models/vllm_patch/recording/server.py +++ b/trinity/common/models/vllm_patch/recording/server.py @@ -7,11 +7,8 @@ and you get the standard vLLM OpenAI server *plus* generation recording, with no edits to vLLM source or to ``api_patch_v17.py``. -Recording wiring (all applied between ``build_app`` and ``serve_http`` because -we own both ``app`` and ``engine_client`` at that point): - 1. ``patch_engine_for_recording`` — instance-level wrap of - ``engine_client.generate`` to force top-k logprobs and record finished - ``RequestOutput`` (covers chat/completion/responses, streaming and not). +Recording wiring: + 1. ``vLLMRolloutModel`` owns the recorder and attaches it to ``async_llm``. 2. ``RecordingIdentityMiddleware`` — in-process ASGI middleware reading ``Authorization: Bearer `` into a contextvar. 3. ``query_router`` — ``/records/*`` endpoints for later analysis. @@ -43,21 +40,21 @@ from vllm.utils.system_utils import set_ulimit from vllm.version import __version__ as VLLM_VERSION -from trinity.common.models.vllm_patch import get_vllm_version -from trinity.common.models.vllm_patch.recording.context import ( - RecordingIdentityMiddleware, -) -from trinity.common.models.vllm_patch.recording.query import query_router -from trinity.common.models.vllm_patch.recording.recorder import ( - Recorder, - patch_engine_for_recording, +from trinity.common.models.recording.context import RecordingIdentityMiddleware +from trinity.common.models.recording.query import ( + RECORDER_STATE_ATTR, + STORE_STATE_ATTR, + query_router, ) -from trinity.common.models.vllm_patch.recording.store import MemoryStore, RecordStore +from trinity.common.models.recording.recorder import Recorder +from trinity.common.models.recording.store import RecordStore +from trinity.common.models.vllm_patch import get_vllm_version +from trinity.common.models.vllm_patch.recording.recorder import TRINITY_RECORDER_ATTR #: Attribute on app.state holding the active RecordStore. -_STORE_STATE_ATTR = "trinity_record_store" +_STORE_STATE_ATTR = STORE_STATE_ATTR #: Attribute on app.state holding the active Recorder. -_RECORDER_STATE_ATTR = "trinity_recorder" +_RECORDER_STATE_ATTR = RECORDER_STATE_ATTR def setup_server_in_ray(args, logger): @@ -108,15 +105,14 @@ def dummy_add_signal_handler(self, *args, **kwargs): def _setup_recording( - args, engine_client, app, logger, ) -> Recorder: """Wire generation recording onto the in-construction server. - Returns the started Recorder (for lifecycle management). This is only - called when recording is on, so there is no disable switch here. + Returns the Recorder owned by ``vLLMRolloutModel``. This is only called when + recording is on, so there is no disable switch here. No static config is threaded in: the chosen-token logprob width is a constant inside the recorder (``_RECORDER_LOGPROB_WIDTH`` — we store only @@ -129,31 +125,18 @@ def _setup_recording( rank), so heavy experience bytes never touch SQL or Ray serialization. Args: - args: Parsed vLLM CLI args. - engine_client: AsyncLLM instance (we own it pre-init_app_state). + engine_client: AsyncLLM instance with ``trinity_recorder`` already set + by ``vLLMRolloutModel``. app: FastAPI app from ``build_app`` (we own it pre-serve_http). logger: Logger. """ - store: RecordStore = MemoryStore() - - # Rank is constant per process; capture once (RequestOutput does not expose - # parallel_config, so we read it from engine_client here, mirroring - # api_patch_v17.py:148). - try: - rank = int(engine_client.vllm_config.parallel_config._api_process_rank) - except Exception: - rank = 0 - - recorder = Recorder( - store=store, - enabled=True, - rank=rank, - engine_client=engine_client, - ) - - # (1) engine-level wrap — before init_app_state so serving objects inherit - # the wrapped reference. Idempotent via the __patched_*__ guard. - patch_engine_for_recording(engine_client, recorder, logger) + recorder = getattr(engine_client, TRINITY_RECORDER_ATTR, None) + if recorder is None: + raise RuntimeError( + "Generation recording API server requires vLLMRolloutModel to install " + "engine_client.trinity_recorder before server startup." + ) + store: RecordStore = recorder.store # (2) in-process middleware: API key -> contextvar. Zero network hop. app.add_middleware(RecordingIdentityMiddleware) @@ -167,7 +150,7 @@ def _setup_recording( logger.info( "Generation recording enabled: store=%s rank=%d", type(store).__name__, - rank, + recorder.rank, ) return recorder @@ -182,8 +165,8 @@ async def run_server_worker_in_ray( """Modified from vllm.entrypoints.openai.api_server.run_server_worker. Differs from api_patch_v17.py only in the recording wiring inserted between - ``build_app`` and ``init_app_state``, plus starting/stopping the recorder - flusher around ``serve_http``. + ``build_app`` and ``init_app_state``. The recorder lifecycle is owned by + ``vLLMRolloutModel``. """ if args.tool_parser_plugin and len(args.tool_parser_plugin) > 3: ToolParserManager.import_tool_parser(args.tool_parser_plugin) @@ -194,7 +177,7 @@ async def run_server_worker_in_ray( app = build_app(args) # --- recording wiring: engine wrap must precede init_app_state ----------- - recorder = _setup_recording(args, engine_client, app, logger) + _setup_recording(engine_client, app, logger) # ------------------------------------------------------------------------ await init_app_state(engine_client, app.state, args) @@ -208,8 +191,6 @@ async def run_server_worker_in_ray( listen_address, ) - recorder.start() - shutdown_task = await serve_http( app, sock=sock, @@ -233,7 +214,6 @@ async def run_server_worker_in_ray( try: await shutdown_task finally: - await recorder.stop() sock.close() diff --git a/trinity/common/models/vllm_patch/recording/store.py b/trinity/common/models/vllm_patch/recording/store.py deleted file mode 100644 index e08a4d449be..00000000000 --- a/trinity/common/models/vllm_patch/recording/store.py +++ /dev/null @@ -1,130 +0,0 @@ -"""Pluggable storage backends for recorded experiences. - -A ``RecordStore`` persists Trinity ``Experience`` objects in the vLLM API -server process. The only backend is ``MemoryStore`` — in-process, keyed by the -recording identity (``info["record_key"]`` = the API key / Ray-injected record -key), falling back to ``eid.suffix`` (the vLLM ``request_id``) when no identity -was supplied. - -The consume side is ``update_reward_by_record_key``: it sets ``reward``/``run``/ -``task`` on every experience in a record-key group, pops the group, and returns -it. This is the in-memory replacement for the old SQL ``HistoryRecorder``- -mediated join — the coordinator calls it (via ``/records/consume_task``) at -finalize time, so heavy experience bytes cross the network exactly once (store -→ coordinator pipeline) and never through Ray. - -Keying: experiences are identified by ``eid.suffix`` (the vLLM ``request_id``) -for traceability, but **grouped** by ``info["record_key"]`` so a whole reward -unit's worth of turns/samples can be reward-updated and consumed together. - -Concurrency: ``append_turn`` is called from a single background flusher task; -``update_reward_by_record_key`` is called from the ``/records/consume_task`` -HTTP handler. Both run in the same asyncio loop, so the dict is single-writer- -safe across these two without a lock. -""" -import abc -from collections import defaultdict -from typing import Optional - -from trinity.common.experience import Experience - -#: Attribute carrying the vLLM request id on each experience's ``info`` dict. -_REQUEST_ID_INFO_KEY = "request_id" -#: Attribute carrying the recording identity (the MemoryStore group key) on -#: each experience's ``info``. -_RECORD_KEY_INFO_KEY = "record_key" - - -class RecordStore(abc.ABC): - """Abstract persistence interface for recorded experiences.""" - - @abc.abstractmethod - async def append_turn(self, exp: Experience) -> None: - """Persist one completed experience.""" - - @abc.abstractmethod - async def update_reward_by_record_key( - self, record_key: str, reward: float, run: int, task: str - ) -> list[Experience]: - """Set reward/run/task on every experience in the group, pop and return it. - - Args: - record_key: The recording identity (group key). When the recorded - experience had no identity, this is its ``eid.suffix``. - reward: Reward to stamp on every experience in the group. - run: Run id to stamp on ``eid.run``. - task: Task id to stamp on ``eid.task``. - - Returns: - The (now reward-stamped) experiences of the group, in insertion - order. Empty list if the group was absent. - """ - - @abc.abstractmethod - async def get_task(self, record_key: str) -> list[Experience]: - """Return all experiences for a record key, in insertion order.""" - - @abc.abstractmethod - async def get_turn(self, record_key: str, request_id: str) -> Optional[Experience]: - """Return a single experience, or None if not found.""" - - @abc.abstractmethod - async def list_tasks(self) -> list[str]: - """Return all known record keys.""" - - @abc.abstractmethod - async def delete_task(self, record_key: str) -> None: - """Drop all experiences for a record key.""" - - -class MemoryStore(RecordStore): - """In-process store. - - Groups experiences by recording identity (``info["record_key"]``) when an - API key / record key was supplied, otherwise each turn is keyed by its own - ``eid.suffix`` (request_id) — so a missing identity never collapses - distinct turns. ``get_turn`` resolves an individual turn by - ``info["request_id"]``. - - Note: per-process under data-parallel serving — each API-server rank holds - only the experiences it served. The coordinator fans out - ``/records/consume_task`` to every rank and merges, so cross-rank - aggregation happens at consume time, not in storage. - """ - - def __init__(self) -> None: - # group key -> [experience, ...] in insertion order. - self._records: dict[str, list[Experience]] = defaultdict(list) - - @staticmethod - def _group_key(exp: Experience) -> str: - record_key = exp.info.get(_RECORD_KEY_INFO_KEY) - return record_key if record_key else exp.eid.suffix - - async def append_turn(self, exp: Experience) -> None: - self._records[self._group_key(exp)].append(exp) - - async def update_reward_by_record_key( - self, record_key: str, reward: float, run: int, task: str - ) -> list[Experience]: - exps = self._records.pop(record_key, []) - for exp in exps: - exp.reward = reward - exp.eid.run = run - exp.eid.task = task - return exps - - async def get_task(self, record_key: str) -> list[Experience]: - return list(self._records.get(record_key, [])) - - async def get_turn(self, record_key: str, request_id: str) -> Optional[Experience]: - for exp in self._records.get(record_key, []): - if exp.info.get(_REQUEST_ID_INFO_KEY) == request_id: - return exp - return None - - async def list_tasks(self) -> list[str]: - return list(self._records.keys()) - - async def delete_task(self, record_key: str) -> None: - self._records.pop(record_key, None) From 8ce1f4bfb29b71aa1906ab9417405bd6b3c83021 Mon Sep 17 00:00:00 2001 From: pxc Date: Thu, 25 Jun 2026 17:25:15 +0800 Subject: [PATCH 09/84] add tests --- tests/common/vllm_test.py | 288 ++++++++++++++++++++++++++++++++++++++ 1 file changed, 288 insertions(+) diff --git a/tests/common/vllm_test.py b/tests/common/vllm_test.py index 596fd823810..242f44bcf2a 100644 --- a/tests/common/vllm_test.py +++ b/tests/common/vllm_test.py @@ -6,9 +6,12 @@ from copy import deepcopy from typing import cast +import httpx +import openai import ray import torch from openai import BadRequestError +from packaging.version import parse as parse_version from parameterized import parameterized_class from transformers import AutoConfig, AutoTokenizer @@ -23,8 +26,10 @@ ) from trinity.common.config import Config from trinity.common.constants import ROLLOUT_WEIGHT_SYNC_GROUP_NAME, SyncMethod +from trinity.common.experience import Experience from trinity.common.models.allocator import Allocator from trinity.common.models.model import ModelWrapper +from trinity.common.models.vllm_patch import get_vllm_version from trinity.manager.synchronizer import Synchronizer DEBUG = False @@ -1806,6 +1811,289 @@ async def test_api_tool_calls(self): ) +class TestRecording(VLLMTestBase): + """Correctness of the in-vLLM generation recording flow (``enable_recording``). + + Verifies that every call path lands its finished turn in the in-process + ``MemoryStore`` under the right ``record_key``, and that + ``POST /records/consume_task`` flushes the recorder, reward-stamps the + whole record-key group, pops it, and returns it as serialized experiences. + + Paths covered (all async): + * Ray-direct ``generate`` / ``chat`` — record_key propagated via + ``record_key_ctx`` (set inside the actor by ``VLLMModel``). + * OpenAI HTTP regular / streaming / tool-call — record_key propagated + via the ``Authorization: Bearer `` header, captured by + ``RecordingIdentityMiddleware``. + + ``enable_recording`` forces ``enable_return_routed_experts`` in the + Allocator, and vLLM's routed-experts capturer raises on a non-MoE model, + so this test requires a MoE checkpoint (``TRINITY_MOE_MODEL_PATH``). + """ + + async def asyncSetUp(self): + if get_vllm_version() < parse_version("0.23.0"): + self.skipTest("generation recording requires vLLM >= 0.23.0") + self.config = get_template_config() + self.config.mode = "explore" + # enable_recording forces enable_return_routed_experts -> needs a MoE + # model (vLLM raises on dense models). Use a Qwen3-MoE checkpoint. + self.config.model.model_path = get_moe_model_path() + self.text_config = _get_text_config(self.config.model.model_path) + self.expected_routed_experts_layers = _count_moe_layers(self.text_config) + self.expected_routed_experts_topk = int(self.text_config.num_experts_per_tok) + self.config.model.custom_chat_template = CHAT_TEMPLATE + self.config.explorer.rollout_model.engine_type = "vllm" + self.config.explorer.rollout_model.engine_num = 1 + self.config.explorer.rollout_model.tensor_parallel_size = 2 + self.config.explorer.rollout_model.chat_template = CHAT_TEMPLATE + # enable_recording requires the OpenAI API server (the recording runner). + self.config.explorer.rollout_model.enable_openai_api = True + self.config.explorer.rollout_model.enable_recording = True + self.config.explorer.rollout_model.enable_expert_parallel = True + # Tool-call coverage; qwen3_coder matches the Qwen3.5 chat template. + self.config.explorer.rollout_model.enable_auto_tool_choice = True + self.config.explorer.rollout_model.tool_call_parser = "qwen3_coder" + self.config.explorer.rollout_model.enable_thinking = False + # History recording is client-side; the in-vLLM recorder is the subject. + self.config.explorer.rollout_model.enable_history = False + self.config.explorer.rollout_model.extra_engine_args = { + "max_num_seqs": 24, + "moe_backend": "triton", + "gdn_prefill_backend": "triton", + } + # check_and_update derives enable_return_routed_experts from this. + self.config.algorithm.enable_router_replay = True + self.config.check_and_update() + + self.engines, self.auxiliary_engines = await create_test_models(self.config) + self.model_wrapper = self.engines[0] + self.api_address = self.model_wrapper.api_address + self.expected_model_version = await self.model_wrapper.model_version_async + self._http = httpx.AsyncClient(timeout=120.0) + self._model_id = None + + async def asyncTearDown(self): + await self._http.aclose() + await super().asyncTearDown() + + # -- /records store query/consume helpers --------------------------------- + + async def _flush(self): + """Drain the recorder without popping anything (an empty consume).""" + resp = await self._http.post( + f"{self.api_address}/records/consume_task", json={"updates": []} + ) + resp.raise_for_status() + + async def _list_record_keys(self): + resp = await self._http.get(f"{self.api_address}/records/tasks") + resp.raise_for_status() + return resp.json()["record_keys"] + + async def _get_task(self, record_key: str) -> dict: + resp = await self._http.get(f"{self.api_address}/records/tasks/{record_key}") + resp.raise_for_status() + return resp.json() + + async def _get_turn_blob(self, record_key: str, request_id: str) -> Experience: + resp = await self._http.get( + f"{self.api_address}/records/tasks/{record_key}/turns/{request_id}/blob" + ) + resp.raise_for_status() + return Experience.deserialize(resp.content) + + async def _consume( + self, record_key: str, reward: float, run: int, task: str + ) -> list[Experience]: + resp = await self._http.post( + f"{self.api_address}/records/consume_task", + json={ + "updates": [{"record_key": record_key, "reward": reward, "run": run, "task": task}] + }, + ) + resp.raise_for_status() + return Experience.deserialize_many(resp.content) + + async def _openai_client(self, record_key: str) -> openai.AsyncOpenAI: + # record_key travels as the Bearer api_key -> RecordingIdentityMiddleware. + return openai.AsyncOpenAI(base_url=f"{self.api_address}/v1", api_key=record_key) + + async def _model_id(self, client: openai.AsyncOpenAI) -> str: + if self._model_id is None: + self._model_id = (await client.models.list()).data[0].id + return self._model_id + + # -- per-recorded-experience invariants ----------------------------------- + + def _assert_recorded_experience(self, exp: Experience, record_key: str): + self.assertEqual(exp.info.get("record_key"), record_key) + self.assertIsNotNone(exp.info.get("request_id")) + self.assertEqual(exp.info.get("rank"), 0) + self.assertEqual(exp.info.get("model_version"), self.expected_model_version) + self.assertGreater(len(exp.tokens), exp.prompt_length) + # The recorder forces top-1 logprobs even when the client omitted them. + self.assertGreater(len(exp.logprobs), 0) + self.assertEqual(len(exp.logprobs), len(exp.tokens) - exp.prompt_length) + self.assertGreater(len(exp.prompt_text), 0) + self.assertGreater(len(exp.response_text), 0) + + def _assert_recorded_routed_experts(self, exp: Experience): + # enable_return_routed_experts is forced on by enable_recording. + self.assertIsNotNone(exp.routed_experts) + re = exp.routed_experts + self.assertEqual(re.dtype, torch.uint8) + self.assertEqual(re.ndim, 3) + self.assertEqual(re.shape[1], self.expected_routed_experts_layers) + self.assertEqual(re.shape[2], self.expected_routed_experts_topk) + + async def test_record(self): # noqa: C901 + messages = [ + {"role": "system", "content": "You are a helpful assistant."}, + {"role": "user", "content": "Say hello in one short sentence."}, + ] + no_think = {"chat_template_kwargs": {"enable_thinking": False}} + + # ===== 1. Ray-direct generate (record_key via record_key_ctx) ===== + rk_gen = "trinity_record_generate" + await self.model_wrapper.generate_async( + ["Hello, world!"], n=1, temperature=1.0, max_tokens=16, record_key=rk_gen + ) + await self._flush() + self.assertIn(rk_gen, await self._list_record_keys()) + task = await self._get_task(rk_gen) + self.assertEqual(len(task["turns"]), 1) + # blob endpoint round-trips a full experience + request_id = task["turns"][0]["info"]["request_id"] + blob_exp = await self._get_turn_blob(rk_gen, request_id) + self._assert_recorded_experience(blob_exp, rk_gen) + self._assert_recorded_routed_experts(blob_exp) + consumed = await self._consume(rk_gen, reward=0.5, run=1, task="t_gen") + self.assertEqual(len(consumed), 1) + self.assertEqual(consumed[0].reward, 0.5) + self.assertEqual(consumed[0].eid.run, 1) + self.assertEqual(consumed[0].eid.task, "t_gen") + self._assert_recorded_experience(consumed[0], rk_gen) + self._assert_recorded_routed_experts(consumed[0]) + self.assertNotIn(rk_gen, await self._list_record_keys()) # popped + + # ===== 2. Ray-direct chat, n=2 (one record-key group, two samples) ===== + rk_chat = "trinity_record_chat" + chat_exps = await self.model_wrapper.chat_async( + messages, n=2, temperature=1.0, max_tokens=16, record_key=rk_chat + ) + self.assertEqual(len(chat_exps), 2) + await self._flush() + task = await self._get_task(rk_chat) + self.assertEqual(len(task["turns"]), 2) + # n=2 of one engine request -> two completions sharing one request_id, + # distinguished by sample_index. + self.assertEqual(sorted(t["info"]["sample_index"] for t in task["turns"]), [0, 1]) + self.assertEqual(len({t["info"]["request_id"] for t in task["turns"]}), 1) + consumed = await self._consume(rk_chat, reward=0.8, run=2, task="t_chat") + self.assertEqual(len(consumed), 2) + for exp in consumed: + self.assertEqual(exp.reward, 0.8) + self.assertEqual(exp.eid.run, 2) + self.assertEqual(exp.eid.task, "t_chat") + self._assert_recorded_experience(exp, rk_chat) + self._assert_recorded_routed_experts(exp) + self.assertNotIn(rk_chat, await self._list_record_keys()) + + # ===== 3. OpenAI regular (HTTP; record_key = Bearer api_key) ===== + rk_oai = "trinity_record_openai" + client = await self._openai_client(rk_oai) + model_id = await self._model_id(client) + resp = await client.chat.completions.create( + model=model_id, + messages=messages, + n=1, + temperature=0.7, + max_tokens=32, + extra_body=no_think, + ) + consumed = await self._consume(rk_oai, reward=0.3, run=3, task="t_oai") + self.assertEqual(len(consumed), 1) + self._assert_recorded_experience(consumed[0], rk_oai) + self._assert_recorded_routed_experts(consumed[0]) + self.assertEqual(consumed[0].response_text, resp.choices[0].message.content) + self.assertNotIn(rk_oai, await self._list_record_keys()) + + # ===== 4. OpenAI streaming (HTTP) ===== + rk_str = "trinity_record_stream" + sclient = await self._openai_client(rk_str) + stream = await sclient.chat.completions.create( + model=model_id, + messages=messages, + n=1, + stream=True, + temperature=0.7, + max_tokens=32, + extra_body=no_think, + ) + content = "" + async for chunk in stream: + delta = chunk.choices[0].delta.content + if delta: + content += delta + self.assertGreater(len(content), 0) + consumed = await self._consume(rk_str, reward=0.1, run=4, task="t_str") + self.assertEqual(len(consumed), 1) + self._assert_recorded_experience(consumed[0], rk_str) + self._assert_recorded_routed_experts(consumed[0]) + self.assertEqual(consumed[0].response_text, content) + self.assertNotIn(rk_str, await self._list_record_keys()) + + # ===== 5. OpenAI tool usage (HTTP) ===== + rk_tool = "trinity_record_tool" + tclient = await self._openai_client(rk_tool) + tools = [ + { + "type": "function", + "function": { + "name": "get_current_weather", + "description": "Get the current weather in a given location", + "parameters": { + "type": "object", + "properties": { + "location": { + "type": "string", + "description": "The city and state, e.g. San Francisco, CA", + } + }, + "required": ["location"], + }, + }, + } + ] + tool_messages = [{"role": "user", "content": "What's the weather like in Boston?"}] + tresp = await tclient.chat.completions.create( + model=model_id, + messages=tool_messages, + tools=tools, + tool_choice="auto", + max_tokens=64, + extra_body=no_think, + ) + consumed = await self._consume(rk_tool, reward=1.0, run=5, task="t_tool") + self.assertEqual(len(consumed), 1) + self._assert_recorded_experience(consumed[0], rk_tool) + self._assert_recorded_routed_experts(consumed[0]) + # The tool-augmented prompt (tool defs rendered by the chat template) + # must be part of the recorded experience. + self.assertIn("get_current_weather", consumed[0].prompt_text) + # If the model emitted a tool call, its function name is in the raw + # recorded response text. + choice = tresp.choices[0] + if choice.finish_reason == "tool_calls" and choice.message.tool_calls: + for tc in choice.message.tool_calls: + self.assertIn(tc.function.name, consumed[0].response_text) + self.assertNotIn(rk_tool, await self._list_record_keys()) + + # ===== global: every group consumed -> store is drained ===== + self.assertEqual(await self._list_record_keys(), []) + + class TestSuperLongGeneration(VLLMTestBase): async def asyncSetUp(self): self.config = get_template_config() From 02eb4e4553a6d515696e8c05226ea460d4c55faf Mon Sep 17 00:00:00 2001 From: "panxuchen.pxc" Date: Thu, 25 Jun 2026 17:35:34 +0800 Subject: [PATCH 10/84] add log --- tests/common/vllm_test.py | 12 ++++++------ trinity/common/models/model.py | 9 +++++++++ trinity/common/models/vllm_model.py | 8 ++++++++ trinity/common/models/vllm_patch/recording/server.py | 6 +++++- 4 files changed, 28 insertions(+), 7 deletions(-) diff --git a/tests/common/vllm_test.py b/tests/common/vllm_test.py index 242f44bcf2a..b537ed45cd2 100644 --- a/tests/common/vllm_test.py +++ b/tests/common/vllm_test.py @@ -1922,7 +1922,7 @@ async def _openai_client(self, record_key: str) -> openai.AsyncOpenAI: async def _model_id(self, client: openai.AsyncOpenAI) -> str: if self._model_id is None: self._model_id = (await client.models.list()).data[0].id - return self._model_id + return self._model_id # type: ignore [return-value] # -- per-recorded-experience invariants ----------------------------------- @@ -1931,12 +1931,12 @@ def _assert_recorded_experience(self, exp: Experience, record_key: str): self.assertIsNotNone(exp.info.get("request_id")) self.assertEqual(exp.info.get("rank"), 0) self.assertEqual(exp.info.get("model_version"), self.expected_model_version) - self.assertGreater(len(exp.tokens), exp.prompt_length) + self.assertGreater(len(exp.tokens), exp.prompt_length) # type: ignore [arg-type] # The recorder forces top-1 logprobs even when the client omitted them. - self.assertGreater(len(exp.logprobs), 0) - self.assertEqual(len(exp.logprobs), len(exp.tokens) - exp.prompt_length) - self.assertGreater(len(exp.prompt_text), 0) - self.assertGreater(len(exp.response_text), 0) + self.assertGreater(len(exp.logprobs), 0) # type: ignore [arg-type] + self.assertEqual(len(exp.logprobs), len(exp.tokens) - exp.prompt_length) # type: ignore [arg-type] + self.assertGreater(len(exp.prompt_text), 0) # type: ignore [arg-type] + self.assertGreater(len(exp.response_text), 0) # type: ignore [arg-type] def _assert_recorded_routed_experts(self, exp: Experience): # enable_return_routed_experts is forced on by enable_recording. diff --git a/trinity/common/models/model.py b/trinity/common/models/model.py index 70e56a367d0..541a5901296 100644 --- a/trinity/common/models/model.py +++ b/trinity/common/models/model.py @@ -138,6 +138,10 @@ def get_api_server_url(self) -> Optional[str]: """Get the API server URL if available.""" return None + def get_api_server_exit_reason(self) -> Optional[str]: + """Return API server exit reason if the background server task has exited.""" + return None + def get_api_key(self) -> str: """Get the API key.""" return "EMPTY" @@ -467,6 +471,11 @@ async def prepare(self) -> None: max_retries = 30 interval = 2 # seconds for i in range(max_retries): + reason = await self.model.get_api_server_exit_reason.remote() + if reason is not None: + raise RuntimeError( + f"API server at {self.api_address} exited before becoming ready: {reason}." + ) try: async with httpx.AsyncClient() as client: response = await client.get(self.api_address + "/health", timeout=5) diff --git a/trinity/common/models/vllm_model.py b/trinity/common/models/vllm_model.py index 116423dca56..eae146a5792 100644 --- a/trinity/common/models/vllm_model.py +++ b/trinity/common/models/vllm_model.py @@ -714,6 +714,14 @@ def get_api_server_url(self) -> Optional[str]: return None return f"http://{self.api_server_host}:{self.api_server_port}" + def get_api_server_exit_reason(self) -> Optional[str]: + if self.api_server is None or not self.api_server.done(): + return None + if self.api_server.cancelled(): + return "cancelled" + exc = self.api_server.exception() + return "unknown error" if exc is None else repr(exc) + async def reset_prefix_cache(self) -> None: await self.async_llm.reset_prefix_cache(reset_running_requests=True) diff --git a/trinity/common/models/vllm_patch/recording/server.py b/trinity/common/models/vllm_patch/recording/server.py index a33a79a2afb..5e5570b3bf6 100644 --- a/trinity/common/models/vllm_patch/recording/server.py +++ b/trinity/common/models/vllm_patch/recording/server.py @@ -225,7 +225,11 @@ async def run_server_in_ray( # Modified from vllm.entrypoints.openai.api_server.run_server listen_address, sock = setup_server_in_ray(args, logger) logger.info("vLLM API server listening on %s", listen_address) - await run_server_worker_in_ray(listen_address, sock, args, engine_client, logger) + try: + await run_server_worker_in_ray(listen_address, sock, args, engine_client, logger) + except Exception: + logger.exception("vLLM recording API server exited before becoming ready") + raise async def run_api_server_with_recording( From db492626cf8f828f7d79664252b2312ca0b8c3b9 Mon Sep 17 00:00:00 2001 From: "panxuchen.pxc" Date: Thu, 25 Jun 2026 17:39:35 +0800 Subject: [PATCH 11/84] fix middleware --- .../common/models/vllm_patch/recording/server.py | 16 +++++++++++++++- 1 file changed, 15 insertions(+), 1 deletion(-) diff --git a/trinity/common/models/vllm_patch/recording/server.py b/trinity/common/models/vllm_patch/recording/server.py index 5e5570b3bf6..9d07c8582d9 100644 --- a/trinity/common/models/vllm_patch/recording/server.py +++ b/trinity/common/models/vllm_patch/recording/server.py @@ -104,6 +104,20 @@ def dummy_add_signal_handler(self, *args, **kwargs): pass +def _add_recording_middleware(app) -> None: + """Install recording middleware before serving, even if vLLM built the stack. + + Some vLLM/FastAPI versions touch ``app.middleware_stack`` inside + ``build_app``. Starlette then rejects ``add_middleware`` with "Cannot add + middleware after an application has started", even though uvicorn has not + started serving yet. Clearing the cached stack lets Starlette rebuild it + with our middleware on first request. + """ + if getattr(app, "middleware_stack", None) is not None: + app.middleware_stack = None + app.add_middleware(RecordingIdentityMiddleware) + + def _setup_recording( engine_client, app, @@ -139,7 +153,7 @@ def _setup_recording( store: RecordStore = recorder.store # (2) in-process middleware: API key -> contextvar. Zero network hop. - app.add_middleware(RecordingIdentityMiddleware) + _add_recording_middleware(app) # (3) query routes mounted on the main app; OpenAI /v1/* surface untouched. app.include_router(query_router) From e6ae7dcd315a3b09748db482bc8da9d42de7ff43 Mon Sep 17 00:00:00 2001 From: "panxuchen.pxc" Date: Thu, 25 Jun 2026 18:03:27 +0800 Subject: [PATCH 12/84] update interface --- tests/common/vllm_test.py | 34 ++++++------ tests/explorer/explorer_test.py | 2 +- tests/trainer/trainer_test.py | 2 +- trinity/common/config.py | 2 +- trinity/common/models/recording/query.py | 54 ++++++++++--------- trinity/common/models/recording/store.py | 36 ++++++++++--- .../models/vllm_patch/recording/server.py | 2 +- trinity/common/workflows/workflow.py | 2 +- trinity/explorer/proxy/service.py | 2 +- trinity/explorer/rollout_coordinator.py | 14 ++--- 10 files changed, 85 insertions(+), 65 deletions(-) diff --git a/tests/common/vllm_test.py b/tests/common/vllm_test.py index b537ed45cd2..499a7ec60eb 100644 --- a/tests/common/vllm_test.py +++ b/tests/common/vllm_test.py @@ -1816,7 +1816,7 @@ class TestRecording(VLLMTestBase): Verifies that every call path lands its finished turn in the in-process ``MemoryStore`` under the right ``record_key``, and that - ``POST /records/consume_task`` flushes the recorder, reward-stamps the + ``POST /records/update_record`` flushes the recorder, reward-stamps the whole record-key group, pops it, and returns it as serialized experiences. Paths covered (all async): @@ -1882,24 +1882,22 @@ async def asyncTearDown(self): async def _flush(self): """Drain the recorder without popping anything (an empty consume).""" resp = await self._http.post( - f"{self.api_address}/records/consume_task", json={"updates": []} + f"{self.api_address}/records/update_record", json={"updates": []} ) resp.raise_for_status() async def _list_record_keys(self): - resp = await self._http.get(f"{self.api_address}/records/tasks") + resp = await self._http.get(f"{self.api_address}/records") resp.raise_for_status() return resp.json()["record_keys"] - async def _get_task(self, record_key: str) -> dict: - resp = await self._http.get(f"{self.api_address}/records/tasks/{record_key}") + async def _get_record_experiences(self, record_key: str) -> dict: + resp = await self._http.get(f"{self.api_address}/records/{record_key}") resp.raise_for_status() return resp.json() - async def _get_turn_blob(self, record_key: str, request_id: str) -> Experience: - resp = await self._http.get( - f"{self.api_address}/records/tasks/{record_key}/turns/{request_id}/blob" - ) + async def _get_request_experience(self, record_key: str, request_id: str) -> Experience: + resp = await self._http.get(f"{self.api_address}/records/{record_key}/request/{request_id}") resp.raise_for_status() return Experience.deserialize(resp.content) @@ -1907,7 +1905,7 @@ async def _consume( self, record_key: str, reward: float, run: int, task: str ) -> list[Experience]: resp = await self._http.post( - f"{self.api_address}/records/consume_task", + f"{self.api_address}/records/update_record", json={ "updates": [{"record_key": record_key, "reward": reward, "run": run, "task": task}] }, @@ -1961,11 +1959,11 @@ async def test_record(self): # noqa: C901 ) await self._flush() self.assertIn(rk_gen, await self._list_record_keys()) - task = await self._get_task(rk_gen) - self.assertEqual(len(task["turns"]), 1) + task = await self._get_record_experiences(rk_gen) + self.assertEqual(len(task["experiences"]), 1) # blob endpoint round-trips a full experience - request_id = task["turns"][0]["info"]["request_id"] - blob_exp = await self._get_turn_blob(rk_gen, request_id) + request_id = task["experiences"][0]["info"]["request_id"] + blob_exp = await self._get_request_experience(rk_gen, request_id) self._assert_recorded_experience(blob_exp, rk_gen) self._assert_recorded_routed_experts(blob_exp) consumed = await self._consume(rk_gen, reward=0.5, run=1, task="t_gen") @@ -1984,12 +1982,12 @@ async def test_record(self): # noqa: C901 ) self.assertEqual(len(chat_exps), 2) await self._flush() - task = await self._get_task(rk_chat) - self.assertEqual(len(task["turns"]), 2) + task = await self._get_record_experiences(rk_chat) + self.assertEqual(len(task["experiences"]), 2) # n=2 of one engine request -> two completions sharing one request_id, # distinguished by sample_index. - self.assertEqual(sorted(t["info"]["sample_index"] for t in task["turns"]), [0, 1]) - self.assertEqual(len({t["info"]["request_id"] for t in task["turns"]}), 1) + self.assertEqual(sorted(t["info"]["sample_index"] for t in task["experiences"]), [0, 1]) + self.assertEqual(len({t["info"]["request_id"] for t in task["experiences"]}), 1) consumed = await self._consume(rk_chat, reward=0.8, run=2, task="t_chat") self.assertEqual(len(consumed), 2) for exp in consumed: diff --git a/tests/explorer/explorer_test.py b/tests/explorer/explorer_test.py index 27203d29268..93eaec0292b 100644 --- a/tests/explorer/explorer_test.py +++ b/tests/explorer/explorer_test.py @@ -440,7 +440,7 @@ def run_agent(proxy_url, model_path: str, stream: bool): @unittest.skip( "serve-mode experience collection moved to the in-vLLM MemoryStore + " - "/records/consume_task flow; the proxy /feedback//commit path and external " + "/records/update_record flow; the proxy /feedback//commit path and external " "reward reporting are being redesigned (see recording refactor plan)." ) class ServeTest(RayUnittestBaseAsync): diff --git a/tests/trainer/trainer_test.py b/tests/trainer/trainer_test.py index ccc16b0d92d..6865aa275ec 100644 --- a/tests/trainer/trainer_test.py +++ b/tests/trainer/trainer_test.py @@ -1100,7 +1100,7 @@ async def run_math_workflow(serve_url: str, task: dict): @unittest.skip( "serve-mode experience collection moved to the in-vLLM MemoryStore + " - "/records/consume_task flow; the proxy /feedback//commit path and external " + "/records/update_record flow; the proxy /feedback//commit path and external " "reward reporting are being redesigned (see recording refactor plan)." ) class TestServeWithTrainer(RayUnittestBaseAsync): diff --git a/trinity/common/config.py b/trinity/common/config.py index 18fb999f20c..e9a6a6d6f73 100644 --- a/trinity/common/config.py +++ b/trinity/common/config.py @@ -597,7 +597,7 @@ class InferenceModelConfig: # recording identity (``record_key``). This is the single switch for the # recording flow — when on, the explorer also consumes from the store: the # WorkflowRunner ships only a small reward map keyed by ``record_key`` and - # the RolloutCoordinator pulls heavy experiences via ``/records/consume_task`` + # the RolloutCoordinator pulls heavy experiences via ``/records/update_record`` # at finalize time. When off (default), runners ship serialized experiences # through the scheduler as before (legacy path). When True, the Allocator # forces ``enable_return_routed_experts``. VLLMModel mirrors the recording diff --git a/trinity/common/models/recording/query.py b/trinity/common/models/recording/query.py index 7d6e2de0cf5..1acbd58eea7 100644 --- a/trinity/common/models/recording/query.py +++ b/trinity/common/models/recording/query.py @@ -1,4 +1,4 @@ -"""HTTP query/consume endpoints over recorded experiences.""" +"""HTTP query/update endpoints over recorded experiences.""" from typing import List @@ -15,15 +15,15 @@ query_router = APIRouter(prefix="/records", tags=["trinity-recording"]) -class _RecordKeyUpdate(BaseModel): +class _RecordUpdate(BaseModel): record_key: str reward: float run: int = 0 task: str = "" -class _ConsumeTaskRequest(BaseModel): - updates: List[_RecordKeyUpdate] +class _UpdateRecordRequest(BaseModel): + updates: List[_RecordUpdate] def _store(request: Request) -> RecordStore: @@ -41,34 +41,27 @@ def _recorder(request: Request) -> Recorder: async def _get_exp(store: RecordStore, record_key: str, request_id: str) -> Experience: - exp = await store.get_turn(record_key, request_id) + exp = await store.get_request_experience(record_key, request_id) if exp is None: raise HTTPException(status_code=404, detail="experience not found") return exp -@query_router.get("/tasks") -async def list_tasks(request: Request) -> dict: +@query_router.get("") +async def list_records(request: Request) -> dict: store = _store(request) - return {"record_keys": await store.list_tasks()} + return {"record_keys": await store.list_records()} -@query_router.get("/tasks/{record_key}") -async def get_task(record_key: str, request: Request) -> dict: +@query_router.get("/{record_key}") +async def get_record_experiences(record_key: str, request: Request) -> dict: store = _store(request) - experiences = await store.get_task(record_key) - return {"record_key": record_key, "turns": [e.to_dict() for e in experiences]} + experiences = await store.get_record_experiences(record_key) + return {"record_key": record_key, "experiences": [e.to_dict() for e in experiences]} -@query_router.get("/tasks/{record_key}/turns/{request_id}") -async def get_turn(record_key: str, request_id: str, request: Request) -> dict: - store = _store(request) - exp = await _get_exp(store, record_key, request_id) - return exp.to_dict() - - -@query_router.get("/tasks/{record_key}/turns/{request_id}/blob") -async def get_turn_blob(record_key: str, request_id: str, request: Request) -> Response: +@query_router.get("/{record_key}/request/{request_id}") +async def get_request_experience(record_key: str, request_id: str, request: Request) -> Response: store = _store(request) exp = await _get_exp(store, record_key, request_id) return Response( @@ -78,15 +71,24 @@ async def get_turn_blob(record_key: str, request_id: str, request: Request) -> R ) -@query_router.delete("/tasks/{record_key}") -async def delete_task(record_key: str, request: Request) -> dict: +@query_router.delete("/{record_key}") +async def delete_record_experiences(record_key: str, request: Request) -> dict: store = _store(request) - await store.delete_task(record_key) + await store.delete_record_experiences(record_key) return {"record_key": record_key, "deleted": True} -@query_router.post("/consume_task") -async def consume_task(req: _ConsumeTaskRequest, request: Request) -> Response: +@query_router.delete("/{record_key}/request/{request_id}") +async def delete_request_experience(record_key: str, request_id: str, request: Request) -> dict: + store = _store(request) + deleted = await store.delete_request_experience(record_key, request_id) + if not deleted: + raise HTTPException(status_code=404, detail="experience not found") + return {"record_key": record_key, "request_id": request_id, "deleted": True} + + +@query_router.post("/update_record") +async def update_record(req: _UpdateRecordRequest, request: Request) -> Response: store = _store(request) recorder = _recorder(request) await recorder.flush() diff --git a/trinity/common/models/recording/store.py b/trinity/common/models/recording/store.py index ee7a764e367..3851e3d0959 100644 --- a/trinity/common/models/recording/store.py +++ b/trinity/common/models/recording/store.py @@ -24,21 +24,27 @@ async def update_reward_by_record_key( """Set reward/run/task on every experience in the group, pop and return it.""" @abc.abstractmethod - async def get_task(self, record_key: str) -> list[Experience]: + async def get_record_experiences(self, record_key: str) -> list[Experience]: """Return all experiences for a record key, in insertion order.""" @abc.abstractmethod - async def get_turn(self, record_key: str, request_id: str) -> Optional[Experience]: + async def get_request_experience( + self, record_key: str, request_id: str + ) -> Optional[Experience]: """Return a single experience, or None if not found.""" @abc.abstractmethod - async def list_tasks(self) -> list[str]: + async def list_records(self) -> list[str]: """Return all known record keys.""" @abc.abstractmethod - async def delete_task(self, record_key: str) -> None: + async def delete_record_experiences(self, record_key: str) -> None: """Drop all experiences for a record key.""" + @abc.abstractmethod + async def delete_request_experience(self, record_key: str, request_id: str) -> bool: + """Drop one experience by request id. Return True if one was deleted.""" + class MemoryStore(RecordStore): """In-process store grouped by recording identity.""" @@ -64,17 +70,31 @@ async def update_reward_by_record_key( exp.eid.task = task return exps - async def get_task(self, record_key: str) -> list[Experience]: + async def get_record_experiences(self, record_key: str) -> list[Experience]: return list(self._records.get(record_key, [])) - async def get_turn(self, record_key: str, request_id: str) -> Optional[Experience]: + async def get_request_experience( + self, record_key: str, request_id: str + ) -> Optional[Experience]: for exp in self._records.get(record_key, []): if exp.info.get(REQUEST_ID_INFO_KEY) == request_id: return exp return None - async def list_tasks(self) -> list[str]: + async def list_records(self) -> list[str]: return list(self._records.keys()) - async def delete_task(self, record_key: str) -> None: + async def delete_record_experiences(self, record_key: str) -> None: self._records.pop(record_key, None) + + async def delete_request_experience(self, record_key: str, request_id: str) -> bool: + exps = self._records.get(record_key) + if not exps: + return False + for index, exp in enumerate(exps): + if exp.info.get(REQUEST_ID_INFO_KEY) == request_id: + del exps[index] + if not exps: + self._records.pop(record_key, None) + return True + return False diff --git a/trinity/common/models/vllm_patch/recording/server.py b/trinity/common/models/vllm_patch/recording/server.py index 9d07c8582d9..27084111098 100644 --- a/trinity/common/models/vllm_patch/recording/server.py +++ b/trinity/common/models/vllm_patch/recording/server.py @@ -135,7 +135,7 @@ def _setup_recording( by VLLMModel at engine creation and in ``sync_model_weights``). The store backend is always the in-process ``MemoryStore``; the coordinator - drains it at finalize time via ``/records/consume_task`` (fanned out per + drains it at finalize time via ``/records/update_record`` (fanned out per rank), so heavy experience bytes never touch SQL or Ray serialization. Args: diff --git a/trinity/common/workflows/workflow.py b/trinity/common/workflows/workflow.py index e7670c79011..328a1d237a8 100644 --- a/trinity/common/workflows/workflow.py +++ b/trinity/common/workflows/workflow.py @@ -232,7 +232,7 @@ def process_messages_to_experience( # captured the authoritative heavy data (real logprobs without an extra # forward, real routed_experts) into the MemoryStore, keyed by the # record_key the chat call carried. Replace this with an in-process - # lookup by record_key (store.get_task / consume), then concatenate the + # lookup by record_key (store.get_record_experiences / update), then concatenate the # session's turns (info["sample_index"] orders them) into one # experience here. Requires threading the per-call record_key down to # this call site. diff --git a/trinity/explorer/proxy/service.py b/trinity/explorer/proxy/service.py index 4a247442410..89eef711ed7 100644 --- a/trinity/explorer/proxy/service.py +++ b/trinity/explorer/proxy/service.py @@ -15,7 +15,7 @@ class ExplorerService: The proxy is a request router + model-weight sync coordinator for serve mode. Experience collection used to live here (SQL-mediated ``/feedback``/``/commit``); it has been removed in favor of the in-vLLM - ``MemoryStore`` + ``/records/consume_task`` flow driven by the rollout + ``MemoryStore`` + ``/records/update_record`` flow driven by the rollout coordinator. Serve-mode external reward reporting is therefore pending (see the recording refactor plan). """ diff --git a/trinity/explorer/rollout_coordinator.py b/trinity/explorer/rollout_coordinator.py index feed993a337..2591cda6192 100644 --- a/trinity/explorer/rollout_coordinator.py +++ b/trinity/explorer/rollout_coordinator.py @@ -73,7 +73,7 @@ def __init__( self.running = False self.detailed_stats = getattr(getattr(config, "monitor", None), "detailed_stats", False) # Lazily-resolved map of rollout engine_id -> API server URL, for the - # recording path's per-rank /records/consume_task fan-out. + # recording path's per-rank /records/update_record fan-out. self._rank_urls: Optional[Dict[int, str]] = None def _enable_recording(self) -> bool: @@ -336,7 +336,7 @@ async def _consume_recorded_experiences(self, payload_chunks: List[bytes]) -> di ``payload_chunks`` are small pickle reward maps produced by the runners (``{"engine_id": int, "updates": [{"record_key", "reward", "run", "task"}]}``). - Group updates by engine, fan out ``POST /records/consume_task`` to each + Group updates by engine, fan out ``POST /records/update_record`` to each rank (which drains its recorder, reward-stamps the matching record-key groups, pops them, and returns ``serialize_many`` bytes), deserialize, and feed the assembled experiences straight into the pipeline — no Ray @@ -357,7 +357,7 @@ async def _consume_recorded_experiences(self, payload_chunks: List[bytes]) -> di rank_urls = self._resolve_rank_urls() async with httpx.AsyncClient(timeout=_CONSUME_TIMEOUT) as client: requests = [ - self._post_consume_task(client, rank_urls[engine_id], updates) + self._post_update_record(client, rank_urls[engine_id], updates) for engine_id, updates in per_engine.items() if engine_id in rank_urls ] @@ -369,21 +369,21 @@ async def _consume_recorded_experiences(self, payload_chunks: List[bytes]) -> di exps.extend(Experience.deserialize_many(resp_bytes)) return await self.experience_pipeline.process_experiences(exps) - async def _post_consume_task( + async def _post_update_record( self, client: httpx.AsyncClient, rank_url: str, updates: List[dict] ) -> bytes: """POST a batch of record-key reward updates to one rank; return heavy bytes.""" try: resp = await client.post( - f"{rank_url}/records/consume_task", + f"{rank_url}/records/update_record", json={"updates": updates}, ) except (httpx.TimeoutException, httpx.RequestError) as exc: - self.logger.error("consume_task to %s failed: %s", rank_url, exc) + self.logger.error("update_record to %s failed: %s", rank_url, exc) return b"" if resp.status_code != 200: self.logger.error( - "consume_task to %s returned %d: %s", + "update_record to %s returned %d: %s", rank_url, resp.status_code, resp.text[:200], From 63a5404b83786781d8d18943b725f622799a9904 Mon Sep 17 00:00:00 2001 From: "panxuchen.pxc" Date: Thu, 25 Jun 2026 18:14:29 +0800 Subject: [PATCH 13/84] fix prompt text --- tests/common/vllm_test.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/tests/common/vllm_test.py b/tests/common/vllm_test.py index 499a7ec60eb..b843c22b937 100644 --- a/tests/common/vllm_test.py +++ b/tests/common/vllm_test.py @@ -1933,7 +1933,11 @@ def _assert_recorded_experience(self, exp: Experience, record_key: str): # The recorder forces top-1 logprobs even when the client omitted them. self.assertGreater(len(exp.logprobs), 0) # type: ignore [arg-type] self.assertEqual(len(exp.logprobs), len(exp.tokens) - exp.prompt_length) # type: ignore [arg-type] - self.assertGreater(len(exp.prompt_text), 0) # type: ignore [arg-type] + # Ray-direct generate may pass token ids into vLLM; in that path + # RequestOutput.prompt is not populated. Recording intentionally does + # not decode token ids back to text on the hot path. + if exp.prompt_text is not None: + self.assertGreater(len(exp.prompt_text), 0) self.assertGreater(len(exp.response_text), 0) # type: ignore [arg-type] def _assert_recorded_routed_experts(self, exp: Experience): From 5ff21b239dd12dac41957aaa18fb5dff8b63ac98 Mon Sep 17 00:00:00 2001 From: "panxuchen.pxc" Date: Thu, 25 Jun 2026 18:36:46 +0800 Subject: [PATCH 14/84] fix streaming recorder --- tests/common/vllm_test.py | 18 ++- .../models/vllm_patch/recording/recorder.py | 141 +++++++++++++++++- 2 files changed, 153 insertions(+), 6 deletions(-) diff --git a/tests/common/vllm_test.py b/tests/common/vllm_test.py index b843c22b937..ac822e20d93 100644 --- a/tests/common/vllm_test.py +++ b/tests/common/vllm_test.py @@ -1839,6 +1839,10 @@ async def asyncSetUp(self): # enable_recording forces enable_return_routed_experts -> needs a MoE # model (vLLM raises on dense models). Use a Qwen3-MoE checkpoint. self.config.model.model_path = get_moe_model_path() + self.tokenizer = AutoTokenizer.from_pretrained( + self.config.model.model_path, + trust_remote_code=True, + ) self.text_config = _get_text_config(self.config.model.model_path) self.expected_routed_experts_layers = _count_moe_layers(self.text_config) self.expected_routed_experts_topk = int(self.text_config.num_experts_per_tok) @@ -1917,7 +1921,7 @@ async def _openai_client(self, record_key: str) -> openai.AsyncOpenAI: # record_key travels as the Bearer api_key -> RecordingIdentityMiddleware. return openai.AsyncOpenAI(base_url=f"{self.api_address}/v1", api_key=record_key) - async def _model_id(self, client: openai.AsyncOpenAI) -> str: + async def _get_model_id(self, client: openai.AsyncOpenAI) -> str: if self._model_id is None: self._model_id = (await client.models.list()).data[0].id return self._model_id # type: ignore [return-value] @@ -1938,7 +1942,11 @@ def _assert_recorded_experience(self, exp: Experience, record_key: str): # not decode token ids back to text on the hot path. if exp.prompt_text is not None: self.assertGreater(len(exp.prompt_text), 0) - self.assertGreater(len(exp.response_text), 0) # type: ignore [arg-type] + # OpenAI streaming clients receive text as delta chunks. The finished + # engine output recorded below may carry an empty native + # CompletionOutput.text even when response token ids are present; avoid + # decoding tokens on the recording hot path just to populate this field. + self.assertIsNotNone(exp.response_text) def _assert_recorded_routed_experts(self, exp: Experience): # enable_return_routed_experts is forced on by enable_recording. @@ -2005,7 +2013,7 @@ async def test_record(self): # noqa: C901 # ===== 3. OpenAI regular (HTTP; record_key = Bearer api_key) ===== rk_oai = "trinity_record_openai" client = await self._openai_client(rk_oai) - model_id = await self._model_id(client) + model_id = await self._get_model_id(client) resp = await client.chat.completions.create( model=model_id, messages=messages, @@ -2043,7 +2051,9 @@ async def test_record(self): # noqa: C901 self.assertEqual(len(consumed), 1) self._assert_recorded_experience(consumed[0], rk_str) self._assert_recorded_routed_experts(consumed[0]) - self.assertEqual(consumed[0].response_text, content) + response_token_ids = consumed[0].tokens[consumed[0].prompt_length :].tolist() + decoded_content = self.tokenizer.decode(response_token_ids, skip_special_tokens=True) + self.assertEqual(decoded_content, content) self.assertNotIn(rk_str, await self._list_record_keys()) # ===== 5. OpenAI tool usage (HTTP) ===== diff --git a/trinity/common/models/vllm_patch/recording/recorder.py b/trinity/common/models/vllm_patch/recording/recorder.py index 585feca4b02..5f12704bbb0 100644 --- a/trinity/common/models/vllm_patch/recording/recorder.py +++ b/trinity/common/models/vllm_patch/recording/recorder.py @@ -19,6 +19,7 @@ """ import functools import logging +from types import SimpleNamespace from typing import Optional from trinity.common.models.recording.context import record_key_ctx @@ -109,6 +110,138 @@ def _build_multi_modal_inputs(engine_client, prompt, output, logger: logging.Log return None +def _completion_index(completion, fallback: int) -> int: + return int(getattr(completion, "index", fallback)) + + +def _list_or_empty(value): + if value is None: + return [] + return list(value) + + +def _is_cumulative(prev_tokens: list[int], cur_tokens: list[int]) -> bool: + return ( + bool(prev_tokens) + and len(cur_tokens) >= len(prev_tokens) + and (cur_tokens[: len(prev_tokens)] == prev_tokens) + ) + + +def _concat_routed_experts(prev, cur): + if cur is None: + return prev + if prev is None: + return cur + try: + import numpy as np + + if isinstance(prev, np.ndarray) or isinstance(cur, np.ndarray): + return np.concatenate([prev, cur], axis=0) + except Exception: + pass + try: + import torch + + if isinstance(prev, torch.Tensor) or isinstance(cur, torch.Tensor): + return torch.cat([torch.as_tensor(prev), torch.as_tensor(cur)], dim=0) + except Exception: + pass + try: + return prev + cur + except Exception: + return cur + + +def _accumulate_request_output(state, output): + if state is None: + state = { + "request_id": output.request_id, + "prompt_token_ids": _list_or_empty(getattr(output, "prompt_token_ids", None)), + "prompt": getattr(output, "prompt", None), + "outputs": {}, + "order": [], + } + elif not state["prompt_token_ids"]: + state["prompt_token_ids"] = _list_or_empty(getattr(output, "prompt_token_ids", None)) + state["prompt"] = state["prompt"] or getattr(output, "prompt", None) + + for fallback_index, completion in enumerate(list(getattr(output, "outputs", None) or [])): + index = _completion_index(completion, fallback_index) + if index not in state["outputs"]: + state["outputs"][index] = { + "token_ids": [], + "logprobs": None, + "text": "", + "routed_experts": None, + } + state["order"].append(index) + + acc = state["outputs"][index] + cur_token_ids = _list_or_empty(getattr(completion, "token_ids", None)) + cumulative = _is_cumulative(acc["token_ids"], cur_token_ids) + if cumulative: + acc["token_ids"] = cur_token_ids + elif cur_token_ids: + acc["token_ids"].extend(cur_token_ids) + + cur_logprobs = getattr(completion, "logprobs", None) + if cur_logprobs is not None: + cur_logprobs = list(cur_logprobs) + if not cur_logprobs: + pass + elif cumulative: + acc["logprobs"] = cur_logprobs + elif acc["logprobs"] is None: + acc["logprobs"] = cur_logprobs + else: + acc["logprobs"].extend(cur_logprobs) + + cur_text = getattr(completion, "text", None) or "" + if cur_text: + if cumulative and cur_text.startswith(acc["text"]): + acc["text"] = cur_text + else: + acc["text"] += cur_text + + cur_routed_experts = getattr(completion, "routed_experts", None) + if cur_routed_experts is not None: + if cumulative: + acc["routed_experts"] = cur_routed_experts + else: + acc["routed_experts"] = _concat_routed_experts( + acc["routed_experts"], + cur_routed_experts, + ) + + return state + + +def _build_record_output(state, last): + if state is None: + return last + completions = [] + for index in state["order"]: + acc = state["outputs"][index] + completions.append( + SimpleNamespace( + index=index, + token_ids=acc["token_ids"], + logprobs=acc["logprobs"], + text=acc["text"], + routed_experts=acc["routed_experts"], + ) + ) + return SimpleNamespace( + request_id=state["request_id"], + prompt_token_ids=state["prompt_token_ids"], + prompt=state["prompt"], + outputs=completions, + finished=getattr(last, "finished", False), + prompt_routed_experts=getattr(last, "prompt_routed_experts", None), + ) + + def patch_engine_for_recording( engine_client, recorder: "Recorder", @@ -155,10 +288,13 @@ async def _patched_generate(*args, **kwargs): ) last = None + accumulated = None # ``current`` is the original *bound* method captured pre-wrap, so it # still resolves ``self`` correctly. Yields RequestOutput unchanged. async for out in current(*args, **kwargs): last = out + if recorder.enabled: + accumulated = _accumulate_request_output(accumulated, out) yield out if recorder.enabled and last is not None and getattr(last, "finished", False): @@ -168,14 +304,15 @@ async def _patched_generate(*args, **kwargs): # opt into grouping this turn, so skip recording entirely. record_key = record_key_ctx.get() if record_key is not None: + record_output = _build_record_output(accumulated, last) multi_modal_inputs = _build_multi_modal_inputs( engine_client, prompt, - last, + record_output, logger, ) recorder.schedule_record( - last, + record_output, record_key, multi_modal_inputs=multi_modal_inputs, ) From 6a0ad4445ee91ff537ca1a7c4990409797f153bf Mon Sep 17 00:00:00 2001 From: "panxuchen.pxc" Date: Thu, 25 Jun 2026 18:39:44 +0800 Subject: [PATCH 15/84] add delta stream --- .../models/vllm_patch/recording/recorder.py | 48 +++++++++---------- 1 file changed, 23 insertions(+), 25 deletions(-) diff --git a/trinity/common/models/vllm_patch/recording/recorder.py b/trinity/common/models/vllm_patch/recording/recorder.py index 5f12704bbb0..5f9351776f3 100644 --- a/trinity/common/models/vllm_patch/recording/recorder.py +++ b/trinity/common/models/vllm_patch/recording/recorder.py @@ -120,12 +120,9 @@ def _list_or_empty(value): return list(value) -def _is_cumulative(prev_tokens: list[int], cur_tokens: list[int]) -> bool: - return ( - bool(prev_tokens) - and len(cur_tokens) >= len(prev_tokens) - and (cur_tokens[: len(prev_tokens)] == prev_tokens) - ) +def _is_delta_output(sampling_params) -> bool: + output_kind = getattr(sampling_params, "output_kind", None) + return getattr(output_kind, "name", output_kind) == "DELTA" def _concat_routed_experts(prev, cur): @@ -153,7 +150,7 @@ def _concat_routed_experts(prev, cur): return cur -def _accumulate_request_output(state, output): +def _accumulate_request_output(state, output, *, is_delta_output: bool): if state is None: state = { "request_id": output.request_id, @@ -179,40 +176,36 @@ def _accumulate_request_output(state, output): acc = state["outputs"][index] cur_token_ids = _list_or_empty(getattr(completion, "token_ids", None)) - cumulative = _is_cumulative(acc["token_ids"], cur_token_ids) - if cumulative: - acc["token_ids"] = cur_token_ids - elif cur_token_ids: + if is_delta_output and cur_token_ids: acc["token_ids"].extend(cur_token_ids) + else: + acc["token_ids"] = cur_token_ids cur_logprobs = getattr(completion, "logprobs", None) if cur_logprobs is not None: cur_logprobs = list(cur_logprobs) if not cur_logprobs: pass - elif cumulative: - acc["logprobs"] = cur_logprobs - elif acc["logprobs"] is None: - acc["logprobs"] = cur_logprobs - else: + elif is_delta_output and acc["logprobs"] is not None: acc["logprobs"].extend(cur_logprobs) + else: + acc["logprobs"] = cur_logprobs cur_text = getattr(completion, "text", None) or "" - if cur_text: - if cumulative and cur_text.startswith(acc["text"]): - acc["text"] = cur_text - else: - acc["text"] += cur_text + if is_delta_output and cur_text: + acc["text"] += cur_text + else: + acc["text"] = cur_text cur_routed_experts = getattr(completion, "routed_experts", None) if cur_routed_experts is not None: - if cumulative: - acc["routed_experts"] = cur_routed_experts - else: + if is_delta_output: acc["routed_experts"] = _concat_routed_experts( acc["routed_experts"], cur_routed_experts, ) + else: + acc["routed_experts"] = cur_routed_experts return state @@ -286,6 +279,7 @@ async def _patched_generate(*args, **kwargs): sampling_params.logprobs = ( max(cur, _RECORDER_LOGPROB_WIDTH) if cur is not None else _RECORDER_LOGPROB_WIDTH ) + is_delta_output = _is_delta_output(sampling_params) last = None accumulated = None @@ -294,7 +288,11 @@ async def _patched_generate(*args, **kwargs): async for out in current(*args, **kwargs): last = out if recorder.enabled: - accumulated = _accumulate_request_output(accumulated, out) + accumulated = _accumulate_request_output( + accumulated, + out, + is_delta_output=is_delta_output, + ) yield out if recorder.enabled and last is not None and getattr(last, "finished", False): From bb2f22cd767feb191d053de1fb7f88de9bd0895c Mon Sep 17 00:00:00 2001 From: "panxuchen.pxc" Date: Thu, 25 Jun 2026 19:31:37 +0800 Subject: [PATCH 16/84] refactor history recording --- trinity/common/models/model.py | 170 +++++------------- trinity/common/models/vllm_model.py | 16 +- .../common/workflows/step_wise_workflow.py | 14 +- trinity/explorer/workflow_runner.py | 68 +++++-- 4 files changed, 121 insertions(+), 147 deletions(-) diff --git a/trinity/common/models/model.py b/trinity/common/models/model.py index 541a5901296..38d76349b48 100644 --- a/trinity/common/models/model.py +++ b/trinity/common/models/model.py @@ -17,12 +17,6 @@ from trinity.common.config import InferenceModelConfig from trinity.common.constants import RunningStatus, SyncMethod from trinity.common.experience import Experience -from trinity.common.models.experience_extraction import ( - HistoryRecordingStream, - convert_api_output_to_experience, - get_routed_experts_layout, -) -from trinity.common.models.mm_utils import should_use_processor, vLLMMultiModalRender from trinity.common.models.utils import get_action_mask_method from trinity.utils.log import get_logger @@ -146,6 +140,12 @@ def get_api_key(self) -> str: """Get the API key.""" return "EMPTY" + async def extract_experience_from_history( + self, record_key: str, clear_history: bool = True + ) -> List[Experience]: + """Extract recorded experiences by record key.""" + raise NotImplementedError("Recording extraction is not supported by this model.") + def get_model_config(self) -> InferenceModelConfig: """Get the model configuration.""" return self.config @@ -437,26 +437,16 @@ def __init__( self.enable_lora = config.enable_lora self.enable_history = config.enable_history self.history = [] + self.recording_history_offsets: Dict[str, int] = {} self.status = RunningStatus.RUNNING self.workflow_state: Dict = {} self.request_count = 0 self.state_lock = asyncio.Lock() - self._routed_experts_layout: Optional[Tuple[int, int]] = None - self._mm_render = None async def prepare(self) -> None: """Prepare some necessary information for the model before inference.""" if not self.config.enable_openai_api: return - if ( - self.config.enable_return_routed_experts - and self.config.engine_type == "sglang" - and self._routed_experts_layout is None - ): - self._routed_experts_layout = get_routed_experts_layout( - self.model_path, - trust_remote_code=self.config.trust_remote_code, - ) if self.api_address is None: if self.model is None: raise ValueError("Cannot get API address from the model.") @@ -497,28 +487,12 @@ def _record_history(self, exps: Union[Experience, List[Experience]]) -> None: else: raise TypeError("Expected Experience or List[Experience], got {}".format(type(exps))) - def _assert_openai_routed_experts_request_supported( - self, extra_body: Dict[str, Any], kwargs: Dict[str, Any] - ) -> None: - """Validate routed_experts constraints for OpenAI-compatible backends.""" - requested_routed_experts = self.config.enable_return_routed_experts or bool( - extra_body.get("return_routed_experts", False) - ) - if requested_routed_experts: - if self.config.engine_type not in {"sglang", "vllm"}: - raise ValueError("Routed experts can only be returned from SGLang or vLLM.") - if kwargs.get("stream", False): - raise ValueError("Routed experts cannot be returned for streaming requests.") - if self.config.engine_type == "sglang" and kwargs.get("n", 1) != 1: - raise ValueError( - "SGLang OpenAI API returns routed_experts at response level only; " - "set n=1 when requesting routed_experts." - ) - @_history_recorder def generate(self, prompts: List[str], **kwargs) -> List[Experience]: """Generate a list of experiences from a list of prompts.""" lora_request = self.get_lora_request() + if self.config.enable_recording and kwargs.get("record_key") is None: + kwargs["record_key"] = self._api_key results = ray.get( [self.model.generate.remote(prompt, lora_request, **kwargs) for prompt in prompts] ) @@ -528,6 +502,8 @@ def generate(self, prompts: List[str], **kwargs) -> List[Experience]: async def generate_async(self, prompts: List[str], **kwargs) -> List[Experience]: """Generate a list of experiences from a list of prompts in async.""" lora_request = await self.get_lora_request_async() + if self.config.enable_recording and kwargs.get("record_key") is None: + kwargs["record_key"] = self._api_key results = await asyncio.gather( *[self.model.generate.remote(prompt, lora_request, **kwargs) for prompt in prompts] ) @@ -539,6 +515,8 @@ def chat( ) -> List[Experience]: """Generate a list of experiences from a list of messages.""" lora_request = self.get_lora_request() + if self.config.enable_recording and record_key is None: + record_key = self._api_key return ray.get( self.model.chat.remote( messages, lora_request=lora_request, record_key=record_key, **kwargs @@ -551,6 +529,8 @@ async def chat_async( ) -> List[Experience]: """Generate a list of experiences from a list of messages in async.""" lora_request = await self.get_lora_request_async() + if self.config.enable_recording and record_key is None: + record_key = self._api_key return await self.model.chat.remote( messages, lora_request=lora_request, record_key=record_key, **kwargs ) @@ -594,6 +574,14 @@ def api_key(self) -> str: """Get the API key.""" return self._api_key + def set_api_key(self, api_key: str) -> None: + """Set the API key used by existing and future OpenAI clients.""" + self._api_key = api_key + if self.openai_client is not None: + self.openai_client.api_key = api_key + if self.openai_async_client is not None: + self.openai_async_client.api_key = api_key + @property def model_version(self) -> int: """Get the version of the model.""" @@ -641,23 +629,6 @@ async def get_lora_request_async(self) -> Any: async def get_message_token_len(self, messages: List[dict]) -> int: return await self.model.get_message_token_len.remote(messages) - def _get_multi_modal_inputs( - self, - *, - messages: List[dict] = None, - tools: Optional[List[dict]] = None, - input_ids: Optional[List[int]] = None, - ) -> Optional[dict[str, torch.Tensor]]: - if should_use_processor(self.model_path): - if self._mm_render is None: - self._mm_render = vLLMMultiModalRender( # TODO: support sglang - self.model_path, - ) - return self._mm_render.build_mm_input_for_training( - messages=messages, tools=tools, input_ids=input_ids - ) - return None - def get_openai_client(self) -> "openai.OpenAI": """Get the openai client. @@ -705,40 +676,6 @@ def chat_completions(*args, **kwargs): return response self.openai_client.chat.completions.create = chat_completions - elif self.enable_history: - # add a decorator to the openai client to record history - - ori_create = self.openai_client.chat.completions.create - - def record_chat_completions(*args, **kwargs): - logprobs = kwargs.pop("logprobs", True) - extra_body = dict(kwargs.pop("extra_body", {})) - if self.config.enable_thinking is not None: - chat_template_kwargs = dict(extra_body.get("chat_template_kwargs", {})) - chat_template_kwargs["enable_thinking"] = self.config.enable_thinking - extra_body["chat_template_kwargs"] = chat_template_kwargs - extra_body["return_token_ids"] = True - if self.config.enable_return_routed_experts: - extra_body["return_routed_experts"] = True - self._assert_openai_routed_experts_request_supported(extra_body, kwargs) - response = ori_create(*args, extra_body=extra_body, logprobs=logprobs, **kwargs) - if kwargs.get("stream", False): - return HistoryRecordingStream(response, self.history, is_async=False) - messages = args[-2] if len(args) > 2 else kwargs.get("messages") - tools = kwargs.get("tools", None) - multi_modal_inputs = self._get_multi_modal_inputs( - messages=messages, tools=tools, input_ids=response.prompt_token_ids - ) - self.history.extend( - convert_api_output_to_experience( - response, - multi_modal_inputs=multi_modal_inputs, - routed_experts_layout=self._routed_experts_layout, - ) - ) - return response - - self.openai_client.chat.completions.create = record_chat_completions setattr(self.openai_client, "model_path", self.config.model_path) return self.openai_client @@ -784,42 +721,6 @@ async def chat_completions(*args, **kwargs): return response self.openai_async_client.chat.completions.create = chat_completions - elif self.enable_history: - # add a decorator to the openai client to record history - - ori_create = self.openai_async_client.chat.completions.create - - async def record_chat_completions(*args, **kwargs): - logprobs = kwargs.pop("logprobs", True) - extra_body = dict(kwargs.pop("extra_body", {})) - if self.config.enable_thinking is not None: - chat_template_kwargs = dict(extra_body.get("chat_template_kwargs", {})) - chat_template_kwargs["enable_thinking"] = self.config.enable_thinking - extra_body["chat_template_kwargs"] = chat_template_kwargs - extra_body["return_token_ids"] = True - if self.config.enable_return_routed_experts: - extra_body["return_routed_experts"] = True - self._assert_openai_routed_experts_request_supported(extra_body, kwargs) - response = await ori_create( - *args, extra_body=extra_body, logprobs=logprobs, **kwargs - ) - if kwargs.get("stream", False): - return HistoryRecordingStream(response, self.history, is_async=True) - messages = args[-2] if len(args) > 2 else kwargs.get("messages") - tools = kwargs.get("tools", None) - multi_modal_inputs = self._get_multi_modal_inputs( - messages=messages, tools=tools, input_ids=response.prompt_token_ids - ) - self.history.extend( - convert_api_output_to_experience( - response, - multi_modal_inputs=multi_modal_inputs, - routed_experts_layout=self._routed_experts_layout, - ) - ) - return response - - self.openai_async_client.chat.completions.create = record_chat_completions # get model_path from the sync openai client to avoid async call here setattr(self.openai_async_client, "model_path", self.config.model_path) return self.openai_async_client @@ -876,8 +777,29 @@ async def sync_model_weights( # update the model path after syncing weights for tinker engine self._model_path = await self.model.get_model_path.remote() - def extract_experience_from_history(self, clear_history: bool = True) -> List[Experience]: + def extract_experience_from_history( + self, clear_history: bool = True, record_key: Optional[str] = None + ) -> List[Experience]: """Extract experiences from the history.""" + if self.config.enable_recording: + if self.model is None: + raise ValueError("Recording extraction requires an inference model actor.") + record_key = record_key or self._api_key + if record_key is None: + raise ValueError("record_key is required when recording is enabled.") + exps = ray.get( + self.model.extract_experience_from_history.remote( + record_key=record_key, + clear_history=False, + ) + ) + offset = self.recording_history_offsets.get(record_key, 0) + if offset > len(exps): + offset = 0 + new_exps = exps[offset:] + if clear_history: + self.recording_history_offsets[record_key] = len(exps) + return new_exps if not self.enable_history: raise ValueError("History recording is not enabled.") exps = [exp for exp in self.history] @@ -896,6 +818,7 @@ async def clean_workflow_state(self) -> None: async with self.state_lock: self.workflow_state = {} self.history.clear() + self.recording_history_offsets.clear() async def shutdown(self) -> None: """Shutdown all underlying model actors cleanly.""" @@ -917,4 +840,5 @@ def clone_with_isolated_history(self) -> "ModelWrapper": new_wrapper.openai_async_client = None new_wrapper.openai_client = None new_wrapper.history = [] + new_wrapper.recording_history_offsets = {} return new_wrapper diff --git a/trinity/common/models/vllm_model.py b/trinity/common/models/vllm_model.py index eae146a5792..d57a2f0ca16 100644 --- a/trinity/common/models/vllm_model.py +++ b/trinity/common/models/vllm_model.py @@ -229,8 +229,8 @@ async def chat( record_key (Optional[str]): Recording identity for the in-vLLM recorder (the MemoryStore group key). Propagated to ``generate`` via ``record_key_ctx`` so the recorder stamps it - into ``info["record_key"]`` without an HTTP hop. None leaves it - to the engine request id (fallback grouping). + into ``info["record_key"]`` without an HTTP hop. None skips + recording. kwargs (dict): A dictionary of sampling parameters. Returns: @@ -372,6 +372,18 @@ async def logprobs( # type: ignore [override] dtype=torch.float32, ) + async def extract_experience_from_history( + self, record_key: str, clear_history: bool = True + ) -> List[Experience]: + """Extract recorded experiences from the in-vLLM recorder store.""" + if self.recorder is None: + raise ValueError("Recording is not enabled for this vLLM model.") + await self.recorder.flush() + exps = await self.recorder.store.get_record_experiences(record_key) + if clear_history: + await self.recorder.store.delete_record_experiences(record_key) + return exps + async def add_lora_adapter(self, lora_request: Any) -> int: """Add a LoRA adapter to the vLLM engine. diff --git a/trinity/common/workflows/step_wise_workflow.py b/trinity/common/workflows/step_wise_workflow.py index 1042866063a..1eee6bdefd5 100644 --- a/trinity/common/workflows/step_wise_workflow.py +++ b/trinity/common/workflows/step_wise_workflow.py @@ -12,9 +12,10 @@ def __init__( self, *, task: Task, model: ModelWrapper, auxiliary_models=None, use_openai_client=True ): super().__init__(task=task, model=model, auxiliary_models=auxiliary_models) - assert model.enable_history, ( - "Rollout Model must have history enabled for step-wise rewards, please " - "set `explorer.rollout_model.enable_history` to `True` in your config." + assert model.enable_history or model.config.enable_recording, ( + "Rollout Model must have history or recording enabled for step-wise rewards, " + "please set `explorer.rollout_model.enable_history` or " + "`explorer.rollout_model.enable_recording` to `True` in your config." ) # use the rollout model's OpenAI client to write your agent application if use_openai_client: @@ -121,9 +122,10 @@ def __init__( self, *, task: Task, model: ModelWrapper, auxiliary_models=None, use_openai_client=True ): super().__init__(task=task, model=model, auxiliary_models=auxiliary_models) - assert model.enable_history, ( - "Rollout Model must have history enabled for step-wise rewards, please " - "set `explorer.rollout_model.enable_history` to `True` in your config." + assert model.enable_history or model.config.enable_recording, ( + "Rollout Model must have history or recording enabled for step-wise rewards, " + "please set `explorer.rollout_model.enable_history` or " + "`explorer.rollout_model.enable_recording` to `True` in your config." ) # use the rollout model's OpenAI client to write your agent application if use_openai_client: diff --git a/trinity/explorer/workflow_runner.py b/trinity/explorer/workflow_runner.py index e1003366e74..2fee317ca8e 100644 --- a/trinity/explorer/workflow_runner.py +++ b/trinity/explorer/workflow_runner.py @@ -67,6 +67,7 @@ def __init__( for index, auxiliary_model_id in enumerate(auxiliary_model_ids or []) ] self.workflow_instance: Workflow = None + self.rollout_model_id = rollout_model_id self.runner_id = runner_id self.runner_state = { "workflow_id": None, @@ -103,9 +104,17 @@ async def prepare(self) -> None: def is_alive(self): return True - def _create_workflow_instance(self, task: Task) -> Workflow: + def _build_record_key(self, task: Task, run_index: int) -> str: + return f"{task.batch_id}/{task.task_id}/{run_index}" + + def _set_record_key(self, model_wrapper: ModelWrapper, record_key: Optional[str]) -> None: + if self._enable_recording() and record_key is not None: + model_wrapper.set_api_key(record_key) + + def _create_workflow_instance(self, task: Task, record_key: Optional[str] = None) -> Workflow: if task.workflow is None: raise ValueError("Workflow is not set in the task.") + self._set_record_key(self.model_wrapper, record_key) if ( self.workflow_instance is None or not self.workflow_instance.__class__ == task.workflow @@ -131,17 +140,21 @@ async def _run_workflow(self, workflow_instance: Workflow) -> List[Experience]: exps = workflow_instance.run() return exps - def _create_isolated_workflow_instance(self, task: Task) -> Workflow: + def _create_isolated_workflow_instance( + self, task: Task, record_key: Optional[str] = None + ) -> Tuple[Workflow, ModelWrapper]: + model_wrapper = ( + self.model_wrapper.clone_with_isolated_history() + if (self.config.explorer.rollout_model.enable_history or self._enable_recording()) + else self.model_wrapper + ) + self._set_record_key(model_wrapper, record_key) wf = task.to_workflow( - ( - self.model_wrapper.clone_with_isolated_history() - if self.config.explorer.rollout_model.enable_history - else self.model_wrapper - ), + model_wrapper, self.auxiliary_model_wrappers, ) wf.enable_recording = self._enable_recording() - return wf + return wf, model_wrapper def _build_execution_result( self, @@ -208,8 +221,17 @@ async def _run_parallel_runs( async def run_single( i: int, ) -> Tuple[bool, List[Experience], Optional[Dict[str, float]], Optional[str]]: - workflow = self._create_isolated_workflow_instance(task) - return await self._execute_single_run(workflow, task, i, run_id_base) + run_index = run_id_base + i + record_key = self._build_record_key(task, run_index) + workflow, model_wrapper = self._create_isolated_workflow_instance(task, record_key) + return await self._execute_single_run( + workflow, + task, + i, + run_id_base, + model_wrapper=model_wrapper, + record_key=record_key, + ) if collect_partial_runs: if use_threads: @@ -262,10 +284,14 @@ async def _execute_single_run( task: Task, run_index: int, run_id_base: int, + model_wrapper: Optional[ModelWrapper] = None, + record_key: Optional[str] = None, ) -> Tuple[bool, List[Experience], Optional[Dict[str, float]], Optional[str]]: st = time.time() - await self.model_wrapper.clean_workflow_state() - self.runner_state["workflow_id"] = f"{task.batch_id}/{task.task_id}/{run_index}" + model_wrapper = model_wrapper or self.model_wrapper + self._set_record_key(model_wrapper, record_key) + await model_wrapper.clean_workflow_state() + self.runner_state["workflow_id"] = self._build_record_key(task, run_id_base + run_index) self.runner_state["terminate_time"] = None self.runner_state["begin_time"] = st try: @@ -296,11 +322,13 @@ async def _run_task( ) -> RunnerExecutionResult: """Init workflow from the task and run it.""" if task.workflow.can_repeat: - workflow_instance = self._create_workflow_instance(task) + record_key = self._build_record_key(task, run_id_base) + workflow_instance = self._create_workflow_instance(task, record_key=record_key) workflow_instance.set_repeat_times(repeat_times, run_id_base) st = time.time() + self._set_record_key(self.model_wrapper, record_key) await self.model_wrapper.clean_workflow_state() - self.runner_state["workflow_id"] = f"{task.batch_id}/{task.task_id}/{run_id_base}" + self.runner_state["workflow_id"] = record_key self.runner_state["terminate_time"] = None self.runner_state["begin_time"] = st exps = await self._run_workflow(workflow_instance) @@ -333,8 +361,16 @@ async def _sequential_run( ) -> RunnerExecutionResult: results = [] for i in range(repeat_times): - workflow = self._create_workflow_instance(task) - result = await self._execute_single_run(workflow, task, i, run_id_base) + run_index = run_id_base + i + record_key = self._build_record_key(task, run_index) + workflow = self._create_workflow_instance(task, record_key=record_key) + result = await self._execute_single_run( + workflow, + task, + i, + run_id_base, + record_key=record_key, + ) results.append(result) if collect_partial_runs: continue From bf931de9b49928a2e640ace50004ed39e30677ea Mon Sep 17 00:00:00 2001 From: pxc Date: Thu, 25 Jun 2026 20:40:40 +0800 Subject: [PATCH 17/84] sglang self recording experiences --- trinity/common/models/sglang_model.py | 113 +++++++- .../models/sglang_patch/recording/__init__.py | 41 +++ .../models/sglang_patch/recording/models.py | 170 +++++++++++ .../models/sglang_patch/recording/recorder.py | 271 ++++++++++++++++++ .../models/sglang_patch/recording/server.py | 80 ++++++ .../models/sglang_patch/server_patch.py | 31 +- 6 files changed, 697 insertions(+), 9 deletions(-) create mode 100644 trinity/common/models/sglang_patch/recording/__init__.py create mode 100644 trinity/common/models/sglang_patch/recording/models.py create mode 100644 trinity/common/models/sglang_patch/recording/recorder.py create mode 100644 trinity/common/models/sglang_patch/recording/server.py diff --git a/trinity/common/models/sglang_model.py b/trinity/common/models/sglang_model.py index b2602daba5a..d5a598daefa 100644 --- a/trinity/common/models/sglang_model.py +++ b/trinity/common/models/sglang_model.py @@ -28,17 +28,26 @@ def __init__(self, server_url: str, api_key: Optional[str], logger: Logger): self.api_key = api_key self.logger = logger + def _auth_header(self, api_key_override: Optional[str]) -> str: + # The record_key travels as the bearer so the server-side + # RecordingIdentityMiddleware stamps it into ``record_key_ctx``. Falls + # back to the client's api_key (which, on the Trinity Ray-direct path, + # equals the record_key injected by ModelWrapper). + token = api_key_override if api_key_override is not None else self.api_key + return f"Bearer {token}" if token else "" + async def _server_call( self, method: Literal["GET", "POST"], endpoint: str, payload: Optional[dict] = None, timeout: float = 60, + api_key_override: Optional[str] = None, ) -> dict: async with httpx.AsyncClient( headers={ "Content-Type": "application/json; charset=utf-8", - "Authorization": f"Bearer {self.api_key}" if self.api_key else "", + "Authorization": self._auth_header(api_key_override), } ) as client: url = f"{self.server_url}{endpoint}" @@ -191,7 +200,9 @@ async def update_weights_from_disk( ) return success - async def generate(self, input_ids: List[int], **kwargs) -> Sequence[dict[str, Any]]: + async def generate( + self, input_ids: List[int], record_key: Optional[str] = None, **kwargs + ) -> Sequence[dict[str, Any]]: sampling_params = { "n": kwargs.get("n", 1), "temperature": kwargs.get("temperature"), @@ -219,6 +230,7 @@ async def generate(self, input_ids: List[int], **kwargs) -> Sequence[dict[str, A "/generate", payload, timeout=kwargs.get("timeout", 300), + api_key_override=record_key, ) if isinstance(response, dict) and response.get("error"): raise RuntimeError(f"Failed to generate with SGLang: {response['error']}") @@ -252,6 +264,8 @@ def __init__( self.api_server_port: Optional[int] = None self.api_server: Optional[asyncio.Task[None]] = None self.api_client: Optional[SGLangClient] = None + self.recorder = None + self.record_store = None self.synchronizer = None self.state_dict_meta: List[Tuple[str, str, Tuple]] = [] self.model_version = 0 @@ -365,7 +379,22 @@ def _extract_routed_experts(self, routed_experts_str: str, total_tokens: int) -> assert routed_experts is not None return routed_experts - async def generate(self, prompt: str, lora_request=None, **kwargs) -> Sequence[Experience]: + async def generate( + self, + prompt: str, + lora_request=None, + record_key: Optional[str] = None, + **kwargs, + ) -> Sequence[Experience]: + """Generate a response from the provided prompt in async. + + When ``record_key`` is set, it is sent as the Authorization bearer so the + server-side recorder groups this turn under that key (the api_key doubles + as the record_key on the Trinity path). The returned experiences are the + client-side copy; the recorded copy is written to the in-process store by + the engine-level recorder and drained via + ``extract_experience_from_history`` — mirroring vLLM. + """ assert self.api_client is not None, "API client must be initialized before calling generate" if self.tokenizer is None: await self._initialize_tokenizer() @@ -379,6 +408,7 @@ async def generate(self, prompt: str, lora_request=None, **kwargs) -> Sequence[E return_logprob = logprobs is not None and logprobs is not False responses = await self.api_client.generate( input_ids=prompt_token_ids, + record_key=record_key, n=kwargs.get("n", 1), temperature=kwargs.get("temperature", self.config.temperature), top_p=kwargs.get("top_p", self.config.top_p), @@ -441,18 +471,24 @@ async def chat( record_key: Optional[str] = None, **kwargs, ) -> Sequence[Experience]: - # ``record_key`` is the in-vLLM recorder's group key; SGLang does not - # support recording, so accept and ignore it here. - del record_key + # ``record_key`` is propagated to ``generate`` so the server-side recorder + # groups this turn under the caller's key (sent as the Authorization + # bearer, same as vLLM's RecordingIdentityMiddleware path). if self.tokenizer is None: await self._initialize_tokenizer() normalized_messages = self._normalize_chat_messages(messages) prompt = self.apply_chat_template(self.tokenizer, normalized_messages) - return await self.generate(prompt=prompt, lora_request=lora_request, **kwargs) + return await self.generate( + prompt=prompt, lora_request=lora_request, record_key=record_key, **kwargs + ) async def logprobs(self, token_ids: List[int], **kwargs) -> torch.Tensor: raise NotImplementedError("SGLangRolloutModel does not support logprobs.") + # NOTE: if implemented later, the auxiliary forward must avoid being + # recorded. Unlike vLLM, ``skip_recording_ctx`` does NOT cross the HTTP + # hop to the server; instead omit ``record_key`` for that call so the + # server-side recorder skips it (record_key is None -> no record). async def convert_messages_to_experience( self, @@ -464,6 +500,24 @@ async def convert_messages_to_experience( "SGLangRolloutModel does not support convert_messages_to_experience." ) + async def extract_experience_from_history( + self, record_key: str, clear_history: bool = True + ) -> List[Experience]: + """Extract recorded experiences from the in-process recorder store. + + Mirrors ``vLLMRolloutModel.extract_experience_from_history``: the + recorder and store live in-process (the SGLang server runs as an + asyncio task in the same loop as this model), so extraction is a direct + store lookup with no HTTP hop. + """ + if self.recorder is None: + raise ValueError("Recording is not enabled for this SGLang model.") + await self.recorder.flush() + exps = await self.recorder.store.get_record_experiences(record_key) + if clear_history: + await self.recorder.store.delete_record_experiences(record_key) + return exps + def _get_api_server_exit_reason(self) -> Optional[str]: if self.api_server is None or not self.api_server.done(): return None @@ -497,6 +551,36 @@ async def run_api_server(self) -> bool: if self.api_server_host is None or self.api_server_port is None: self.api_server_host, self.api_server_port = self.get_available_address() + + # When recording is on, own the recorder/store here so they can be drained + # in-process via ``extract_experience_from_history``. They are wired onto + # the embedded server (engine wrap + middleware + query routes) inside + # ``get_api_server`` -> ``setup_sglang_recording``. + record_store = None + recorder = None + routed_experts_layout = None + if self.config.enable_recording: + from trinity.common.models.recording.recorder import Recorder + from trinity.common.models.recording.store import MemoryStore + from trinity.common.models.sglang_patch.recording.models import ( + build_sglang_experience, + ) + + record_store = MemoryStore() + recorder = Recorder( + store=record_store, + build_experiences=build_sglang_experience, + enabled=True, + rank=0, + engine_client=None, + ) + # Decode layout for base64-str routed experts (None for non-MoE + # models -> routed_experts stays None in the recorded Experience). + if self.config.enable_return_routed_experts: + layout = self._get_routed_experts_layout() + if layout is not None: + routed_experts_layout = (layout[0], layout[1]) + self.api_server = get_api_server( host=self.api_server_host, port=self.api_server_port, @@ -518,8 +602,16 @@ async def run_api_server(self) -> bool: master_addr=self.master_addr, master_port=self.master_port, enable_return_routed_experts=self.config.enable_return_routed_experts, + enable_recording=self.config.enable_recording, + recorder=recorder, + record_store=record_store, + routed_experts_layout=routed_experts_layout, logger=self.logger, ) + # ``setup_sglang_recording`` (called inside get_api_server) owns the + # recorder handle we passed in; keep references for in-process draining. + self.recorder = recorder + self.record_store = record_store server_url = f"http://{self.api_server_host}:{self.api_server_port}" self.api_client = SGLangClient( server_url=server_url, @@ -550,6 +642,13 @@ async def shutdown(self) -> None: self.api_server = None self.api_client = None self._has_weight_update_group = False + if self.recorder is not None: + try: + await self.recorder.stop() + except Exception as e: + self.logger.error("Error while stopping SGLang recorder: %s", e) + self.recorder = None + self.record_store = None async def sync_model_weights( self, diff --git a/trinity/common/models/sglang_patch/recording/__init__.py b/trinity/common/models/sglang_patch/recording/__init__.py new file mode 100644 index 00000000000..260314194c1 --- /dev/null +++ b/trinity/common/models/sglang_patch/recording/__init__.py @@ -0,0 +1,41 @@ +"""SGLang generation recording adapter. + +Re-exports the SGLang-specific pieces (``build_sglang_experience``, +``create_sglang_recorder``, ``setup_sglang_recording``) and the engine-agnostic +core symbols shared with the vLLM recording path +(``trinity.common.models.recording``). +""" +from trinity.common.models.recording.context import ( # noqa: F401 + RecordingIdentityMiddleware, + get_recording_record_key, + record_key_ctx, + skip_recording_ctx, +) +from trinity.common.models.recording.query import query_router # noqa: F401 +from trinity.common.models.recording.recorder import Recorder # noqa: F401 +from trinity.common.models.recording.store import MemoryStore, RecordStore # noqa: F401 +from trinity.common.models.sglang_patch.recording.models import ( # noqa: F401 + build_sglang_experience, +) +from trinity.common.models.sglang_patch.recording.recorder import ( # noqa: F401 + create_sglang_recorder, + patch_tokenizer_manager_for_recording, +) +from trinity.common.models.sglang_patch.recording.server import ( # noqa: F401 + setup_sglang_recording, +) + +__all__ = [ + "MemoryStore", + "RecordStore", + "Recorder", + "RecordingIdentityMiddleware", + "build_sglang_experience", + "create_sglang_recorder", + "get_recording_record_key", + "patch_tokenizer_manager_for_recording", + "query_router", + "record_key_ctx", + "setup_sglang_recording", + "skip_recording_ctx", +] diff --git a/trinity/common/models/sglang_patch/recording/models.py b/trinity/common/models/sglang_patch/recording/models.py new file mode 100644 index 00000000000..6dcc2a277ab --- /dev/null +++ b/trinity/common/models/sglang_patch/recording/models.py @@ -0,0 +1,170 @@ +"""Build Trinity ``Experience`` objects from a finished SGLang ``ret``. + +Mirrors ``trinity/common/models/vllm_patch/recording/models.py`` but for the +SGLang output shape. A SGLang ``ret`` is the dict (or list of dicts for ``n > 1`` +/ batch) yielded by ``tokenizer_manager.generate_request``. Each item carries +``output_ids``, ``text`` and a ``meta_info`` dict with ``id``, ``prompt_tokens``, +``output_token_logprobs``, ``routed_experts`` and ``weight_version``. + +The finished ``ret`` also carries ``prompt_token_ids`` because the recorder +wrapper forces ``obj.return_prompt_token_ids = True`` (see ``recorder.py``), so +the recorded Experience gets the real prompt tokens without reconstructing them +from the request. + +Field mapping (SGLang ``ret`` -> ``Experience``): + meta_info.id -> eid.suffix (traceability; task/run/reward assigned by + ``MemoryStore.update_reward_by_record_key`` at consume) + record_key -> info["record_key"] (the MemoryStore group key) + sample index -> info["sample_index"] (position within the n set) + prompt_token_ids -> tokens (prompt) + prompt_length + output_ids -> tokens (response) + output_token_logprobs -> Experience.logprobs (flat ``[resp_length]``; SGLang + returns ``(logprob, *_)`` tuples per token) + routed_experts -> Experience.routed_experts (uint8 tensor, decoded with + the model's ``(num_layers, topk)`` layout when base64-str) + meta_info.weight_version -> info["model_version"] +""" +from typing import Any, List, Optional, Tuple + +import torch + +from trinity.common.experience import EID, Experience +from trinity.common.models.experience_extraction import decode_sglang_routed_experts + + +def _extract_output_logprobs(meta_info: dict) -> List[float]: + """Pull the chosen-token logprob at each response position. + + SGLang ``output_token_logprobs`` is a list of ``(logprob, *_)`` tuples (one + per generated token). Mirrors ``SGLangRolloutModel._extract_output_logprobs``. + """ + output_token_logprobs = meta_info.get("output_token_logprobs") or [] + return [float(logprob) for logprob, *_ in output_token_logprobs] + + +def _extract_routed_experts( + routed_experts_value: Any, + total_tokens: int, + routed_experts_layout: Optional[Tuple[int, int]], +) -> Optional[torch.Tensor]: + if routed_experts_value is None: + return None + if isinstance(routed_experts_value, str): + if routed_experts_layout is None: + return None + return decode_sglang_routed_experts( + routed_experts_value, + total_tokens, + layout=routed_experts_layout, + ) + return torch.tensor(routed_experts_value, dtype=torch.uint8) + + +def build_sglang_experience( + ret: Any, + record_key: Optional[str], + *, + rank: int, + timestamp: str, + model_version: Optional[Any] = None, + include_routed_experts: bool = True, + routed_experts_layout: Optional[Tuple[int, int]] = None, +) -> List[Experience]: + """Build Trinity ``Experience`` objects from a finished SGLang ``ret``. + + One experience per output (``n > 1`` / batch is captured in full). Each + shares ``eid.suffix = meta_info.id`` and ``info["record_key"] = record_key``; + ``info["sample_index"]`` distinguishes samples within the group. + + Args: + ret: A finished SGLang result — a dict, or a list of dicts for ``n > 1`` + / batch. Each dict has ``output_ids``/``text``/``meta_info`` and + (when the wrapper forced ``return_prompt_token_ids``) a + ``prompt_token_ids`` list. + record_key: The recording identity (Authorization bearer / record key); + the MemoryStore group key. + rank: Data-parallel serving rank. + timestamp: UTC ISO-8601 string (caller-stamped to keep this pure). + model_version: Checkpoint version fallback; overridden by + ``meta_info.weight_version`` when present. + include_routed_experts: Whether routed experts should be copied. + routed_experts_layout: ``(num_layers, topk)`` for decoding base64-str + routed experts (from ``BaseInferenceModel._get_routed_experts_layout``). + + Returns: + One ``Experience`` per non-degenerate output. Empty list if the request + had no prompt tokens or no output with response tokens. + """ + ret_list = ret if isinstance(ret, list) else [ret] + if not ret_list: + return [] + + experiences: List[Experience] = [] + for sample_index, item in enumerate(ret_list): + if not isinstance(item, dict): + continue + meta_info = item.get("meta_info") or {} + prompt_token_ids = list(item.get("prompt_token_ids") or []) + if not prompt_token_ids: + # No prompt tokens captured (return_prompt_token_ids not honored); + # cannot build a valid single-turn Experience. + continue + + response_token_ids = list(item.get("output_ids") or []) + if not response_token_ids: + # Fall back to re-encoding text if the engine omitted output_ids. + response_text = item.get("text") or "" + if response_text: + # The recorder runs in-process but has no tokenizer handle here; + # output_ids should normally be present, so just skip otherwise. + response_token_ids = [] + if not response_token_ids: + continue + + prompt_length = int(meta_info.get("prompt_tokens") or len(prompt_token_ids)) + # Guard against an inconsistent count: prefer the real token list length. + if prompt_length <= 0 or prompt_length > len(prompt_token_ids): + prompt_length = len(prompt_token_ids) + + response_logprobs = torch.tensor( + _extract_output_logprobs(meta_info), + dtype=torch.float32, + ) + + routed_experts = None + if include_routed_experts: + routed_experts = _extract_routed_experts( + meta_info.get("routed_experts"), + total_tokens=len(prompt_token_ids) + len(response_token_ids), + routed_experts_layout=routed_experts_layout, + ) + + request_id = str(meta_info.get("id") or "") + resolved_model_version = meta_info.get("weight_version") + if resolved_model_version is None: + resolved_model_version = model_version + + eid = EID(suffix=request_id) + info = { + "request_id": request_id, + "record_key": record_key, + "sample_index": sample_index, + "rank": rank, + "timestamp": timestamp, + "endpoint": "sglang", + "model_version": resolved_model_version, + } + + experiences.append( + Experience( + eid=eid, + tokens=torch.tensor(prompt_token_ids + response_token_ids, dtype=torch.int32), + logprobs=response_logprobs, + prompt_length=prompt_length, + prompt_text=item.get("prompt_text"), + response_text=item.get("text"), + routed_experts=routed_experts, + info=info, + ) + ) + return experiences diff --git a/trinity/common/models/sglang_patch/recording/recorder.py b/trinity/common/models/sglang_patch/recording/recorder.py new file mode 100644 index 00000000000..9f2a760d06f --- /dev/null +++ b/trinity/common/models/sglang_patch/recording/recorder.py @@ -0,0 +1,271 @@ +"""Engine-level wrap that records finished SGLang turns into the shared store. + +Mirrors ``trinity/common/models/vllm_patch/recording/recorder.py`` but adapts to +the SGLang output path. vLLM wraps ``engine_client.generate`` (the in-process +engine boundary); SGLang's single convergence point for ``/generate``, +``/v1/chat/completions`` and ``/invocations`` is +``tokenizer_manager.generate_request`` (an async generator yielding ``ret`` +dicts), so that is what we wrap here — instance-level, idempotent. + +Two adaptations forced by the SGLang shape (see the plan for detail): + +1. Trigger on the **finished yield**, not on generator exhaustion. SGLang's + non-stream ``/generate`` handler pulls exactly one item via ``__anext__()`` + and never exhausts the generator, so a "record after the loop" trigger would + never fire. We detect finished via ``ret["meta_info"]["finish_reason"] is not + None`` and ``schedule_record`` *before* yielding that finished ``ret``. + +2. Force ``obj.return_logprob = True`` and ``obj.return_prompt_token_ids = True`` + so the finished ``ret`` always carries the chosen-token logprobs and the full + prompt token ids (the latter is stashed onto ``out_dict`` by + ``tokenizer_manager`` only when ``return_prompt_token_ids`` is set). This is + transparent to OpenAI clients (the chat serving layer gates its + ``prompt_token_ids`` on the separate ``return_token_ids`` flag) and only adds + an ignored field to ``/generate`` JSON responses. +""" +import functools +import logging +from typing import Any, List, Optional, Tuple + +from trinity.common.models.recording.context import record_key_ctx +from trinity.common.models.recording.recorder import Recorder +from trinity.common.models.recording.store import MemoryStore, RecordStore +from trinity.common.models.sglang_patch.recording.models import build_sglang_experience + +#: Guard attribute marking the wrapped generate_request (mirrors vLLM's style). +_PATCHED_FLAG = "__patched_sglang_recording__" +TRINITY_RECORDER_ATTR = "trinity_recorder" +TRINITY_RECORD_STORE_ATTR = "trinity_record_store" + + +def _get_obj(args, kwargs): + """Extract the ``GenerateReqInput``/``EmbeddingReqInput`` argument. + + ``generate_request(self, obj, request=None)`` is wrapped as an instance + attribute, so ``self`` is absent and ``args`` map 1:1 to the protocol + (``obj`` is ``args[0]``; ``obj`` may also be passed as ``obj=``). + """ + if "obj" in kwargs: + return kwargs["obj"] + if args: + return args[0] + return None + + +def _force_record_fields(obj: Any) -> None: + """Force logprob + prompt-token-id capture for recording (transparent).""" + if obj is None: + return + # return_logprob may be a list for batched requests; broadcast True. + if hasattr(obj, "return_logprob"): + cur = getattr(obj, "return_logprob", None) + if isinstance(cur, list): + obj.return_logprob = [True] * len(cur) + else: + obj.return_logprob = True + if hasattr(obj, "return_prompt_token_ids"): + obj.return_prompt_token_ids = True + + +def _normalize_ret(out: Any) -> List[dict]: + """A ``ret`` is a dict (n=1) or a list of dicts (n>1 / batch).""" + if isinstance(out, list): + return [item for item in out if isinstance(item, dict)] + if isinstance(out, dict): + return [out] + return [] + + +def _is_finished(out: Any) -> bool: + """True if any output carries a non-None ``finish_reason``. + + ``tokenizer_manager`` sets ``meta_info["finish_reason"] = + recv_obj.finished_reasons[i]`` and ``finished_reasons[i]`` is ``None`` until + the request is done, so this is the reliable finished signal. + """ + for item in _normalize_ret(out): + meta_info = item.get("meta_info") or {} + if meta_info.get("finish_reason") is not None: + return True + return False + + +def _monotonic_extend_list(acc: list, cur: Optional[list]) -> list: + """Replace ``acc`` with ``cur`` if ``cur`` is cumulative (starts with ``acc`` + as a prefix), otherwise extend. Handles both cumulative-streaming chunks + (each carries the full-so-far ids) and incremental-streaming deltas. + """ + if not cur: + return acc + if not acc: + return list(cur) + if len(cur) >= len(acc) and list(cur[: len(acc)]) == acc: + return list(cur) # cumulative: cur supersedes acc + return acc + list(cur) # delta: extend + + +def _monotonic_extend_text(acc: str, cur: Optional[str]) -> str: + """Same discipline for ``text`` (string prefix check). ``cur`` may be ``None`` + for non-incremental intermediate chunks (deferred text).""" + if not cur: + return acc + if not acc: + return cur + if cur.startswith(acc): + return cur # cumulative + return acc + cur # delta + + +def _accumulate_ret(state: dict, order: list, out: Any) -> Tuple[dict, list]: + """Merge a yielded ``ret`` into the per-output accumulator. + + ``state`` maps output index -> accumulated fields; ``order`` preserves + first-seen order so the reconstructed ``ret`` keeps sample indexing. + """ + items = _normalize_ret(out) + for idx, item in enumerate(items): + meta_info = item.get("meta_info") or {} + if idx not in state: + state[idx] = { + "output_ids": [], + "text": "", + "output_token_logprobs": [], + "routed_experts": None, + "prompt_token_ids": None, + "meta_info": {}, + } + order.append(idx) + acc = state[idx] + acc["output_ids"] = _monotonic_extend_list(acc["output_ids"], item.get("output_ids")) + acc["text"] = _monotonic_extend_text(acc["text"], item.get("text")) + acc["output_token_logprobs"] = _monotonic_extend_list( + acc["output_token_logprobs"], + (meta_info.get("output_token_logprobs")), + ) + routed = meta_info.get("routed_experts") + if routed is not None: + acc["routed_experts"] = routed # latest non-None (final chunk is full) + prompt_ids = item.get("prompt_token_ids") + if prompt_ids: + acc["prompt_token_ids"] = list(prompt_ids) + # Keep the latest meta_info (carries id/finish_reason/weight_version/ + # prompt_tokens); output_token_logprobs/routed_experts are overridden + # below from the accumulated fields. + acc["meta_info"] = meta_info + return state, order + + +def _build_ret(state: dict, order: list) -> List[dict]: + """Reconstruct a finished ``ret`` list from accumulated per-output state.""" + out_list: List[dict] = [] + for idx in order: + acc = state[idx] + meta_info = dict(acc["meta_info"]) + # Override with the fully-accumulated fields (streaming deltas merged). + meta_info["output_token_logprobs"] = acc["output_token_logprobs"] + if acc["routed_experts"] is not None: + meta_info["routed_experts"] = acc["routed_experts"] + out_list.append( + { + "output_ids": acc["output_ids"], + "text": acc["text"], + "prompt_token_ids": acc["prompt_token_ids"], + "meta_info": meta_info, + } + ) + return out_list + + +def create_sglang_recorder( + tokenizer_manager, + logger: logging.Logger, + *, + store: Optional[RecordStore] = None, + recorder: Optional[Recorder] = None, + enabled: bool = True, + routed_experts_layout: Optional[Tuple[int, int]] = None, +) -> Recorder: + """Create/accept and install a SGLang-backed recorder on ``tokenizer_manager``. + + The caller (``SGLangRolloutModel``) may pre-create and own ``recorder`` so it + can drain it in-process via ``extract_experience_from_history``; this + function wires it onto the ``tokenizer_manager`` (engine_client) and patches + ``generate_request``. Idempotent. + """ + existing = getattr(tokenizer_manager, TRINITY_RECORDER_ATTR, None) + if existing is not None: + return existing + + if recorder is None: + recorder = Recorder( + store=store or MemoryStore(), + build_experiences=build_sglang_experience, + enabled=enabled, + engine_client=tokenizer_manager, + ) + else: + # The model owns the recorder; let it read model_version off the engine + # if needed (build_sglang_experience prefers meta_info.weight_version). + recorder.engine_client = tokenizer_manager + if store is not None and recorder.store is None: + recorder.store = store + + patch_tokenizer_manager_for_recording( + tokenizer_manager, recorder, logger, routed_experts_layout=routed_experts_layout + ) + setattr(tokenizer_manager, TRINITY_RECORDER_ATTR, recorder) + setattr(tokenizer_manager, TRINITY_RECORD_STORE_ATTR, recorder.store) + return recorder + + +def patch_tokenizer_manager_for_recording( + tokenizer_manager, + recorder: "Recorder", + logger: logging.Logger, + *, + routed_experts_layout: Optional[Tuple[int, int]] = None, +) -> None: + """Wrap ``tokenizer_manager.generate_request`` in place to record turns. + + Instance-level: only this server's tokenizer_manager is affected. Must run + before the server starts serving (the serving objects hold the same instance, + so the wrap is inherited). + """ + current = getattr(tokenizer_manager, "generate_request", None) + if current is None: + raise RuntimeError( + "SGLang recording patch failed: tokenizer_manager.generate_request not found" + ) + if getattr(current, _PATCHED_FLAG, False): + return + + @functools.wraps(current) + async def _patched_generate_request(*args, **kwargs): + obj = _get_obj(args, kwargs) + if recorder.enabled and obj is not None: + _force_record_fields(obj) + + state: dict = {} + order: list = [] + # ``current`` is the original *bound* method captured pre-wrap, so it + # still resolves ``self`` correctly. Yields each ret unchanged. + async for out in current(*args, **kwargs): + if recorder.enabled: + state, order = _accumulate_ret(state, order, out) + # Trigger on the finished yield (not on generator exhaustion): the + # non-stream /generate consumer pulls only once via __anext__(). + if recorder.enabled and _is_finished(out): + record_key = record_key_ctx.get() + if record_key is not None and state: + reconstructed = _build_ret(state, order) + recorder.schedule_record( + reconstructed, + record_key, + include_routed_experts=True, + routed_experts_layout=routed_experts_layout, + ) + yield out + + setattr(_patched_generate_request, _PATCHED_FLAG, True) + tokenizer_manager.generate_request = _patched_generate_request + logger.info("Patched SGLang tokenizer_manager.generate_request for generation recording") diff --git a/trinity/common/models/sglang_patch/recording/server.py b/trinity/common/models/sglang_patch/recording/server.py new file mode 100644 index 00000000000..f8e5ffbf158 --- /dev/null +++ b/trinity/common/models/sglang_patch/recording/server.py @@ -0,0 +1,80 @@ +"""Wiring that installs SGLang generation recording onto the embedded HTTP server. + +Mirrors ``trinity/common/models/vllm_patch/recording/server.py:_setup_recording``: +(1) the engine wrap (``create_sglang_recorder``), (2) ``RecordingIdentityMiddleware`` +— an in-process ASGI middleware reading ``Authorization: Bearer `` +into a contextvar, and (3) ``query_router`` — ``/records/*`` endpoints. + +Called from ``sglang_patch.server_patch.get_api_server`` after the +``tokenizer_manager`` is created and **before** the uvicorn task starts serving, +so the middleware/router are mounted on ``app`` in time. The recorder and store +are owned by ``SGLangRolloutModel`` (passed in) so it can drain them in-process +via ``extract_experience_from_history``; they are also stashed on ``app.state`` +for the ``query_router`` HTTP drain path used by the coordinator. +""" +import logging +from typing import Optional, Tuple + +from trinity.common.models.recording.context import RecordingIdentityMiddleware +from trinity.common.models.recording.query import ( + RECORDER_STATE_ATTR, + STORE_STATE_ATTR, + query_router, +) +from trinity.common.models.recording.recorder import Recorder +from trinity.common.models.recording.store import RecordStore +from trinity.common.models.sglang_patch.recording.recorder import create_sglang_recorder + + +def _add_recording_middleware(app) -> None: + """Install recording middleware before serving, even if SGLang built the stack. + + Starlette rejects ``add_middleware`` after ``middleware_stack`` is built with + "Cannot add middleware after an application has started". Clearing the cached + stack lets Starlette rebuild it with our middleware on first request (same + defensive pattern as the vLLM recording patch). + """ + if getattr(app, "middleware_stack", None) is not None: + app.middleware_stack = None + app.add_middleware(RecordingIdentityMiddleware) + + +def setup_sglang_recording( + tokenizer_manager, + app, + logger: logging.Logger, + *, + recorder: Optional[Recorder] = None, + store: Optional[RecordStore] = None, + routed_experts_layout: Optional[Tuple[int, int]] = None, +) -> Recorder: + """Wire generation recording onto the in-construction SGLang server. + + Only called when recording is on. The recorder is started here (its flusher + task lives in the server's event loop, same loop as ``SGLangRolloutModel``). + """ + recorder = create_sglang_recorder( + tokenizer_manager, + logger, + store=store, + recorder=recorder, + enabled=True, + routed_experts_layout=routed_experts_layout, + ) + + # (2) in-process middleware: Authorization bearer -> record_key contextvar. + _add_recording_middleware(app) + + # (3) query routes mounted on the main app; OpenAI /v1/* surface untouched. + app.include_router(query_router) + + setattr(app.state, STORE_STATE_ATTR, recorder.store) + setattr(app.state, RECORDER_STATE_ATTR, recorder) + + recorder.start() + + logger.info( + "SGLang generation recording enabled: store=%s", + type(recorder.store).__name__, + ) + return recorder diff --git a/trinity/common/models/sglang_patch/server_patch.py b/trinity/common/models/sglang_patch/server_patch.py index d60a1c1594e..4c14ba332a2 100644 --- a/trinity/common/models/sglang_patch/server_patch.py +++ b/trinity/common/models/sglang_patch/server_patch.py @@ -4,7 +4,7 @@ import os import time from logging import Logger -from typing import Any, Callable, Coroutine, Dict, List, Optional +from typing import Any, Callable, Coroutine, Dict, List, Optional, Tuple import uvicorn from fastapi import FastAPI, Response @@ -304,6 +304,10 @@ def get_api_server( master_addr: Optional[str], master_port: Optional[int], logger: Logger, + enable_recording: bool = False, + recorder: Optional[Any] = None, + record_store: Optional[Any] = None, + routed_experts_layout: Optional[Tuple[int, int]] = None, ) -> "asyncio.Task[None]": _apply_openai_api_monkey_patch() @@ -329,7 +333,14 @@ def get_api_server( enable_return_routed_experts=enable_return_routed_experts, skip_server_warmup=True, disable_piecewise_cuda_graph=True, - api_key=api_key, + # When recording is on, disable SGLang's api_key auth so the + # Authorization bearer is used purely as the per-task record_key (read + # by RecordingIdentityMiddleware). Trinity's record_key is per-task + # ("batch_id/task_id/run_index") and differs from the single configured + # api_key, so the auth middleware would otherwise 401-reject it. This + # mirrors vLLM, whose recording server sets no api_key auth. The + # embedded server is localhost/in-Ray-actor, so auth is not needed. + api_key=None if enable_recording else api_key, nnodes=nnodes, node_rank=node_rank, dist_init_addr=( @@ -365,6 +376,22 @@ def get_api_server( logger=logger, ) + # Wire generation recording before the uvicorn task starts serving. The + # recorder/store are owned by ``SGLangRolloutModel``; this installs the + # engine wrap on ``tokenizer_manager``, ``RecordingIdentityMiddleware`` and + # ``query_router`` on ``app``, and stashes store/recorder on ``app.state``. + if enable_recording: + from trinity.common.models.sglang_patch.recording import setup_sglang_recording + + setup_sglang_recording( + tokenizer_manager, + app, + logger, + recorder=recorder, + store=record_store, + routed_experts_layout=routed_experts_layout, + ) + config = uvicorn.Config( app, host=server_args.host, From a274ed32e6701c4597ac09d63462c6f693b54393 Mon Sep 17 00:00:00 2001 From: pxc Date: Thu, 25 Jun 2026 21:10:25 +0800 Subject: [PATCH 18/84] add sglang tests --- tests/common/sglang_test.py | 286 ++++++++++++++++++++++++++ trinity/common/models/model.py | 18 +- trinity/common/models/sglang_model.py | 21 +- trinity/common/models/vllm_model.py | 12 -- 4 files changed, 306 insertions(+), 31 deletions(-) diff --git a/tests/common/sglang_test.py b/tests/common/sglang_test.py index 6c4254b371b..277c0a04dae 100644 --- a/tests/common/sglang_test.py +++ b/tests/common/sglang_test.py @@ -1,5 +1,7 @@ import asyncio +import httpx +import openai import torch from parameterized import parameterized_class from transformers import AutoConfig, AutoTokenizer @@ -11,6 +13,7 @@ get_moe_model_path, get_template_config, ) +from trinity.common.experience import Experience from trinity.common.models.allocator import Allocator @@ -361,3 +364,286 @@ async def test_chat_completions(self): ) else: self.assertEqual(len(self.model_wrapper.history), 0) + + +class TestRecording(RayUnittestBaseAsync): + """Correctness of the in-SGLang generation recording flow (``enable_recording``). + + Mirrors ``tests/common/vllm_test.py::TestRecording``. Verifies that every + call path lands its finished turn in the in-process ``MemoryStore`` under + the right ``record_key``, and that ``POST /records/update_record`` flushes + the recorder, reward-stamps the whole record-key group, pops it, and returns + it as serialized experiences. + + Paths covered (all async): + * Ray-direct ``generate`` / ``chat`` — SGLang's Ray-direct path is over + HTTP (unlike vLLM's in-process call), so ``record_key`` travels as the + ``Authorization: Bearer `` header. + * OpenAI HTTP regular / streaming / tool-augmented — same bearer path. + + Recording disables SGLang's api_key auth middleware (Option A, see + ``sglang_patch/server_patch.py``), so the bearer is used purely as the + per-task ``record_key`` (captured by ``RecordingIdentityMiddleware``), + matching vLLM (which sets no api_key auth in recording mode). + + ``enable_recording`` forces ``enable_return_routed_experts`` in the + Allocator, so this test uses a MoE checkpoint (``get_moe_model_path``). + """ + + async def asyncSetUp(self): + self.config = get_template_config() + self.config.mode = "explore" + # enable_recording forces enable_return_routed_experts -> needs a MoE + # model (otherwise routed_experts is absent and the shape asserts below + # would fail). Use a Qwen3-MoE checkpoint. + self.config.model.model_path = get_moe_model_path() + self.tokenizer = AutoTokenizer.from_pretrained( + self.config.model.model_path, + trust_remote_code=True, + ) + self.text_config = _get_text_config(self.config.model.model_path) + self.expected_routed_experts_layers = int(self.text_config.num_hidden_layers) + self.expected_routed_experts_topk = int(self.text_config.num_experts_per_tok) + self.config.model.custom_chat_template = CHAT_TEMPLATE + self.config.explorer.rollout_model.engine_type = "sglang" + self.config.explorer.rollout_model.engine_num = 1 + self.config.explorer.rollout_model.tensor_parallel_size = 2 + self.config.explorer.rollout_model.chat_template = CHAT_TEMPLATE + # enable_recording requires the OpenAI API server (the recording runner). + self.config.explorer.rollout_model.enable_openai_api = True + self.config.explorer.rollout_model.enable_recording = True + self.config.explorer.rollout_model.enable_expert_parallel = True + # History recording is client-side; the in-SGLang recorder is the subject. + self.config.explorer.rollout_model.enable_history = False + self.config.explorer.rollout_model.base_port = 13400 + self.config.check_and_update() + + allocator = Allocator(self.config.explorer) + rollout_models, _ = await allocator.create_all_models() + self.model_wrapper = rollout_models[0] + self.api_address = self.model_wrapper.api_address + self._http = httpx.AsyncClient(timeout=120.0) + self._model_id = None + + async def asyncTearDown(self): + await self._http.aclose() + await self.model_wrapper.shutdown() + await super().asyncTearDown() + + # -- /records store query/consume helpers --------------------------------- + + async def _flush(self): + """Drain the recorder without popping anything (an empty consume).""" + resp = await self._http.post( + f"{self.api_address}/records/update_record", json={"updates": []} + ) + resp.raise_for_status() + + async def _list_record_keys(self): + resp = await self._http.get(f"{self.api_address}/records") + resp.raise_for_status() + return resp.json()["record_keys"] + + async def _get_record_experiences(self, record_key: str) -> dict: + resp = await self._http.get(f"{self.api_address}/records/{record_key}") + resp.raise_for_status() + return resp.json() + + async def _get_request_experience(self, record_key: str, request_id: str) -> Experience: + resp = await self._http.get(f"{self.api_address}/records/{record_key}/request/{request_id}") + resp.raise_for_status() + return Experience.deserialize(resp.content) + + async def _consume( + self, record_key: str, reward: float, run: int, task: str + ) -> list[Experience]: + resp = await self._http.post( + f"{self.api_address}/records/update_record", + json={ + "updates": [{"record_key": record_key, "reward": reward, "run": run, "task": task}] + }, + ) + resp.raise_for_status() + return Experience.deserialize_many(resp.content) + + async def _openai_client(self, record_key: str) -> openai.AsyncOpenAI: + # record_key travels as the Bearer api_key -> RecordingIdentityMiddleware. + return openai.AsyncOpenAI(base_url=f"{self.api_address}/v1", api_key=record_key) + + async def _get_model_id(self, client: openai.AsyncOpenAI) -> str: + if self._model_id is None: + self._model_id = (await client.models.list()).data[0].id + return self._model_id # type: ignore [return-value] + + # -- per-recorded-experience invariants ----------------------------------- + + def _assert_recorded_experience(self, exp: Experience, record_key: str): + self.assertEqual(exp.info.get("record_key"), record_key) + self.assertIsNotNone(exp.info.get("request_id")) + self.assertEqual(exp.info.get("rank"), 0) + # SGLang stamps meta_info.weight_version ("default" until a weight sync); + # unlike vLLM it is a server-tracked string, not the model_version int. + self.assertIsNotNone(exp.info.get("model_version")) + self.assertGreater(len(exp.tokens), exp.prompt_length) # type: ignore [arg-type] + # The recorder forces return_logprob=True even when the client omitted it. + self.assertGreater(len(exp.logprobs), 0) # type: ignore [arg-type] + self.assertEqual(len(exp.logprobs), len(exp.tokens) - exp.prompt_length) # type: ignore [arg-type] + # SGLang's ret does not carry prompt text, so prompt_text is None on the + # recording hot path (decode token ids lazily where a check is needed). + if exp.prompt_text is not None: + self.assertGreater(len(exp.prompt_text), 0) + self.assertIsNotNone(exp.response_text) + + def _assert_recorded_routed_experts(self, exp: Experience): + # enable_return_routed_experts is forced on by enable_recording. + self.assertIsNotNone(exp.routed_experts) + re = exp.routed_experts + self.assertEqual(re.dtype, torch.uint8) + self.assertEqual(re.ndim, 3) + self.assertEqual(re.shape[1], self.expected_routed_experts_layers) + self.assertEqual(re.shape[2], self.expected_routed_experts_topk) + + async def test_record(self): # noqa: C901 + messages = [ + {"role": "system", "content": "You are a helpful assistant."}, + {"role": "user", "content": "Say hello in one short sentence."}, + ] + + # ===== 1. Ray-direct generate (record_key via Authorization bearer) ===== + rk_gen = "trinity_record_generate" + await self.model_wrapper.generate_async( + ["Hello, world!"], n=1, temperature=1.0, max_tokens=16, record_key=rk_gen + ) + await self._flush() + self.assertIn(rk_gen, await self._list_record_keys()) + task = await self._get_record_experiences(rk_gen) + self.assertEqual(len(task["experiences"]), 1) + # blob endpoint round-trips a full experience + request_id = task["experiences"][0]["info"]["request_id"] + blob_exp = await self._get_request_experience(rk_gen, request_id) + self._assert_recorded_experience(blob_exp, rk_gen) + self._assert_recorded_routed_experts(blob_exp) + consumed = await self._consume(rk_gen, reward=0.5, run=1, task="t_gen") + self.assertEqual(len(consumed), 1) + self.assertEqual(consumed[0].reward, 0.5) + self.assertEqual(consumed[0].eid.run, 1) + self.assertEqual(consumed[0].eid.task, "t_gen") + self._assert_recorded_experience(consumed[0], rk_gen) + self._assert_recorded_routed_experts(consumed[0]) + self.assertNotIn(rk_gen, await self._list_record_keys()) # popped + + # ===== 2. Ray-direct chat, n=2 (one record-key group, two samples) ===== + rk_chat = "trinity_record_chat" + chat_exps = await self.model_wrapper.chat_async( + messages, n=2, temperature=1.0, max_tokens=16, record_key=rk_chat + ) + self.assertEqual(len(chat_exps), 2) + await self._flush() + task = await self._get_record_experiences(rk_chat) + self.assertEqual(len(task["experiences"]), 2) + # SGLang expands n=2 parallel sampling into two scheduler requests, + # each with its own rid -> two distinct request_ids. The list position + # becomes sample_index (0, 1) to order the two samples within the + # record-key group (the GRPO-style grouping the store is built for). + # (vLLM differs: n=2 shares one request_id with sample_index 0/1.) + self.assertEqual(sorted(t["info"]["sample_index"] for t in task["experiences"]), [0, 1]) + self.assertEqual(len({t["info"]["request_id"] for t in task["experiences"]}), 2) + consumed = await self._consume(rk_chat, reward=0.8, run=2, task="t_chat") + self.assertEqual(len(consumed), 2) + for exp in consumed: + self.assertEqual(exp.reward, 0.8) + self.assertEqual(exp.eid.run, 2) + self.assertEqual(exp.eid.task, "t_chat") + self._assert_recorded_experience(exp, rk_chat) + self._assert_recorded_routed_experts(exp) + self.assertNotIn(rk_chat, await self._list_record_keys()) + + # ===== 3. OpenAI regular (HTTP; record_key = Bearer api_key) ===== + rk_oai = "trinity_record_openai" + client = await self._openai_client(rk_oai) + model_id = await self._get_model_id(client) + resp = await client.chat.completions.create( + model=model_id, + messages=messages, + n=1, + temperature=0.7, + max_tokens=32, + ) + consumed = await self._consume(rk_oai, reward=0.3, run=3, task="t_oai") + self.assertEqual(len(consumed), 1) + self._assert_recorded_experience(consumed[0], rk_oai) + self._assert_recorded_routed_experts(consumed[0]) + # No reasoning_parser is configured, so message.content == ret.text. + self.assertEqual(consumed[0].response_text, resp.choices[0].message.content) + self.assertNotIn(rk_oai, await self._list_record_keys()) + + # ===== 4. OpenAI streaming (HTTP) ===== + rk_str = "trinity_record_stream" + sclient = await self._openai_client(rk_str) + stream = await sclient.chat.completions.create( + model=model_id, + messages=messages, + n=1, + stream=True, + temperature=0.7, + max_tokens=32, + ) + content = "" + async for chunk in stream: + delta = chunk.choices[0].delta.content + if delta: + content += delta + self.assertGreater(len(content), 0) + consumed = await self._consume(rk_str, reward=0.1, run=4, task="t_str") + self.assertEqual(len(consumed), 1) + self._assert_recorded_experience(consumed[0], rk_str) + self._assert_recorded_routed_experts(consumed[0]) + response_token_ids = consumed[0].tokens[consumed[0].prompt_length :].tolist() + decoded_content = self.tokenizer.decode(response_token_ids, skip_special_tokens=True) + self.assertEqual(decoded_content, content) + self.assertNotIn(rk_str, await self._list_record_keys()) + + # ===== 5. OpenAI tool-augmented (HTTP) ===== + rk_tool = "trinity_record_tool" + tclient = await self._openai_client(rk_tool) + tools = [ + { + "type": "function", + "function": { + "name": "get_current_weather", + "description": "Get the current weather in a given location", + "parameters": { + "type": "object", + "properties": { + "location": { + "description": "The city and state, e.g. San Francisco, CA", + "type": "string", + } + }, + "required": ["location"], + }, + }, + } + ] + tool_messages = [{"role": "user", "content": "What's the weather like in Boston?"}] + tresp = await tclient.chat.completions.create( + model=model_id, + messages=tool_messages, + tools=tools, + tool_choice="none", + max_tokens=32, + ) + consumed = await self._consume(rk_tool, reward=1.0, run=5, task="t_tool") + self.assertEqual(len(consumed), 1) + self._assert_recorded_experience(consumed[0], rk_tool) + self._assert_recorded_routed_experts(consumed[0]) + # The tool-augmented prompt (tool defs rendered by the chat template) + # must be part of the recorded experience. SGLang's ret does not carry + # prompt text, so decode the recorded tokens to check. + decoded = self.tokenizer.decode(consumed[0].tokens.tolist(), skip_special_tokens=False) + self.assertIn("get_current_weather", decoded) + self.assertEqual(consumed[0].response_text, tresp.choices[0].message.content) + self.assertNotIn(rk_tool, await self._list_record_keys()) + + # ===== global: every group consumed -> store is drained ===== + self.assertEqual(await self._list_record_keys(), []) diff --git a/trinity/common/models/model.py b/trinity/common/models/model.py index 38d76349b48..10a5bcfbbcb 100644 --- a/trinity/common/models/model.py +++ b/trinity/common/models/model.py @@ -143,8 +143,22 @@ def get_api_key(self) -> str: async def extract_experience_from_history( self, record_key: str, clear_history: bool = True ) -> List[Experience]: - """Extract recorded experiences by record key.""" - raise NotImplementedError("Recording extraction is not supported by this model.") + """Extract recorded experiences by record key from the in-process store. + + Both vLLM and SGLang keep the recorder and its store in-process (the + engine / embedded HTTP server runs in the same event loop as the model), + so extraction is a direct store lookup with no HTTP hop. Subclasses that + enable recording must set ``self.recorder`` (a ``Recorder`` whose + ``.store`` is a ``RecordStore``); this base implementation is shared. + """ + recorder = getattr(self, "recorder", None) + if recorder is None: + raise ValueError("Recording is not enabled for this model.") + await recorder.flush() + exps = await recorder.store.get_record_experiences(record_key) + if clear_history: + await recorder.store.delete_record_experiences(record_key) + return exps def get_model_config(self) -> InferenceModelConfig: """Get the model configuration.""" diff --git a/trinity/common/models/sglang_model.py b/trinity/common/models/sglang_model.py index d5a598daefa..efe9dd101d3 100644 --- a/trinity/common/models/sglang_model.py +++ b/trinity/common/models/sglang_model.py @@ -500,23 +500,10 @@ async def convert_messages_to_experience( "SGLangRolloutModel does not support convert_messages_to_experience." ) - async def extract_experience_from_history( - self, record_key: str, clear_history: bool = True - ) -> List[Experience]: - """Extract recorded experiences from the in-process recorder store. - - Mirrors ``vLLMRolloutModel.extract_experience_from_history``: the - recorder and store live in-process (the SGLang server runs as an - asyncio task in the same loop as this model), so extraction is a direct - store lookup with no HTTP hop. - """ - if self.recorder is None: - raise ValueError("Recording is not enabled for this SGLang model.") - await self.recorder.flush() - exps = await self.recorder.store.get_record_experiences(record_key) - if clear_history: - await self.recorder.store.delete_record_experiences(record_key) - return exps + # ``extract_experience_from_history`` is implemented on the shared + # ``InferenceModel`` base; ``self.recorder`` is installed by ``run_api_server`` + # when recording is on (the recorder/store live in-process with the embedded + # SGLang server, same as vLLM). def _get_api_server_exit_reason(self) -> Optional[str]: if self.api_server is None or not self.api_server.done(): diff --git a/trinity/common/models/vllm_model.py b/trinity/common/models/vllm_model.py index d57a2f0ca16..539042df0fd 100644 --- a/trinity/common/models/vllm_model.py +++ b/trinity/common/models/vllm_model.py @@ -372,18 +372,6 @@ async def logprobs( # type: ignore [override] dtype=torch.float32, ) - async def extract_experience_from_history( - self, record_key: str, clear_history: bool = True - ) -> List[Experience]: - """Extract recorded experiences from the in-vLLM recorder store.""" - if self.recorder is None: - raise ValueError("Recording is not enabled for this vLLM model.") - await self.recorder.flush() - exps = await self.recorder.store.get_record_experiences(record_key) - if clear_history: - await self.recorder.store.delete_record_experiences(record_key) - return exps - async def add_lora_adapter(self, lora_request: Any) -> int: """Add a LoRA adapter to the vLLM engine. From 35952ae3bb93903ac97fb9f8c1330b3d06023605 Mon Sep 17 00:00:00 2001 From: pxc Date: Fri, 26 Jun 2026 10:29:02 +0800 Subject: [PATCH 19/84] fix sglang tests --- tests/common/sglang_test.py | 29 ++++++++++++++----- trinity/common/models/sglang_model.py | 1 + .../models/sglang_patch/recording/recorder.py | 21 ++++++++++++-- .../models/sglang_patch/server_patch.py | 5 ++++ 4 files changed, 46 insertions(+), 10 deletions(-) diff --git a/tests/common/sglang_test.py b/tests/common/sglang_test.py index 277c0a04dae..19bf61e431b 100644 --- a/tests/common/sglang_test.py +++ b/tests/common/sglang_test.py @@ -413,6 +413,13 @@ async def asyncSetUp(self): self.config.explorer.rollout_model.enable_openai_api = True self.config.explorer.rollout_model.enable_recording = True self.config.explorer.rollout_model.enable_expert_parallel = True + # Tool-call parsing coverage (qwen3_coder matches the Qwen3.5 chat + # template). SGLang enables tool calling via tool_call_parser (no + # separate enable_auto_tool_choice flag); enable_auto_tool_choice is + # set for parity with the vLLM TestRecording config. + self.config.explorer.rollout_model.enable_auto_tool_choice = True + self.config.explorer.rollout_model.tool_call_parser = "qwen3_coder" + self.config.explorer.rollout_model.enable_thinking = False # History recording is client-side; the in-SGLang recorder is the subject. self.config.explorer.rollout_model.enable_history = False self.config.explorer.rollout_model.base_port = 13400 @@ -603,7 +610,7 @@ async def test_record(self): # noqa: C901 self.assertEqual(decoded_content, content) self.assertNotIn(rk_str, await self._list_record_keys()) - # ===== 5. OpenAI tool-augmented (HTTP) ===== + # ===== 5. OpenAI tool-call parsing (HTTP) ===== rk_tool = "trinity_record_tool" tclient = await self._openai_client(rk_tool) tools = [ @@ -626,23 +633,31 @@ async def test_record(self): # noqa: C901 } ] tool_messages = [{"role": "user", "content": "What's the weather like in Boston?"}] + no_think = {"chat_template_kwargs": {"enable_thinking": False}} tresp = await tclient.chat.completions.create( model=model_id, messages=tool_messages, tools=tools, - tool_choice="none", - max_tokens=32, + tool_choice="auto", + max_tokens=64, + extra_body=no_think, ) consumed = await self._consume(rk_tool, reward=1.0, run=5, task="t_tool") self.assertEqual(len(consumed), 1) self._assert_recorded_experience(consumed[0], rk_tool) self._assert_recorded_routed_experts(consumed[0]) - # The tool-augmented prompt (tool defs rendered by the chat template) - # must be part of the recorded experience. SGLang's ret does not carry - # prompt text, so decode the recorded tokens to check. + # tool_choice != "none" -> SGLang renders the tool defs into the prompt + # (serving_chat._process_messages), so the recorded prompt tokens carry + # the tool name. SGLang's ret does not carry prompt text, so decode. decoded = self.tokenizer.decode(consumed[0].tokens.tolist(), skip_special_tokens=False) self.assertIn("get_current_weather", decoded) - self.assertEqual(consumed[0].response_text, tresp.choices[0].message.content) + # If the model emitted a tool call, its function name is in the raw + # recorded response text (ret.text), which the qwen3_coder parser also + # surfaces as choice.message.tool_calls. + choice = tresp.choices[0] + if choice.finish_reason == "tool_calls" and choice.message.tool_calls: + for tc in choice.message.tool_calls: + self.assertIn(tc.function.name, consumed[0].response_text) self.assertNotIn(rk_tool, await self._list_record_keys()) # ===== global: every group consumed -> store is drained ===== diff --git a/trinity/common/models/sglang_model.py b/trinity/common/models/sglang_model.py index efe9dd101d3..96881a58a64 100644 --- a/trinity/common/models/sglang_model.py +++ b/trinity/common/models/sglang_model.py @@ -593,6 +593,7 @@ async def run_api_server(self) -> bool: recorder=recorder, record_store=record_store, routed_experts_layout=routed_experts_layout, + tool_call_parser=self.config.tool_call_parser, logger=self.logger, ) # ``setup_sglang_recording`` (called inside get_api_server) owns the diff --git a/trinity/common/models/sglang_patch/recording/recorder.py b/trinity/common/models/sglang_patch/recording/recorder.py index 9f2a760d06f..be93485775e 100644 --- a/trinity/common/models/sglang_patch/recording/recorder.py +++ b/trinity/common/models/sglang_patch/recording/recorder.py @@ -52,8 +52,15 @@ def _get_obj(args, kwargs): return None -def _force_record_fields(obj: Any) -> None: - """Force logprob + prompt-token-id capture for recording (transparent).""" +def _force_record_fields(obj: Any, *, force_routed_experts: bool) -> None: + """Force logprob + prompt-token-id (+ routed-expert) capture for recording. + + Transparent to clients: the OpenAI serving layer gates its response + ``logprobs`` / ``prompt_token_ids`` / ``sglext.routed_experts`` emission on + the *ChatCompletionRequest* flags (unchanged); we only flip the + ``GenerateReqInput`` flags the scheduler reads, so the recorded ``ret`` gains + these fields while HTTP responses stay the same. + """ if obj is None: return # return_logprob may be a list for batched requests; broadcast True. @@ -65,6 +72,14 @@ def _force_record_fields(obj: Any) -> None: obj.return_logprob = True if hasattr(obj, "return_prompt_token_ids"): obj.return_prompt_token_ids = True + # The scheduler only returns routed_experts when the per-request flag is set + # (scheduler.py: ``if recv_req.return_routed_experts``), even though the + # model runner computes them whenever the server flag is on. The chat path + # defaults this to False, so force it here when the server is MoE-enabled + # (signaled by a non-None routed_experts_layout) so the recorded experience + # carries routed_experts on every path, not just Ray-direct /generate. + if force_routed_experts and hasattr(obj, "return_routed_experts"): + obj.return_routed_experts = True def _normalize_ret(out: Any) -> List[dict]: @@ -243,7 +258,7 @@ def patch_tokenizer_manager_for_recording( async def _patched_generate_request(*args, **kwargs): obj = _get_obj(args, kwargs) if recorder.enabled and obj is not None: - _force_record_fields(obj) + _force_record_fields(obj, force_routed_experts=routed_experts_layout is not None) state: dict = {} order: list = [] diff --git a/trinity/common/models/sglang_patch/server_patch.py b/trinity/common/models/sglang_patch/server_patch.py index 4c14ba332a2..5233204adc8 100644 --- a/trinity/common/models/sglang_patch/server_patch.py +++ b/trinity/common/models/sglang_patch/server_patch.py @@ -308,6 +308,7 @@ def get_api_server( recorder: Optional[Any] = None, record_store: Optional[Any] = None, routed_experts_layout: Optional[Tuple[int, int]] = None, + tool_call_parser: Optional[str] = None, ) -> "asyncio.Task[None]": _apply_openai_api_monkey_patch() @@ -341,6 +342,10 @@ def get_api_server( # mirrors vLLM, whose recording server sets no api_key auth. The # embedded server is localhost/in-Ray-actor, so auth is not needed. api_key=None if enable_recording else api_key, + # SGLang enables tool calling via tool_call_parser (no separate + # enable_auto_tool_choice flag in this version). Only render/parse tools + # when a parser is configured, matching vLLM's enable_auto_tool_choice. + tool_call_parser=tool_call_parser, nnodes=nnodes, node_rank=node_rank, dist_init_addr=( From 6d6ce69cc226eeacc19cb47c78dc24930cefa31b Mon Sep 17 00:00:00 2001 From: "panxuchen.pxc" Date: Fri, 26 Jun 2026 10:30:22 +0800 Subject: [PATCH 20/84] remove enable recording --- tests/common/vllm_test.py | 10 +-- trinity/common/config.py | 29 +++---- trinity/common/config_validator.py | 4 + trinity/common/models/allocator.py | 15 ++-- trinity/common/models/model.py | 84 ++++++------------- trinity/common/models/vllm_model.py | 2 +- trinity/common/models/vllm_patch/__init__.py | 6 +- .../common/workflows/step_wise_workflow.py | 14 ++-- trinity/common/workflows/workflow.py | 2 +- trinity/explorer/rollout_coordinator.py | 2 +- trinity/explorer/workflow_runner.py | 2 +- 11 files changed, 68 insertions(+), 102 deletions(-) diff --git a/tests/common/vllm_test.py b/tests/common/vllm_test.py index ac822e20d93..1ce6a1cd014 100644 --- a/tests/common/vllm_test.py +++ b/tests/common/vllm_test.py @@ -1812,7 +1812,7 @@ async def test_api_tool_calls(self): class TestRecording(VLLMTestBase): - """Correctness of the in-vLLM generation recording flow (``enable_recording``). + """Correctness of the in-vLLM generation recording flow (``enable_history``). Verifies that every call path lands its finished turn in the in-process ``MemoryStore`` under the right ``record_key``, and that @@ -1826,7 +1826,7 @@ class TestRecording(VLLMTestBase): via the ``Authorization: Bearer `` header, captured by ``RecordingIdentityMiddleware``. - ``enable_recording`` forces ``enable_return_routed_experts`` in the + ``enable_history`` forces ``enable_return_routed_experts`` in the Allocator, and vLLM's routed-experts capturer raises on a non-MoE model, so this test requires a MoE checkpoint (``TRINITY_MOE_MODEL_PATH``). """ @@ -1836,7 +1836,7 @@ async def asyncSetUp(self): self.skipTest("generation recording requires vLLM >= 0.23.0") self.config = get_template_config() self.config.mode = "explore" - # enable_recording forces enable_return_routed_experts -> needs a MoE + # enable_history forces enable_return_routed_experts -> needs a MoE # model (vLLM raises on dense models). Use a Qwen3-MoE checkpoint. self.config.model.model_path = get_moe_model_path() self.tokenizer = AutoTokenizer.from_pretrained( @@ -1851,9 +1851,9 @@ async def asyncSetUp(self): self.config.explorer.rollout_model.engine_num = 1 self.config.explorer.rollout_model.tensor_parallel_size = 2 self.config.explorer.rollout_model.chat_template = CHAT_TEMPLATE - # enable_recording requires the OpenAI API server (the recording runner). + # enable_history requires the OpenAI API server (the recording runner). self.config.explorer.rollout_model.enable_openai_api = True - self.config.explorer.rollout_model.enable_recording = True + self.config.explorer.rollout_model.enable_history = True self.config.explorer.rollout_model.enable_expert_parallel = True # Tool-call coverage; qwen3_coder matches the Qwen3.5 chat template. self.config.explorer.rollout_model.enable_auto_tool_choice = True diff --git a/trinity/common/config.py b/trinity/common/config.py index e9a6a6d6f73..92f89cf4bf9 100644 --- a/trinity/common/config.py +++ b/trinity/common/config.py @@ -567,7 +567,19 @@ class InferenceModelConfig: # For Qwen3 enable_thinking: Optional[bool] = None - # For history recording + # For experience recording. When enabled on vLLM rollout models, the engine + # wraps ``engine_client.generate`` and writes each finished turn as a + # Trinity ``Experience`` to the in-process ``MemoryStore``, keyed by the + # recording identity (``record_key``). This is the single switch for the + # recording flow — when on, the explorer also consumes from the store: the + # WorkflowRunner ships only a small reward map keyed by ``record_key`` and + # the RolloutCoordinator pulls heavy experiences via ``/records/update_record`` + # at finalize time. When off (default), runners ship serialized experiences + # through the scheduler as before (legacy path). When True, the Allocator + # forces ``enable_return_routed_experts``. VLLMModel mirrors the recording + # config onto the engine instance for the recorder to read. The capture + # width (top-k logprobs) reuses ``logprobs`` below (default 1). Requires + # ``enable_openai_api=True`` (the recording runner is the API server). enable_history: bool = False # For OpenAI API @@ -591,21 +603,6 @@ class InferenceModelConfig: node_rank: int = 0 enable_return_routed_experts: bool = False - # Turn on in-vLLM generation recording for the OpenAI API serving path: the - # engine wraps ``engine_client.generate`` and writes each finished turn as a - # Trinity ``Experience`` to the in-process ``MemoryStore``, keyed by the - # recording identity (``record_key``). This is the single switch for the - # recording flow — when on, the explorer also consumes from the store: the - # WorkflowRunner ships only a small reward map keyed by ``record_key`` and - # the RolloutCoordinator pulls heavy experiences via ``/records/update_record`` - # at finalize time. When off (default), runners ship serialized experiences - # through the scheduler as before (legacy path). When True, the Allocator - # forces ``enable_return_routed_experts``. VLLMModel mirrors the recording - # config onto the engine instance for the recorder to read. The capture - # width (top-k logprobs) reuses ``logprobs`` below (default 1). Requires - # ``enable_openai_api=True`` (the recording runner is the API server). - enable_recording: bool = False - # Buffer size (MB) for batched NCCL weight sync. Controls peak GPU memory during sync. weight_sync_buffer_size: int = 1024 # MB diff --git a/trinity/common/config_validator.py b/trinity/common/config_validator.py index 6bb1bfd6c7e..f6b7e909063 100644 --- a/trinity/common/config_validator.py +++ b/trinity/common/config_validator.py @@ -705,6 +705,10 @@ def validate(self, config: Config) -> None: if config.mode == "serve": # in 'serve' mode, we always enable openai api for rollout model config.explorer.rollout_model.enable_openai_api = True + if config.explorer.rollout_model.engine_type.startswith("vllm"): + # enable_history is the single switch for vLLM recording. + if config.explorer.rollout_model.enable_history: + config.explorer.rollout_model.enable_openai_api = True self._validate_inference_parallel_config(config.explorer.rollout_model, "rollout_model") # auxiliary models for aux_model in config.explorer.auxiliary_models: diff --git a/trinity/common/models/allocator.py b/trinity/common/models/allocator.py index bdc202cb24b..f4eae293a35 100644 --- a/trinity/common/models/allocator.py +++ b/trinity/common/models/allocator.py @@ -82,12 +82,15 @@ async def create_engine( config = deepcopy(config) config.engine_id = engine_id - # In-vLLM recording: force routed_experts capture (the engine reads - # this at build time — see vllm_model.py). VLLMModel mirrors the - # recording config onto the engine instance for the recorder to read. - # No env / runtime_env involved. - if config.enable_recording: - config.enable_return_routed_experts = True + if config.engine_type.startswith("vllm"): + # enable_history is the single switch for vLLM recording. + if config.enable_history: + config.enable_openai_api = True + # In-vLLM recording: force routed_experts capture (the engine reads + # this at build time — see vllm_model.py). VLLMModel mirrors the + # recording config onto the engine instance for the recorder to read. + # No env / runtime_env involved. + config.enable_return_routed_experts = True actor_bundle_lists = [] for node_id in range(config.nnodes): diff --git a/trinity/common/models/model.py b/trinity/common/models/model.py index 38d76349b48..805aa7e8b1a 100644 --- a/trinity/common/models/model.py +++ b/trinity/common/models/model.py @@ -366,24 +366,6 @@ async def convert_messages_to_experience( ) -def _history_recorder(func): - """Decorator to record history of the model calls.""" - - async def async_wrapper(self, *args, **kwargs): - result = await func(self, *args, **kwargs) - if self.enable_history: - self._record_history(result) - return result - - def sync_wrapper(self, *args, **kwargs): - result = func(self, *args, **kwargs) - if self.enable_history: - self._record_history(result) - return result - - return async_wrapper if asyncio.iscoroutinefunction(func) else sync_wrapper - - class ModelWrapper: """A wrapper for the InferenceModel Ray Actor""" @@ -436,7 +418,6 @@ def __init__( self.logger = get_logger(__name__) self.enable_lora = config.enable_lora self.enable_history = config.enable_history - self.history = [] self.recording_history_offsets: Dict[str, int] = {} self.status = RunningStatus.RUNNING self.workflow_state: Dict = {} @@ -478,44 +459,32 @@ async def prepare(self) -> None: f"API server at {self.api_address} not ready after {max_retries} attempts." ) - def _record_history(self, exps: Union[Experience, List[Experience]]) -> None: - """Record experiences to history.""" - if isinstance(exps, Experience): - self.history.append(exps) - elif isinstance(exps, list): - self.history.extend(exps) - else: - raise TypeError("Expected Experience or List[Experience], got {}".format(type(exps))) - - @_history_recorder def generate(self, prompts: List[str], **kwargs) -> List[Experience]: """Generate a list of experiences from a list of prompts.""" lora_request = self.get_lora_request() - if self.config.enable_recording and kwargs.get("record_key") is None: + if self.config.enable_history and kwargs.get("record_key") is None: kwargs["record_key"] = self._api_key results = ray.get( [self.model.generate.remote(prompt, lora_request, **kwargs) for prompt in prompts] ) return [exp for exps in results for exp in exps] - @_history_recorder async def generate_async(self, prompts: List[str], **kwargs) -> List[Experience]: """Generate a list of experiences from a list of prompts in async.""" lora_request = await self.get_lora_request_async() - if self.config.enable_recording and kwargs.get("record_key") is None: + if self.config.enable_history and kwargs.get("record_key") is None: kwargs["record_key"] = self._api_key results = await asyncio.gather( *[self.model.generate.remote(prompt, lora_request, **kwargs) for prompt in prompts] ) return [exp for exps in results for exp in exps] - @_history_recorder def chat( self, messages: List[dict], record_key: Optional[str] = None, **kwargs ) -> List[Experience]: """Generate a list of experiences from a list of messages.""" lora_request = self.get_lora_request() - if self.config.enable_recording and record_key is None: + if self.config.enable_history and record_key is None: record_key = self._api_key return ray.get( self.model.chat.remote( @@ -523,13 +492,12 @@ def chat( ) ) - @_history_recorder async def chat_async( self, messages: List[dict], record_key: Optional[str] = None, **kwargs ) -> List[Experience]: """Generate a list of experiences from a list of messages in async.""" lora_request = await self.get_lora_request_async() - if self.config.enable_recording and record_key is None: + if self.config.enable_history and record_key is None: record_key = self._api_key return await self.model.chat.remote( messages, lora_request=lora_request, record_key=record_key, **kwargs @@ -671,6 +639,7 @@ def chat_completions(*args, **kwargs): ) ) response = chat_response.pop() + # TODO: Tinker legacy recording path - should be migrated to engine-side recording if self.enable_history: self.history.extend(chat_response) return response @@ -716,6 +685,7 @@ async def chat_completions(*args, **kwargs): **kwargs, ) response = chat_response.pop() + # TODO: Tinker legacy recording path - should be migrated to engine-side recording if self.enable_history: self.history.extend(chat_response) return response @@ -781,31 +751,26 @@ def extract_experience_from_history( self, clear_history: bool = True, record_key: Optional[str] = None ) -> List[Experience]: """Extract experiences from the history.""" - if self.config.enable_recording: - if self.model is None: - raise ValueError("Recording extraction requires an inference model actor.") - record_key = record_key or self._api_key - if record_key is None: - raise ValueError("record_key is required when recording is enabled.") - exps = ray.get( - self.model.extract_experience_from_history.remote( - record_key=record_key, - clear_history=False, - ) - ) - offset = self.recording_history_offsets.get(record_key, 0) - if offset > len(exps): - offset = 0 - new_exps = exps[offset:] - if clear_history: - self.recording_history_offsets[record_key] = len(exps) - return new_exps if not self.enable_history: raise ValueError("History recording is not enabled.") - exps = [exp for exp in self.history] + if self.model is None: + raise ValueError("Recording extraction requires an inference model actor.") + record_key = record_key or self._api_key + if record_key is None: + raise ValueError("record_key is required when recording is enabled.") + exps = ray.get( + self.model.extract_experience_from_history.remote( + record_key=record_key, + clear_history=False, + ) + ) + offset = self.recording_history_offsets.get(record_key, 0) + if offset > len(exps): + offset = 0 + new_exps = exps[offset:] if clear_history: - self.history.clear() - return exps + self.recording_history_offsets[record_key] = len(exps) + return new_exps # Workflow state management methods async def set_workflow_state(self, state: Dict) -> None: @@ -817,7 +782,7 @@ async def clean_workflow_state(self) -> None: """Clean the state of workflow using the model.""" async with self.state_lock: self.workflow_state = {} - self.history.clear() + # TODO: clear engine side experiences self.recording_history_offsets.clear() async def shutdown(self) -> None: @@ -839,6 +804,5 @@ def clone_with_isolated_history(self) -> "ModelWrapper": new_wrapper = copy.copy(self) new_wrapper.openai_async_client = None new_wrapper.openai_client = None - new_wrapper.history = [] new_wrapper.recording_history_offsets = {} return new_wrapper diff --git a/trinity/common/models/vllm_model.py b/trinity/common/models/vllm_model.py index d57a2f0ca16..b5c15ec2056 100644 --- a/trinity/common/models/vllm_model.py +++ b/trinity/common/models/vllm_model.py @@ -191,7 +191,7 @@ async def prepare(self) -> None: # attribute experiences to the right policy without an extra # launch-time parameter. Updated in sync_model_weights. self.async_llm.trinity_model_version = self.model_version - if self.config.enable_recording: + if self.config.enable_history: from trinity.common.models.vllm_patch.recording.recorder import ( TRINITY_MM_RENDER_ATTR, create_vllm_recorder, diff --git a/trinity/common/models/vllm_patch/__init__.py b/trinity/common/models/vllm_patch/__init__.py index cf54d3dec2d..3f2452700e9 100644 --- a/trinity/common/models/vllm_patch/__init__.py +++ b/trinity/common/models/vllm_patch/__init__.py @@ -158,10 +158,10 @@ def get_api_server( # wires the in-process MemoryStore + engine wrap itself — no static config # needs threading (the logprob width is a recorder-internal constant, and # the checkpoint version is read live off the engine). - recording = bool(config.enable_recording) and vllm_version >= VLLM_VERSION_0230 - if config.enable_recording and not recording: + recording = bool(config.enable_history) and vllm_version >= VLLM_VERSION_0230 + if config.enable_history and not recording: logger.warning( - "enable_recording is on but vLLM %s < 0.23.0; recording disabled", + "enable_history is on but vLLM %s < 0.23.0; recording disabled", vllm.__version__, ) diff --git a/trinity/common/workflows/step_wise_workflow.py b/trinity/common/workflows/step_wise_workflow.py index 1eee6bdefd5..a03757244b1 100644 --- a/trinity/common/workflows/step_wise_workflow.py +++ b/trinity/common/workflows/step_wise_workflow.py @@ -12,10 +12,9 @@ def __init__( self, *, task: Task, model: ModelWrapper, auxiliary_models=None, use_openai_client=True ): super().__init__(task=task, model=model, auxiliary_models=auxiliary_models) - assert model.enable_history or model.config.enable_recording, ( - "Rollout Model must have history or recording enabled for step-wise rewards, " - "please set `explorer.rollout_model.enable_history` or " - "`explorer.rollout_model.enable_recording` to `True` in your config." + assert model.enable_history, ( + "Rollout Model must have history enabled for step-wise rewards, " + "please set `explorer.rollout_model.enable_history` to `True` in your config." ) # use the rollout model's OpenAI client to write your agent application if use_openai_client: @@ -122,10 +121,9 @@ def __init__( self, *, task: Task, model: ModelWrapper, auxiliary_models=None, use_openai_client=True ): super().__init__(task=task, model=model, auxiliary_models=auxiliary_models) - assert model.enable_history or model.config.enable_recording, ( - "Rollout Model must have history or recording enabled for step-wise rewards, " - "please set `explorer.rollout_model.enable_history` or " - "`explorer.rollout_model.enable_recording` to `True` in your config." + assert model.enable_history, ( + "Rollout Model must have history enabled for step-wise rewards, " + "please set `explorer.rollout_model.enable_history` to `True` in your config." ) # use the rollout model's OpenAI client to write your agent application if use_openai_client: diff --git a/trinity/common/workflows/workflow.py b/trinity/common/workflows/workflow.py index 328a1d237a8..83e1d8dad81 100644 --- a/trinity/common/workflows/workflow.py +++ b/trinity/common/workflows/workflow.py @@ -93,7 +93,7 @@ class Workflow: # call carries a per-sample ``record_key`` (the MemoryStore group key) so # the runner can later report ``{record_key: reward}`` and the coordinator # can join reward inside the store. Set by the WorkflowRunner from - # ``config.explorer.rollout_model.enable_recording``. See ``SimpleWorkflow`` + # ``config.explorer.rollout_model.enable_history``. See ``SimpleWorkflow`` # for the per-sample (n=1) loop this triggers. enable_recording: bool = False diff --git a/trinity/explorer/rollout_coordinator.py b/trinity/explorer/rollout_coordinator.py index 2591cda6192..922d6afd325 100644 --- a/trinity/explorer/rollout_coordinator.py +++ b/trinity/explorer/rollout_coordinator.py @@ -78,7 +78,7 @@ def __init__( def _enable_recording(self) -> bool: """Whether the recording-consume path is active for train batches.""" - return bool(self.config.explorer.rollout_model.enable_recording) + return bool(self.config.explorer.rollout_model.enable_history) def _resolve_rank_urls(self) -> Dict[int, str]: """Resolve each rollout engine's API server URL via named Ray actors. diff --git a/trinity/explorer/workflow_runner.py b/trinity/explorer/workflow_runner.py index 2fee317ca8e..581f3a2d147 100644 --- a/trinity/explorer/workflow_runner.py +++ b/trinity/explorer/workflow_runner.py @@ -131,7 +131,7 @@ def _create_workflow_instance(self, task: Task, record_key: Optional[str] = None return self.workflow_instance def _enable_recording(self) -> bool: - return bool(self.config.explorer.rollout_model.enable_recording) + return bool(self.config.explorer.rollout_model.enable_history) async def _run_workflow(self, workflow_instance: Workflow) -> List[Experience]: if workflow_instance.asynchronous: From 82d75baad29cd604f8d2de604f97a8aa937528ea Mon Sep 17 00:00:00 2001 From: "panxuchen.pxc" Date: Fri, 26 Jun 2026 11:25:04 +0800 Subject: [PATCH 21/84] fix models --- tests/common/experience_extraction_test.py | 103 ------- tests/common/sglang_test.py | 3 +- tests/common/vllm_test.py | 3 +- trinity/common/config_validator.py | 2 +- trinity/common/experience.py | 4 +- .../common/models/experience_extraction.py | 283 ------------------ trinity/common/models/model.py | 10 +- trinity/common/models/recording/__init__.py | 14 +- trinity/common/models/recording/query.py | 10 +- trinity/common/models/recording/recorder.py | 2 + trinity/common/models/sglang_model.py | 39 ++- .../models/sglang_patch/recording/models.py | 4 +- .../models/sglang_patch/recording/recorder.py | 8 +- .../models/sglang_patch/recording/server.py | 40 +-- .../models/sglang_patch/server_patch.py | 6 +- trinity/common/models/tinker_model.py | 55 +++- .../models/vllm_patch/recording/models.py | 2 +- .../models/vllm_patch/recording/recorder.py | 18 +- .../models/vllm_patch/recording/server.py | 47 +-- .../workflows/on_policy_distill_workflow.py | 1 + 20 files changed, 149 insertions(+), 505 deletions(-) delete mode 100644 tests/common/experience_extraction_test.py delete mode 100644 trinity/common/models/experience_extraction.py diff --git a/tests/common/experience_extraction_test.py b/tests/common/experience_extraction_test.py deleted file mode 100644 index dea984f6e48..00000000000 --- a/tests/common/experience_extraction_test.py +++ /dev/null @@ -1,103 +0,0 @@ -import io -from types import SimpleNamespace -from unittest import TestCase - -import numpy as np -import pybase64 -import torch - -from trinity.common.models.experience_extraction import convert_api_output_to_experience - - -class TestExperienceExtraction(TestCase): - def test_convert_completion_output_extracts_sglang_routed_experts(self): - routed_experts = torch.tensor( - [ - [[1, 2], [3, 4]], - [[5, 6], [7, 8]], - [[9, 10], [11, 12]], - ], - dtype=torch.int32, - ) - routed_experts_b64 = pybase64.b64encode(routed_experts.numpy().tobytes()).decode("utf-8") - output = SimpleNamespace( - model="mock-moe-model", - prompt_token_ids=[10, 11], - sglext={"routed_experts": routed_experts_b64}, - choices=[ - SimpleNamespace( - token_ids=[12, 13], - message=SimpleNamespace(content="done"), - logprobs=SimpleNamespace( - content=[SimpleNamespace(logprob=-0.1), SimpleNamespace(logprob=-0.2)] - ), - ) - ], - ) - - experiences = convert_api_output_to_experience(output, routed_experts_layout=(2, 2)) - - self.assertEqual(len(experiences), 1) - exp = experiences[0] - self.assertEqual(exp.prompt_length, 2) - self.assertEqual(exp.response_text, "done") - self.assertTrue(torch.equal(exp.logprobs, torch.tensor([-0.1, -0.2], dtype=torch.float32))) - self.assertIsNotNone(exp.routed_experts) - self.assertEqual(exp.routed_experts.dtype, torch.uint8) - self.assertEqual(tuple(exp.routed_experts.shape), (3, 2, 2)) - self.assertTrue(torch.equal(exp.routed_experts, routed_experts.to(torch.uint8))) - - def test_convert_completion_output_ignores_invalid_routed_experts_shape(self): - output = SimpleNamespace( - model="mock-moe-model", - prompt_token_ids=[10, 11], - sglext={"routed_experts": "aW52YWxpZA=="}, - choices=[ - SimpleNamespace( - token_ids=[12, 13], - message=SimpleNamespace(content="done"), - logprobs=None, - ) - ], - ) - - experiences = convert_api_output_to_experience(output, routed_experts_layout=(2, 2)) - - self.assertEqual(len(experiences), 1) - self.assertIsNone(experiences[0].routed_experts) - - def test_convert_completion_output_extracts_vllm_routed_experts(self): - routed_experts = np.array( - [ - [[1, 2], [3, 4]], - [[5, 6], [7, 8]], - [[9, 10], [11, 12]], - ], - dtype=np.uint8, - ) - buffer = io.BytesIO() - np.save(buffer, routed_experts) - routed_experts_b64 = pybase64.b64encode(buffer.getvalue()).decode("utf-8") - output = SimpleNamespace( - model="mock-moe-model", - prompt_token_ids=[10, 11], - choices=[ - SimpleNamespace( - token_ids=[12, 13], - message=SimpleNamespace(content="done"), - logprobs=SimpleNamespace( - content=[SimpleNamespace(logprob=-0.1), SimpleNamespace(logprob=-0.2)] - ), - routed_experts=routed_experts_b64, - ) - ], - ) - - experiences = convert_api_output_to_experience(output, routed_experts_layout=(2, 2)) - - self.assertEqual(len(experiences), 1) - exp = experiences[0] - self.assertIsNotNone(exp.routed_experts) - self.assertEqual(exp.routed_experts.dtype, torch.uint8) - self.assertEqual(tuple(exp.routed_experts.shape), (3, 2, 2)) - self.assertTrue(torch.equal(exp.routed_experts, torch.tensor(routed_experts))) diff --git a/tests/common/sglang_test.py b/tests/common/sglang_test.py index 19bf61e431b..7730452e5cf 100644 --- a/tests/common/sglang_test.py +++ b/tests/common/sglang_test.py @@ -499,7 +499,7 @@ def _assert_recorded_experience(self, exp: Experience, record_key: str): # recording hot path (decode token ids lazily where a check is needed). if exp.prompt_text is not None: self.assertGreater(len(exp.prompt_text), 0) - self.assertIsNotNone(exp.response_text) + self.assertGreater(len(exp.response_text), 0) def _assert_recorded_routed_experts(self, exp: Experience): # enable_return_routed_experts is forced on by enable_recording. @@ -608,6 +608,7 @@ async def test_record(self): # noqa: C901 response_token_ids = consumed[0].tokens[consumed[0].prompt_length :].tolist() decoded_content = self.tokenizer.decode(response_token_ids, skip_special_tokens=True) self.assertEqual(decoded_content, content) + self.assertEqual(consumed[0].response_text, content) self.assertNotIn(rk_str, await self._list_record_keys()) # ===== 5. OpenAI tool-call parsing (HTTP) ===== diff --git a/tests/common/vllm_test.py b/tests/common/vllm_test.py index 1ce6a1cd014..70263cbe757 100644 --- a/tests/common/vllm_test.py +++ b/tests/common/vllm_test.py @@ -1946,7 +1946,7 @@ def _assert_recorded_experience(self, exp: Experience, record_key: str): # engine output recorded below may carry an empty native # CompletionOutput.text even when response token ids are present; avoid # decoding tokens on the recording hot path just to populate this field. - self.assertIsNotNone(exp.response_text) + self.assertGreater(len(exp.response_text), 0) def _assert_recorded_routed_experts(self, exp: Experience): # enable_return_routed_experts is forced on by enable_recording. @@ -2054,6 +2054,7 @@ async def test_record(self): # noqa: C901 response_token_ids = consumed[0].tokens[consumed[0].prompt_length :].tolist() decoded_content = self.tokenizer.decode(response_token_ids, skip_special_tokens=True) self.assertEqual(decoded_content, content) + self.assertEqual(consumed[0].response_text, content) self.assertNotIn(rk_str, await self._list_record_keys()) # ===== 5. OpenAI tool usage (HTTP) ===== diff --git a/trinity/common/config_validator.py b/trinity/common/config_validator.py index f6b7e909063..f9f1052ce8a 100644 --- a/trinity/common/config_validator.py +++ b/trinity/common/config_validator.py @@ -651,7 +651,7 @@ class ExplorerConfigValidator(ConfigValidator): over-rollout ratio validation, and LoRA configuration processing. """ - def validate(self, config: Config) -> None: + def validate(self, config: Config) -> None: # noqa: C901 """Validate and configure explorer-specific settings. - Inherits model configuration from the global model config to rollout models diff --git a/trinity/common/experience.py b/trinity/common/experience.py index 1a104055ac9..45c4c2903ce 100644 --- a/trinity/common/experience.py +++ b/trinity/common/experience.py @@ -131,7 +131,7 @@ class Experience: ) eid: EID = field(default_factory=EID) # Unique identifier for the experience - tokens: Optional[Tensor] = None # [seq_length] + tokens: Tensor = field(default_factory=lambda: torch.tensor([])) # [seq_length] prompt_length: int = 1 # Length of the prompt in tokens, used for generating attention masks logprobs: Optional[Tensor] = None # [resp_length] reward: Optional[float] = None @@ -149,7 +149,7 @@ class Experience: ) # Metrics associated with the experience, directly used by the monitor # for single-turn experiences - response_text: Optional[str] = None # Text of the response + response_text: str = "" # Text of the response prompt_text: Optional[str] = None # Text of the prompt # for multi-turn experiences diff --git a/trinity/common/models/experience_extraction.py b/trinity/common/models/experience_extraction.py deleted file mode 100644 index e70eb4dffdf..00000000000 --- a/trinity/common/models/experience_extraction.py +++ /dev/null @@ -1,283 +0,0 @@ -import io -from typing import Any, Dict, List, Optional, Sequence, Tuple - -import numpy as np -import pybase64 -import torch -from torch import Tensor -from transformers import AutoConfig - -from trinity.common.experience import Experience -from trinity.common.models.mm_utils import combine_output_token_ids - - -def get_routed_experts_layout( - model_path: str, trust_remote_code: bool = True -) -> Optional[Tuple[int, int]]: - hf_config = AutoConfig.from_pretrained(model_path, trust_remote_code=trust_remote_code) - text_config = getattr(hf_config, "text_config", hf_config) - num_layers = getattr(text_config, "num_hidden_layers", None) - topk = getattr(text_config, "num_experts_per_tok", None) - if num_layers is None or topk is None: - return None - return int(num_layers), int(topk) - - -def decode_sglang_routed_experts( - routed_experts_value: Any, - total_tokens: int, - layout: Tuple[int, int], -) -> Optional[Tensor]: - if routed_experts_value is None: - return None - if isinstance(routed_experts_value, torch.Tensor): - return routed_experts_value.to(torch.uint8) - if not isinstance(routed_experts_value, str): - return torch.tensor(routed_experts_value, dtype=torch.uint8) - - decoded = pybase64.b64decode_as_bytearray(routed_experts_value) - routed_experts = torch.frombuffer(decoded, dtype=torch.int32) - num_layers, topk = layout - seq_length = max(total_tokens - 1, 0) - expected_numel = seq_length * num_layers * topk - if routed_experts.numel() != expected_numel: - raise ValueError( - "Unexpected routed_experts size from SGLang: " - f"expected {expected_numel} elements for shape ({seq_length}, {num_layers}, {topk}), " - f"got {routed_experts.numel()}" - ) - return routed_experts.reshape(seq_length, num_layers, topk).to(torch.uint8) - - -def decode_vllm_routed_experts(routed_experts_value: str | None) -> Optional[Tensor]: - if routed_experts_value is None: - return None - - decoded = pybase64.b64decode_as_bytearray(routed_experts_value) - routed_experts = np.load(io.BytesIO(decoded), allow_pickle=False) - return torch.as_tensor(routed_experts, dtype=torch.uint8) - - -def convert_api_output_to_experience( - output, - multi_modal_inputs: Optional[dict[str, torch.Tensor]] = None, - routed_experts_layout: Optional[Tuple[int, int]] = None, -) -> List[Experience]: - """Convert a non-stream API output to a list of experiences. - - Args: - output: Completion output from API client. - multi_modal_inputs: Optional training-time multimodal tensors aligned - with the prompt tokens. - routed_experts_layout: Optional `(num_layers, topk)` layout used to - decode routed experts. - """ - return _convert_completion_output_to_experience( - output, - multi_modal_inputs=multi_modal_inputs, - routed_experts_layout=routed_experts_layout, - ) - - -class HistoryRecordingStream: # TODO: add multi-modal support - def __init__(self, stream, history: List[Experience], is_async: bool = False) -> None: - self._stream = stream - self._history = history - self._chunks = [] - self._recorded = False - self._is_async = is_async - if is_async: - self._iterator = stream.__aiter__() - else: - self._iterator = iter(stream) - - def __iter__(self): - if self._is_async: - raise TypeError("Use 'async for' for async streams.") - return self - - def __next__(self): - if self._is_async: - raise TypeError("Use 'async for' for async streams.") - try: - chunk = next(self._iterator) - except StopIteration: - self._record_history_once() - raise - self._chunks.append(chunk) - return chunk - - def close(self) -> None: - if self._is_async: - raise TypeError("Use 'aclose' for async streams.") - self._record_history_once() - close_fn = getattr(self._stream, "close", None) - if callable(close_fn): - close_fn() - - def __aiter__(self): - if not self._is_async: - raise TypeError("Use 'for' for sync streams.") - return self - - async def __anext__(self): - if not self._is_async: - raise TypeError("Use 'for' for sync streams.") - try: - chunk = await self._iterator.__anext__() - except StopAsyncIteration: - self._record_history_once() - raise - self._chunks.append(chunk) - return chunk - - async def aclose(self) -> None: - if not self._is_async: - raise TypeError("Use 'close' for sync streams.") - self._record_history_once() - close_fn = getattr(self._stream, "aclose", None) - if callable(close_fn): - close_result = close_fn() - if hasattr(close_result, "__await__"): - await close_result - return - close_fn = getattr(self._stream, "close", None) - if callable(close_fn): - close_fn() - - def _record_history_once(self) -> None: - if self._recorded: - return - self._recorded = True - if self._chunks: - self._history.extend(_convert_stream_chunks_to_experience(self._chunks)) - - def __getattr__(self, name: str): - return getattr(self._stream, name) - - -def _convert_completion_output_to_experience( - output, - multi_modal_inputs: Optional[dict[str, torch.Tensor]] = None, - routed_experts_layout: Optional[Tuple[int, int]] = None, -) -> List[Experience]: - return [ - Experience( - tokens=torch.cat( - ( - torch.tensor(output.prompt_token_ids, dtype=torch.int32), - torch.tensor(choice.token_ids, dtype=torch.int32), - ) - ), - logprobs=extract_logprobs(choice), - prompt_length=len(output.prompt_token_ids), - response_text=getattr(choice.message, "content", None), - routed_experts=_extract_completion_routed_experts( - output, - choice, - total_tokens=len(output.prompt_token_ids) + len(choice.token_ids), - routed_experts_layout=routed_experts_layout, - ), - multi_modal_inputs=combine_output_token_ids(choice.token_ids, multi_modal_inputs), - ) - for choice in output.choices - ] - - -def _convert_stream_chunks_to_experience(chunks: Sequence[Any]) -> List[Experience]: - prompt_token_ids: Optional[List[int]] = None - by_choice: Dict[int, Dict[str, Any]] = {} - - for chunk in chunks: - if prompt_token_ids is None and hasattr(chunk, "prompt_token_ids"): - chunk_prompt_token_ids = getattr(chunk, "prompt_token_ids", None) - if chunk_prompt_token_ids is not None: - prompt_token_ids = list(chunk_prompt_token_ids) - - for choice in getattr(chunk, "choices", []) or []: - idx = getattr(choice, "index", 0) - if idx not in by_choice: - by_choice[idx] = { - "token_ids": [], - "logprobs": [], - "response_text_parts": [], - } - data = by_choice[idx] - - token_ids = getattr(choice, "token_ids", None) - if token_ids is not None: - data["token_ids"].extend(token_ids) - - choice_logprobs = getattr(choice, "logprobs", None) - if ( - choice_logprobs is not None - and getattr(choice_logprobs, "content", None) is not None - ): - for token_logprob in choice_logprobs.content: - data["logprobs"].append(token_logprob.logprob) - if token_ids is None: - token_id = getattr(token_logprob, "token_id", None) - if token_id is not None: - data["token_ids"].append(token_id) - - delta = getattr(choice, "delta", None) - if delta is not None: - delta_content = getattr(delta, "content", None) - if isinstance(delta_content, str) and len(delta_content) > 0: - data["response_text_parts"].append(delta_content) - - prompt_token_ids = prompt_token_ids or [] - exps: List[Experience] = [] - for idx in sorted(by_choice.keys()): - data = by_choice[idx] - response_token_ids = data["token_ids"] - if len(response_token_ids) == 0: - continue - response_text = "".join(data["response_text_parts"]) - exps.append( - Experience( - tokens=torch.tensor(prompt_token_ids + response_token_ids, dtype=torch.int32), - logprobs=torch.tensor(data["logprobs"], dtype=torch.float32), - prompt_length=len(prompt_token_ids), - response_text=response_text, - ) - ) - return exps - - -def _extract_completion_routed_experts( - output, - choice, - total_tokens: int, - routed_experts_layout: Optional[Tuple[int, int]] = None, -) -> Optional[Tensor]: - routed_experts_value = getattr(choice, "routed_experts", None) - if routed_experts_value is not None: - try: - return decode_vllm_routed_experts(routed_experts_value) - except (ValueError, OSError): - return None - - if routed_experts_layout is None: - return None - - if not hasattr(output, "sglext") or "routed_experts" not in output.sglext: - return None - routed_experts_value = output.sglext.get("routed_experts", None) - try: - return decode_sglang_routed_experts( - routed_experts_value, - total_tokens, - layout=routed_experts_layout, - ) - except ValueError: - return None - - -def extract_logprobs(choice) -> Tensor: - if not hasattr(choice, "logprobs") or choice.logprobs is None: - return torch.tensor([], dtype=torch.float32) - return torch.tensor( - [logprob.logprob for logprob in choice.logprobs.content], - dtype=torch.float32, - ) diff --git a/trinity/common/models/model.py b/trinity/common/models/model.py index 76096b9a656..faf284973e7 100644 --- a/trinity/common/models/model.py +++ b/trinity/common/models/model.py @@ -5,7 +5,7 @@ import copy import socket from abc import ABC, abstractmethod -from typing import TYPE_CHECKING, Any, Dict, List, Optional, Sequence, Tuple, Union +from typing import TYPE_CHECKING, Any, Dict, List, Optional, Sequence, Tuple import httpx import ray @@ -649,13 +649,11 @@ def chat_completions(*args, **kwargs): messages=messages, with_chat_completion=True, return_token_ids=self.enable_history, + record_key=(self._api_key if self.enable_history else None), **kwargs, ) ) response = chat_response.pop() - # TODO: Tinker legacy recording path - should be migrated to engine-side recording - if self.enable_history: - self.history.extend(chat_response) return response self.openai_client.chat.completions.create = chat_completions @@ -696,12 +694,10 @@ async def chat_completions(*args, **kwargs): messages=messages, with_chat_completion=True, return_token_ids=self.enable_history, + record_key=(self._api_key if self.enable_history else None), **kwargs, ) response = chat_response.pop() - # TODO: Tinker legacy recording path - should be migrated to engine-side recording - if self.enable_history: - self.history.extend(chat_response) return response self.openai_async_client.chat.completions.create = chat_completions diff --git a/trinity/common/models/recording/__init__.py b/trinity/common/models/recording/__init__.py index a9c0b3f105e..6b585dd2648 100644 --- a/trinity/common/models/recording/__init__.py +++ b/trinity/common/models/recording/__init__.py @@ -7,7 +7,15 @@ record_key_ctx, skip_recording_ctx, ) -from trinity.common.models.recording.recorder import Recorder +from trinity.common.models.recording.recorder import ( + TRINITY_RECORD_STORE_ATTR, + TRINITY_RECORDER_ATTR, + Recorder, +) +from trinity.common.models.recording.server import ( + add_recording_middleware, + mount_recording_api, +) from trinity.common.models.recording.store import MemoryStore, RecordStore __all__ = [ @@ -15,8 +23,12 @@ "Recorder", "RecordingIdentityMiddleware", "RecordStore", + "TRINITY_RECORD_STORE_ATTR", + "TRINITY_RECORDER_ATTR", + "add_recording_middleware", "extract_bearer_token", "get_recording_record_key", + "mount_recording_api", "record_key_ctx", "skip_recording_ctx", ] diff --git a/trinity/common/models/recording/query.py b/trinity/common/models/recording/query.py index 1acbd58eea7..98dc33dbdba 100644 --- a/trinity/common/models/recording/query.py +++ b/trinity/common/models/recording/query.py @@ -6,11 +6,15 @@ from pydantic import BaseModel from trinity.common.experience import Experience -from trinity.common.models.recording.recorder import Recorder +from trinity.common.models.recording.recorder import ( + TRINITY_RECORD_STORE_ATTR, + TRINITY_RECORDER_ATTR, + Recorder, +) from trinity.common.models.recording.store import RecordStore -STORE_STATE_ATTR = "trinity_record_store" -RECORDER_STATE_ATTR = "trinity_recorder" +STORE_STATE_ATTR = TRINITY_RECORD_STORE_ATTR +RECORDER_STATE_ATTR = TRINITY_RECORDER_ATTR query_router = APIRouter(prefix="/records", tags=["trinity-recording"]) diff --git a/trinity/common/models/recording/recorder.py b/trinity/common/models/recording/recorder.py index 1515f11a7b2..fae8caa378e 100644 --- a/trinity/common/models/recording/recorder.py +++ b/trinity/common/models/recording/recorder.py @@ -11,6 +11,8 @@ from trinity.common.models.recording.store import RecordStore MODEL_VERSION_ATTR = "trinity_model_version" +TRINITY_RECORDER_ATTR = "trinity_recorder" +TRINITY_RECORD_STORE_ATTR = "trinity_record_store" BuildExperiencesFn = Callable[..., Sequence[Experience]] diff --git a/trinity/common/models/sglang_model.py b/trinity/common/models/sglang_model.py index 96881a58a64..45e213504fa 100644 --- a/trinity/common/models/sglang_model.py +++ b/trinity/common/models/sglang_model.py @@ -9,17 +9,43 @@ from typing import Any, List, Literal, Optional, Sequence, Tuple import httpx +import pybase64 import torch from transformers import AutoTokenizer from trinity.common.config import InferenceModelConfig from trinity.common.constants import ROLLOUT_WEIGHT_SYNC_GROUP_NAME, SyncMethod from trinity.common.experience import Experience -from trinity.common.models.experience_extraction import decode_sglang_routed_experts from trinity.common.models.model import BaseInferenceModel from trinity.manager.synchronizer import Synchronizer +def decode_sglang_routed_experts( + routed_experts_value: Any, + total_tokens: int, + layout: Tuple[int, int], +) -> Optional[torch.Tensor]: + if routed_experts_value is None: + return None + if isinstance(routed_experts_value, torch.Tensor): + return routed_experts_value.to(torch.uint8) + if not isinstance(routed_experts_value, str): + return torch.tensor(routed_experts_value, dtype=torch.uint8) + + decoded = pybase64.b64decode_as_bytearray(routed_experts_value) + routed_experts = torch.frombuffer(decoded, dtype=torch.int32) + num_layers, topk = layout + seq_length = max(total_tokens - 1, 0) + expected_numel = seq_length * num_layers * topk + if routed_experts.numel() != expected_numel: + raise ValueError( + "Unexpected routed_experts size from SGLang: " + f"expected {expected_numel} elements for shape ({seq_length}, {num_layers}, {topk}), " + f"got {routed_experts.numel()}" + ) + return routed_experts.reshape(seq_length, num_layers, topk).to(torch.uint8) + + class SGLangClient: """A simple http client to interact with the SGLang API server.""" @@ -71,12 +97,7 @@ async def _server_call( async def health_check(self) -> bool: try: - async with httpx.AsyncClient( - headers={ - "Content-Type": "application/json; charset=utf-8", - "Authorization": f"Bearer {self.api_key}" if self.api_key else "", - } - ) as client: + async with httpx.AsyncClient() as client: response = await client.get(f"{self.server_url}/health", timeout=5) return response.status_code == 200 except Exception as e: @@ -546,7 +567,7 @@ async def run_api_server(self) -> bool: record_store = None recorder = None routed_experts_layout = None - if self.config.enable_recording: + if self.config.enable_history: from trinity.common.models.recording.recorder import Recorder from trinity.common.models.recording.store import MemoryStore from trinity.common.models.sglang_patch.recording.models import ( @@ -589,7 +610,7 @@ async def run_api_server(self) -> bool: master_addr=self.master_addr, master_port=self.master_port, enable_return_routed_experts=self.config.enable_return_routed_experts, - enable_recording=self.config.enable_recording, + enable_history=self.config.enable_history, recorder=recorder, record_store=record_store, routed_experts_layout=routed_experts_layout, diff --git a/trinity/common/models/sglang_patch/recording/models.py b/trinity/common/models/sglang_patch/recording/models.py index 6dcc2a277ab..8478a48941e 100644 --- a/trinity/common/models/sglang_patch/recording/models.py +++ b/trinity/common/models/sglang_patch/recording/models.py @@ -29,7 +29,7 @@ import torch from trinity.common.experience import EID, Experience -from trinity.common.models.experience_extraction import decode_sglang_routed_experts +from trinity.common.models.sglang_model import decode_sglang_routed_experts def _extract_output_logprobs(meta_info: dict) -> List[float]: @@ -162,7 +162,7 @@ def build_sglang_experience( logprobs=response_logprobs, prompt_length=prompt_length, prompt_text=item.get("prompt_text"), - response_text=item.get("text"), + response_text=item.get("text") or "", routed_experts=routed_experts, info=info, ) diff --git a/trinity/common/models/sglang_patch/recording/recorder.py b/trinity/common/models/sglang_patch/recording/recorder.py index be93485775e..23253e3bc80 100644 --- a/trinity/common/models/sglang_patch/recording/recorder.py +++ b/trinity/common/models/sglang_patch/recording/recorder.py @@ -28,14 +28,16 @@ from typing import Any, List, Optional, Tuple from trinity.common.models.recording.context import record_key_ctx -from trinity.common.models.recording.recorder import Recorder +from trinity.common.models.recording.recorder import ( + TRINITY_RECORD_STORE_ATTR, + TRINITY_RECORDER_ATTR, + Recorder, +) from trinity.common.models.recording.store import MemoryStore, RecordStore from trinity.common.models.sglang_patch.recording.models import build_sglang_experience #: Guard attribute marking the wrapped generate_request (mirrors vLLM's style). _PATCHED_FLAG = "__patched_sglang_recording__" -TRINITY_RECORDER_ATTR = "trinity_recorder" -TRINITY_RECORD_STORE_ATTR = "trinity_record_store" def _get_obj(args, kwargs): diff --git a/trinity/common/models/sglang_patch/recording/server.py b/trinity/common/models/sglang_patch/recording/server.py index f8e5ffbf158..54a2e711ca7 100644 --- a/trinity/common/models/sglang_patch/recording/server.py +++ b/trinity/common/models/sglang_patch/recording/server.py @@ -15,30 +15,12 @@ import logging from typing import Optional, Tuple -from trinity.common.models.recording.context import RecordingIdentityMiddleware -from trinity.common.models.recording.query import ( - RECORDER_STATE_ATTR, - STORE_STATE_ATTR, - query_router, -) from trinity.common.models.recording.recorder import Recorder +from trinity.common.models.recording.server import mount_recording_api from trinity.common.models.recording.store import RecordStore from trinity.common.models.sglang_patch.recording.recorder import create_sglang_recorder -def _add_recording_middleware(app) -> None: - """Install recording middleware before serving, even if SGLang built the stack. - - Starlette rejects ``add_middleware`` after ``middleware_stack`` is built with - "Cannot add middleware after an application has started". Clearing the cached - stack lets Starlette rebuild it with our middleware on first request (same - defensive pattern as the vLLM recording patch). - """ - if getattr(app, "middleware_stack", None) is not None: - app.middleware_stack = None - app.add_middleware(RecordingIdentityMiddleware) - - def setup_sglang_recording( tokenizer_manager, app, @@ -62,19 +44,11 @@ def setup_sglang_recording( routed_experts_layout=routed_experts_layout, ) - # (2) in-process middleware: Authorization bearer -> record_key contextvar. - _add_recording_middleware(app) - - # (3) query routes mounted on the main app; OpenAI /v1/* surface untouched. - app.include_router(query_router) - - setattr(app.state, STORE_STATE_ATTR, recorder.store) - setattr(app.state, RECORDER_STATE_ATTR, recorder) - - recorder.start() - - logger.info( - "SGLang generation recording enabled: store=%s", - type(recorder.store).__name__, + mount_recording_api( + app, + recorder, + logger, + engine_name="SGLang", + start_recorder=True, ) return recorder diff --git a/trinity/common/models/sglang_patch/server_patch.py b/trinity/common/models/sglang_patch/server_patch.py index 5233204adc8..46f1c1ad5f3 100644 --- a/trinity/common/models/sglang_patch/server_patch.py +++ b/trinity/common/models/sglang_patch/server_patch.py @@ -304,7 +304,7 @@ def get_api_server( master_addr: Optional[str], master_port: Optional[int], logger: Logger, - enable_recording: bool = False, + enable_history: bool = False, recorder: Optional[Any] = None, record_store: Optional[Any] = None, routed_experts_layout: Optional[Tuple[int, int]] = None, @@ -341,7 +341,7 @@ def get_api_server( # api_key, so the auth middleware would otherwise 401-reject it. This # mirrors vLLM, whose recording server sets no api_key auth. The # embedded server is localhost/in-Ray-actor, so auth is not needed. - api_key=None if enable_recording else api_key, + api_key=None if enable_history else api_key, # SGLang enables tool calling via tool_call_parser (no separate # enable_auto_tool_choice flag in this version). Only render/parse tools # when a parser is configured, matching vLLM's enable_auto_tool_choice. @@ -385,7 +385,7 @@ def get_api_server( # recorder/store are owned by ``SGLangRolloutModel``; this installs the # engine wrap on ``tokenizer_manager``, ``RecordingIdentityMiddleware`` and # ``query_router`` on ``app``, and stashes store/recorder on ``app.state``. - if enable_recording: + if enable_history: from trinity.common.models.sglang_patch.recording import setup_sglang_recording setup_sglang_recording( diff --git a/trinity/common/models/tinker_model.py b/trinity/common/models/tinker_model.py index 5d05db240a2..2228fa7a6e3 100644 --- a/trinity/common/models/tinker_model.py +++ b/trinity/common/models/tinker_model.py @@ -12,9 +12,32 @@ from trinity.common.constants import SyncMethod from trinity.common.experience import Experience from trinity.common.models.model import BaseInferenceModel +from trinity.common.models.recording.recorder import MODEL_VERSION_ATTR, Recorder +from trinity.common.models.recording.store import MemoryStore from trinity.manager.synchronizer import Synchronizer +def _build_tinker_experiences( + experiences: Sequence[Experience], + record_key: str, + *, + rank: int, + timestamp: str, + model_version: Optional[int] = None, + request_id: str, +) -> Sequence[Experience]: + for index, exp in enumerate(experiences): + if exp.info is None: + exp.info = {} + exp.info["record_key"] = record_key + exp.info["request_id"] = f"{request_id}:{index}" + exp.info["rank"] = rank + exp.info["timestamp"] = timestamp + if model_version is not None: + exp.info["model_version"] = model_version + return experiences + + class TinkerModel(BaseInferenceModel): def __init__( self, @@ -25,6 +48,17 @@ def __init__( self.synchronizer = Synchronizer.get_actor(namespace=ray.get_runtime_context().namespace) self.model = None self.model_path = config.model_path + self.request_id = 0 + self.recorder = None + if self.config.enable_history: + self.recorder = Recorder( + store=MemoryStore(), + build_experiences=_build_tinker_experiences, + enabled=True, + rank=0, + engine_client=self, + ) + setattr(self, MODEL_VERSION_ATTR, self.model_version) async def _initialize_tokenizer(self) -> None: """Initialize the tokenizer.""" @@ -48,11 +82,15 @@ async def _generate_internal(self, prompt: dict, **kwargs) -> types.SampleRespon topk_prompt_logprobs=kwargs.get("topk_prompt_logprobs", self.config.logprobs), ) - async def generate(self, prompt: str, **kwargs) -> Sequence[Experience]: + async def generate(self, prompt: str, lora_request=None, **kwargs) -> Sequence[Experience]: """Generate a responses from a prompt in async.""" if self.tokenizer is None: await self._initialize_tokenizer() + record_key = kwargs.pop("record_key", None) + request_id = str(self.request_id) + self.request_id += 1 + returned_seq, is_valid = self._handle_prompt_truncation(prompt, **kwargs) if not is_valid: return returned_seq # is_valid is False: returned_seq is a list of dummy experiences @@ -118,6 +156,13 @@ async def generate(self, prompt: str, **kwargs) -> Sequence[Experience]: ) experiences.append(chat_completion) + if self.recorder is not None and record_key is not None: + self.recorder.schedule_record( + experiences[: len(output.sequences)], + record_key, + request_id=request_id, + ) + return experiences async def chat(self, messages: List[dict], **kwargs) -> Sequence[Experience]: @@ -149,6 +194,8 @@ async def prepare(self) -> None: base_model=self.config.model_path, ) await self._initialize_tokenizer() + if self.recorder is not None: + self.recorder.start() async def sync_model_weights( self, @@ -162,6 +209,7 @@ async def sync_model_weights( model_path=remote_sampler_path, ) self.model_path = remote_sampler_path + setattr(self, MODEL_VERSION_ATTR, self.model_version) return model_version def get_model_version(self) -> int: @@ -189,3 +237,8 @@ def get_api_key(self): def get_model_path(self) -> Optional[str]: """Get the latest sampler weight path.""" return self.model_path + + async def shutdown(self) -> None: + if self.recorder is not None: + await self.recorder.stop() + self.recorder = None diff --git a/trinity/common/models/vllm_patch/recording/models.py b/trinity/common/models/vllm_patch/recording/models.py index b6817c40d28..2cad37febcd 100644 --- a/trinity/common/models/vllm_patch/recording/models.py +++ b/trinity/common/models/vllm_patch/recording/models.py @@ -211,7 +211,7 @@ def build_experience( prompt_length=prompt_length, routed_experts=routed_experts, prompt_text=prompt_text if prompt_text is not None else output.prompt, - response_text=completion.text, + response_text=getattr(completion, "text", None) or "", multi_modal_inputs=combine_output_token_ids( response_token_ids, multi_modal_inputs, diff --git a/trinity/common/models/vllm_patch/recording/recorder.py b/trinity/common/models/vllm_patch/recording/recorder.py index 5f9351776f3..1c14bbe49a4 100644 --- a/trinity/common/models/vllm_patch/recording/recorder.py +++ b/trinity/common/models/vllm_patch/recording/recorder.py @@ -23,7 +23,11 @@ from typing import Optional from trinity.common.models.recording.context import record_key_ctx -from trinity.common.models.recording.recorder import Recorder +from trinity.common.models.recording.recorder import ( + TRINITY_RECORD_STORE_ATTR, + TRINITY_RECORDER_ATTR, + Recorder, +) from trinity.common.models.recording.store import MemoryStore, RecordStore from trinity.common.models.vllm_patch.recording.models import build_experience @@ -41,8 +45,6 @@ #: to thread a knob through the launcher. The engine's ``max_logprobs`` cap #: (default 20, set at engine build) already covers it. _RECORDER_LOGPROB_WIDTH = 1 -TRINITY_RECORDER_ATTR = "trinity_recorder" -TRINITY_RECORD_STORE_ATTR = "trinity_record_store" TRINITY_MM_RENDER_ATTR = "trinity_mm_render" @@ -176,8 +178,9 @@ def _accumulate_request_output(state, output, *, is_delta_output: bool): acc = state["outputs"][index] cur_token_ids = _list_or_empty(getattr(completion, "token_ids", None)) - if is_delta_output and cur_token_ids: - acc["token_ids"].extend(cur_token_ids) + if is_delta_output: + if cur_token_ids: + acc["token_ids"].extend(cur_token_ids) else: acc["token_ids"] = cur_token_ids @@ -192,8 +195,9 @@ def _accumulate_request_output(state, output, *, is_delta_output: bool): acc["logprobs"] = cur_logprobs cur_text = getattr(completion, "text", None) or "" - if is_delta_output and cur_text: - acc["text"] += cur_text + if is_delta_output: + if cur_text: + acc["text"] += cur_text else: acc["text"] = cur_text diff --git a/trinity/common/models/vllm_patch/recording/server.py b/trinity/common/models/vllm_patch/recording/server.py index 27084111098..b58706c26e8 100644 --- a/trinity/common/models/vllm_patch/recording/server.py +++ b/trinity/common/models/vllm_patch/recording/server.py @@ -40,21 +40,9 @@ from vllm.utils.system_utils import set_ulimit from vllm.version import __version__ as VLLM_VERSION -from trinity.common.models.recording.context import RecordingIdentityMiddleware -from trinity.common.models.recording.query import ( - RECORDER_STATE_ATTR, - STORE_STATE_ATTR, - query_router, -) -from trinity.common.models.recording.recorder import Recorder -from trinity.common.models.recording.store import RecordStore +from trinity.common.models.recording.recorder import TRINITY_RECORDER_ATTR, Recorder +from trinity.common.models.recording.server import mount_recording_api from trinity.common.models.vllm_patch import get_vllm_version -from trinity.common.models.vllm_patch.recording.recorder import TRINITY_RECORDER_ATTR - -#: Attribute on app.state holding the active RecordStore. -_STORE_STATE_ATTR = STORE_STATE_ATTR -#: Attribute on app.state holding the active Recorder. -_RECORDER_STATE_ATTR = RECORDER_STATE_ATTR def setup_server_in_ray(args, logger): @@ -104,20 +92,6 @@ def dummy_add_signal_handler(self, *args, **kwargs): pass -def _add_recording_middleware(app) -> None: - """Install recording middleware before serving, even if vLLM built the stack. - - Some vLLM/FastAPI versions touch ``app.middleware_stack`` inside - ``build_app``. Starlette then rejects ``add_middleware`` with "Cannot add - middleware after an application has started", even though uvicorn has not - started serving yet. Clearing the cached stack lets Starlette rebuild it - with our middleware on first request. - """ - if getattr(app, "middleware_stack", None) is not None: - app.middleware_stack = None - app.add_middleware(RecordingIdentityMiddleware) - - def _setup_recording( engine_client, app, @@ -150,22 +124,7 @@ def _setup_recording( "Generation recording API server requires vLLMRolloutModel to install " "engine_client.trinity_recorder before server startup." ) - store: RecordStore = recorder.store - - # (2) in-process middleware: API key -> contextvar. Zero network hop. - _add_recording_middleware(app) - - # (3) query routes mounted on the main app; OpenAI /v1/* surface untouched. - app.include_router(query_router) - - setattr(app.state, _STORE_STATE_ATTR, store) - setattr(app.state, _RECORDER_STATE_ATTR, recorder) - - logger.info( - "Generation recording enabled: store=%s rank=%d", - type(store).__name__, - recorder.rank, - ) + mount_recording_api(app, recorder, logger, engine_name="vLLM") return recorder diff --git a/trinity/common/workflows/on_policy_distill_workflow.py b/trinity/common/workflows/on_policy_distill_workflow.py index e9251033741..f84818d35c3 100644 --- a/trinity/common/workflows/on_policy_distill_workflow.py +++ b/trinity/common/workflows/on_policy_distill_workflow.py @@ -117,6 +117,7 @@ async def run_async(self) -> List[Experience]: resp_start = response.prompt_length - 1 teacher_resp_logprobs = teacher_logprobs[resp_start:] student_resp_logprobs = response.logprobs + assert student_resp_logprobs is not None, "Student logprobs should not be None." # Verify lengths match (they should be equal for the same token sequence) assert len(teacher_resp_logprobs) == len(student_resp_logprobs), ( From 5cc3f050c4e479cac1c4ba2d375acd3caac616b6 Mon Sep 17 00:00:00 2001 From: "panxuchen.pxc" Date: Fri, 26 Jun 2026 11:29:13 +0800 Subject: [PATCH 22/84] add recording server --- trinity/common/models/recording/server.py | 49 +++++++++++++++++++++++ 1 file changed, 49 insertions(+) create mode 100644 trinity/common/models/recording/server.py diff --git a/trinity/common/models/recording/server.py b/trinity/common/models/recording/server.py new file mode 100644 index 00000000000..42de0e909cd --- /dev/null +++ b/trinity/common/models/recording/server.py @@ -0,0 +1,49 @@ +"""Shared HTTP server wiring for generation recording.""" + +import logging + +from trinity.common.models.recording.context import RecordingIdentityMiddleware +from trinity.common.models.recording.query import ( + RECORDER_STATE_ATTR, + STORE_STATE_ATTR, + query_router, +) +from trinity.common.models.recording.recorder import Recorder + + +def add_recording_middleware(app) -> None: + """Install recording middleware before serving. + + Some FastAPI/Starlette integrations build ``middleware_stack`` before + uvicorn starts serving. Clearing the cached stack lets Starlette rebuild it + with our middleware on first request. + """ + if getattr(app, "middleware_stack", None) is not None: + app.middleware_stack = None + app.add_middleware(RecordingIdentityMiddleware) + + +def mount_recording_api( + app, + recorder: Recorder, + logger: logging.Logger, + *, + engine_name: str, + start_recorder: bool = False, +) -> None: + """Mount recording middleware/query routes and expose state to handlers.""" + add_recording_middleware(app) + app.include_router(query_router) + + setattr(app.state, STORE_STATE_ATTR, recorder.store) + setattr(app.state, RECORDER_STATE_ATTR, recorder) + + if start_recorder: + recorder.start() + + logger.info( + "%s generation recording enabled: store=%s rank=%d", + engine_name, + type(recorder.store).__name__, + recorder.rank, + ) From 936da532fb0044092c5d334fbe6d103e5be20db1 Mon Sep 17 00:00:00 2001 From: "panxuchen.pxc" Date: Fri, 26 Jun 2026 11:32:31 +0800 Subject: [PATCH 23/84] fix sglang --- trinity/common/config_validator.py | 6 ++++-- trinity/common/models/allocator.py | 10 +++++----- 2 files changed, 9 insertions(+), 7 deletions(-) diff --git a/trinity/common/config_validator.py b/trinity/common/config_validator.py index f9f1052ce8a..a22b18b022d 100644 --- a/trinity/common/config_validator.py +++ b/trinity/common/config_validator.py @@ -705,8 +705,10 @@ def validate(self, config: Config) -> None: # noqa: C901 if config.mode == "serve": # in 'serve' mode, we always enable openai api for rollout model config.explorer.rollout_model.enable_openai_api = True - if config.explorer.rollout_model.engine_type.startswith("vllm"): - # enable_history is the single switch for vLLM recording. + if config.explorer.rollout_model.engine_type.startswith("vllm") or ( + config.explorer.rollout_model.engine_type == "sglang" + ): + # enable_history is the single switch for engine-side recording. if config.explorer.rollout_model.enable_history: config.explorer.rollout_model.enable_openai_api = True self._validate_inference_parallel_config(config.explorer.rollout_model, "rollout_model") diff --git a/trinity/common/models/allocator.py b/trinity/common/models/allocator.py index f4eae293a35..089a9907669 100644 --- a/trinity/common/models/allocator.py +++ b/trinity/common/models/allocator.py @@ -82,13 +82,13 @@ async def create_engine( config = deepcopy(config) config.engine_id = engine_id - if config.engine_type.startswith("vllm"): - # enable_history is the single switch for vLLM recording. + if config.engine_type.startswith("vllm") or config.engine_type == "sglang": + # enable_history is the single switch for engine-side recording. if config.enable_history: config.enable_openai_api = True - # In-vLLM recording: force routed_experts capture (the engine reads - # this at build time — see vllm_model.py). VLLMModel mirrors the - # recording config onto the engine instance for the recorder to read. + # Recording: force routed_experts capture (the engine reads + # this at build time). The model mirrors the recording config + # onto the engine/server instance for the recorder to read. # No env / runtime_env involved. config.enable_return_routed_experts = True From 2228d6594845bd989f9e9cba9a140ea5be655be5 Mon Sep 17 00:00:00 2001 From: "panxuchen.pxc" Date: Fri, 26 Jun 2026 11:39:14 +0800 Subject: [PATCH 24/84] remove redundant fields --- tests/common/sglang_test.py | 15 +++++++-------- tests/common/vllm_test.py | 2 +- .../common/workflows/envs/alfworld/RAFT_utils.py | 2 +- trinity/common/workflows/workflow.py | 8 ++++---- trinity/explorer/rollout_coordinator.py | 4 ++-- trinity/explorer/workflow_runner.py | 12 ++++++------ 6 files changed, 21 insertions(+), 22 deletions(-) diff --git a/tests/common/sglang_test.py b/tests/common/sglang_test.py index 7730452e5cf..a22cdeb46b7 100644 --- a/tests/common/sglang_test.py +++ b/tests/common/sglang_test.py @@ -367,7 +367,7 @@ async def test_chat_completions(self): class TestRecording(RayUnittestBaseAsync): - """Correctness of the in-SGLang generation recording flow (``enable_recording``). + """Correctness of the in-SGLang generation recording flow (``enable_history``). Mirrors ``tests/common/vllm_test.py::TestRecording``. Verifies that every call path lands its finished turn in the in-process ``MemoryStore`` under @@ -386,14 +386,14 @@ class TestRecording(RayUnittestBaseAsync): per-task ``record_key`` (captured by ``RecordingIdentityMiddleware``), matching vLLM (which sets no api_key auth in recording mode). - ``enable_recording`` forces ``enable_return_routed_experts`` in the + ``enable_history`` forces ``enable_return_routed_experts`` in the Allocator, so this test uses a MoE checkpoint (``get_moe_model_path``). """ async def asyncSetUp(self): self.config = get_template_config() self.config.mode = "explore" - # enable_recording forces enable_return_routed_experts -> needs a MoE + # enable_history forces enable_return_routed_experts -> needs a MoE # model (otherwise routed_experts is absent and the shape asserts below # would fail). Use a Qwen3-MoE checkpoint. self.config.model.model_path = get_moe_model_path() @@ -409,9 +409,9 @@ async def asyncSetUp(self): self.config.explorer.rollout_model.engine_num = 1 self.config.explorer.rollout_model.tensor_parallel_size = 2 self.config.explorer.rollout_model.chat_template = CHAT_TEMPLATE - # enable_recording requires the OpenAI API server (the recording runner). + # enable_history requires the OpenAI API server (the recording runner). self.config.explorer.rollout_model.enable_openai_api = True - self.config.explorer.rollout_model.enable_recording = True + self.config.explorer.rollout_model.enable_history = True self.config.explorer.rollout_model.enable_expert_parallel = True # Tool-call parsing coverage (qwen3_coder matches the Qwen3.5 chat # template). SGLang enables tool calling via tool_call_parser (no @@ -420,8 +420,7 @@ async def asyncSetUp(self): self.config.explorer.rollout_model.enable_auto_tool_choice = True self.config.explorer.rollout_model.tool_call_parser = "qwen3_coder" self.config.explorer.rollout_model.enable_thinking = False - # History recording is client-side; the in-SGLang recorder is the subject. - self.config.explorer.rollout_model.enable_history = False + # The in-SGLang recorder is the subject. self.config.explorer.rollout_model.base_port = 13400 self.config.check_and_update() @@ -502,7 +501,7 @@ def _assert_recorded_experience(self, exp: Experience, record_key: str): self.assertGreater(len(exp.response_text), 0) def _assert_recorded_routed_experts(self, exp: Experience): - # enable_return_routed_experts is forced on by enable_recording. + # enable_return_routed_experts is forced on by enable_history. self.assertIsNotNone(exp.routed_experts) re = exp.routed_experts self.assertEqual(re.dtype, torch.uint8) diff --git a/tests/common/vllm_test.py b/tests/common/vllm_test.py index 70263cbe757..0aa88c6c2a9 100644 --- a/tests/common/vllm_test.py +++ b/tests/common/vllm_test.py @@ -1949,7 +1949,7 @@ def _assert_recorded_experience(self, exp: Experience, record_key: str): self.assertGreater(len(exp.response_text), 0) def _assert_recorded_routed_experts(self, exp: Experience): - # enable_return_routed_experts is forced on by enable_recording. + # enable_return_routed_experts is forced on by enable_history. self.assertIsNotNone(exp.routed_experts) re = exp.routed_experts self.assertEqual(re.dtype, torch.uint8) diff --git a/trinity/common/workflows/envs/alfworld/RAFT_utils.py b/trinity/common/workflows/envs/alfworld/RAFT_utils.py index 8fe5c9f6b51..cbb4639becc 100644 --- a/trinity/common/workflows/envs/alfworld/RAFT_utils.py +++ b/trinity/common/workflows/envs/alfworld/RAFT_utils.py @@ -113,7 +113,7 @@ async def process_messages_to_experience_async(model, messages, info=None) -> Ex info = {} try: - # TODO(recording): when enable_recording is on, replace this + # TODO(recording): when enable_history is on, replace this # client-side conversion with a MemoryStore lookup by the session's # record_key (concatenate turns via info["sample_index"]); see # workflow.process_messages_to_experience. diff --git a/trinity/common/workflows/workflow.py b/trinity/common/workflows/workflow.py index 83e1d8dad81..ca4732878f6 100644 --- a/trinity/common/workflows/workflow.py +++ b/trinity/common/workflows/workflow.py @@ -95,7 +95,7 @@ class Workflow: # can join reward inside the store. Set by the WorkflowRunner from # ``config.explorer.rollout_model.enable_history``. See ``SimpleWorkflow`` # for the per-sample (n=1) loop this triggers. - enable_recording: bool = False + enable_history: bool = False def __init__( self, @@ -227,7 +227,7 @@ def _build_experience_from_converted( def process_messages_to_experience( self, messages, reward, info={}, truncate_status=None ) -> Experience: - # TODO(recording): when enable_recording is on, this client-side + # TODO(recording): when enable_history is on, this client-side # conversion is redundant — the vLLM recorder's build_experience already # captured the authoritative heavy data (real logprobs without an extra # forward, real routed_experts) into the MemoryStore, keyed by the @@ -350,7 +350,7 @@ def run(self) -> List[Experience]: messages = self.format_messages() self.logger.debug("start chat") - if self.enable_recording: + if self.enable_history: return self._run_recorded(messages) responses = self.model.chat(messages, **self.rollout_args) return self._attach_rewards(responses, base=self.run_id_base) @@ -379,7 +379,7 @@ async def run_async(self) -> List[Experience]: messages = self.format_messages() self.logger.info("start chat") - if self.enable_recording: + if self.enable_history: return await self._run_recorded_async(messages) responses = await self.model.chat_async(messages, **self.rollout_args) return self._attach_rewards(responses, base=self.run_id_base) diff --git a/trinity/explorer/rollout_coordinator.py b/trinity/explorer/rollout_coordinator.py index 922d6afd325..6667b3facba 100644 --- a/trinity/explorer/rollout_coordinator.py +++ b/trinity/explorer/rollout_coordinator.py @@ -76,7 +76,7 @@ def __init__( # recording path's per-rank /records/update_record fan-out. self._rank_urls: Optional[Dict[int, str]] = None - def _enable_recording(self) -> bool: + def _enable_history_recording(self) -> bool: """Whether the recording-consume path is active for train batches.""" return bool(self.config.explorer.rollout_model.enable_history) @@ -319,7 +319,7 @@ async def _finalize_train_batch( batch_state.state = BatchLifecycleState.FINALIZING try: - if self._enable_recording(): + if self._enable_history_recording(): pipeline_metrics = await self._consume_recorded_experiences(payload_chunks) else: pipeline_metrics = await self.process_experiences(payload_chunks) diff --git a/trinity/explorer/workflow_runner.py b/trinity/explorer/workflow_runner.py index 581f3a2d147..fe8519b7bfa 100644 --- a/trinity/explorer/workflow_runner.py +++ b/trinity/explorer/workflow_runner.py @@ -108,7 +108,7 @@ def _build_record_key(self, task: Task, run_index: int) -> str: return f"{task.batch_id}/{task.task_id}/{run_index}" def _set_record_key(self, model_wrapper: ModelWrapper, record_key: Optional[str]) -> None: - if self._enable_recording() and record_key is not None: + if self._enable_history_recording() and record_key is not None: model_wrapper.set_api_key(record_key) def _create_workflow_instance(self, task: Task, record_key: Optional[str] = None) -> Workflow: @@ -127,10 +127,10 @@ def _create_workflow_instance(self, task: Task, record_key: Optional[str] = None ) else: self.workflow_instance.reset(task) - self.workflow_instance.enable_recording = self._enable_recording() + self.workflow_instance.enable_history = self._enable_history_recording() return self.workflow_instance - def _enable_recording(self) -> bool: + def _enable_history_recording(self) -> bool: return bool(self.config.explorer.rollout_model.enable_history) async def _run_workflow(self, workflow_instance: Workflow) -> List[Experience]: @@ -145,7 +145,7 @@ def _create_isolated_workflow_instance( ) -> Tuple[Workflow, ModelWrapper]: model_wrapper = ( self.model_wrapper.clone_with_isolated_history() - if (self.config.explorer.rollout_model.enable_history or self._enable_recording()) + if self.config.explorer.rollout_model.enable_history else self.model_wrapper ) self._set_record_key(model_wrapper, record_key) @@ -153,7 +153,7 @@ def _create_isolated_workflow_instance( model_wrapper, self.auxiliary_model_wrappers, ) - wf.enable_recording = self._enable_recording() + wf.enable_history = self._enable_history_recording() return wf, model_wrapper def _build_execution_result( @@ -466,7 +466,7 @@ async def run_task( # (TODO: wire eval consume-and-discard so eval turns don't leak # in the store). For now, return no payload. return status, b"" - elif self._enable_recording(): + elif self._enable_history_recording(): # Recording path: ship only the small reward map keyed by the # per-sample record_key the workflow stamped on each exp. The # heavy experiences live in the vLLM MemoryStore and are pulled From 3122b55f418c957c94829520ad97812da53d0557 Mon Sep 17 00:00:00 2001 From: "panxuchen.pxc" Date: Fri, 26 Jun 2026 11:46:52 +0800 Subject: [PATCH 25/84] fix vllm test --- tests/common/vllm_test.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/tests/common/vllm_test.py b/tests/common/vllm_test.py index 0aa88c6c2a9..f39be3e2a00 100644 --- a/tests/common/vllm_test.py +++ b/tests/common/vllm_test.py @@ -1859,8 +1859,7 @@ async def asyncSetUp(self): self.config.explorer.rollout_model.enable_auto_tool_choice = True self.config.explorer.rollout_model.tool_call_parser = "qwen3_coder" self.config.explorer.rollout_model.enable_thinking = False - # History recording is client-side; the in-vLLM recorder is the subject. - self.config.explorer.rollout_model.enable_history = False + # The in-vLLM recorder is the subject. self.config.explorer.rollout_model.extra_engine_args = { "max_num_seqs": 24, "moe_backend": "triton", From 579d1890e32f7b8c479c000fc0e520156b1f0d52 Mon Sep 17 00:00:00 2001 From: pxc Date: Fri, 26 Jun 2026 16:42:04 +0800 Subject: [PATCH 26/84] fix tests --- trinity/common/config.py | 12 +++++++---- trinity/common/models/allocator.py | 14 +++++++------ trinity/common/models/vllm_model.py | 21 ++++++++++++++++--- .../models/vllm_patch/recording/server.py | 12 ++++++----- 4 files changed, 41 insertions(+), 18 deletions(-) diff --git a/trinity/common/config.py b/trinity/common/config.py index 92f89cf4bf9..98701d7fb25 100644 --- a/trinity/common/config.py +++ b/trinity/common/config.py @@ -576,10 +576,14 @@ class InferenceModelConfig: # the RolloutCoordinator pulls heavy experiences via ``/records/update_record`` # at finalize time. When off (default), runners ship serialized experiences # through the scheduler as before (legacy path). When True, the Allocator - # forces ``enable_return_routed_experts``. VLLMModel mirrors the recording - # config onto the engine instance for the recorder to read. The capture - # width (top-k logprobs) reuses ``logprobs`` below (default 1). Requires - # ``enable_openai_api=True`` (the recording runner is the API server). + # forces ``enable_openai_api`` (the recorder runs in the API server). + # VLLMModel mirrors the recording config onto the engine instance for the + # recorder to read. The capture width (top-k logprobs) reuses ``logprobs`` + # below (default 1). Requires ``enable_openai_api=True`` (the recording + # runner is the API server). Routed-experts capture is opt-in via + # ``enable_router_replay`` (mirrored to ``enable_return_routed_experts`` in + # ``config_validator``); it is not implied by ``enable_history``, so dense + # models can record history too. enable_history: bool = False # For OpenAI API diff --git a/trinity/common/models/allocator.py b/trinity/common/models/allocator.py index 089a9907669..3d53380698a 100644 --- a/trinity/common/models/allocator.py +++ b/trinity/common/models/allocator.py @@ -83,14 +83,16 @@ async def create_engine( config.engine_id = engine_id if config.engine_type.startswith("vllm") or config.engine_type == "sglang": - # enable_history is the single switch for engine-side recording. + # enable_history is the single switch for engine-side recording. The + # recorder runs inside the OpenAI API server, so it must be on. + # ``enable_return_routed_experts`` is NOT forced here: it is driven + # by the user's ``enable_router_replay`` (see ``config_validator``), + # so dense models can record history without vLLM's routed-experts + # capturer (which raises on configs lacking ``num_experts_per_tok``). + # The recorder simply leaves ``Experience.routed_experts`` as None + # when the engine did not capture any. if config.enable_history: config.enable_openai_api = True - # Recording: force routed_experts capture (the engine reads - # this at build time). The model mirrors the recording config - # onto the engine/server instance for the recorder to read. - # No env / runtime_env involved. - config.enable_return_routed_experts = True actor_bundle_lists = [] for node_id in range(config.nnodes): diff --git a/trinity/common/models/vllm_model.py b/trinity/common/models/vllm_model.py index ac0a8147fbe..f79997643ba 100644 --- a/trinity/common/models/vllm_model.py +++ b/trinity/common/models/vllm_model.py @@ -16,6 +16,10 @@ from trinity.common.models.mm_utils import vLLMMultiModalRender from trinity.common.models.model import BaseInferenceModel from trinity.common.models.recording.context import record_key_ctx, skip_recording_ctx +from trinity.common.models.recording.store import ( + RECORD_KEY_INFO_KEY, + REQUEST_ID_INFO_KEY, +) from trinity.common.models.vllm_patch import get_vllm_version from trinity.common.models.vllm_patch.recording.models import build_experience @@ -285,9 +289,20 @@ async def generate( returned_seq, is_valid = self._handle_prompt_truncation(prompt, **kwargs) # type: ignore if not is_valid: - return ( - returned_seq # is_valid is False: returned_seq is a list of dummy experiences - ) + # Prompt was truncated: ``_handle_prompt_truncation`` returns + # dummy (masked) experiences and we skip real generation. The + # engine-level recorder only captures actual generations, so + # persist these dummies directly under the record_key — masked + # experiences must still be tracked for history extraction and + # the buffer/trainer (they are popped by record_key on consume). + if self.recorder is not None and record_key is not None: + for exp in returned_seq: + exp.info[RECORD_KEY_INFO_KEY] = record_key + exp.info[REQUEST_ID_INFO_KEY] = exp.eid.suffix + exp.info["rank"] = self.recorder.rank + exp.info["model_version"] = self.model_version + await self.recorder.store.append_turn(exp) + return returned_seq prompt = { "prompt_token_ids": returned_seq } # is_valid is True: returned_seq is token_ids diff --git a/trinity/common/models/vllm_patch/recording/server.py b/trinity/common/models/vllm_patch/recording/server.py index b58706c26e8..4e8f4446b09 100644 --- a/trinity/common/models/vllm_patch/recording/server.py +++ b/trinity/common/models/vllm_patch/recording/server.py @@ -260,11 +260,13 @@ async def run_api_server_with_recording( # path happens in VLLMModel (via EngineArgs), *before* this runner gets the # already-built ``async_llm``. Adding ``--enable-return-routed-experts`` / # ``--max-logprobs`` here would be inert (init_app_state does not read them). - # The Allocator therefore forces ``InferenceModelConfig.enable_return_routed_experts - # = True`` when recording is on, and the engine's default ``max_logprobs=20`` - # covers the recorder's top-k (``VLLM_RECORD_TOPK``, default 1). To record - # routed_experts, the engine must be built with that flag on — the launcher - # is responsible for that, not these CLI args. + # ``enable_return_routed_experts`` is opt-in via ``enable_router_replay`` + # (mirrored in ``config_validator``); it is not implied by recording, so a + # dense model records history with ``routed_experts=None``. The engine's + # default ``max_logprobs=20`` covers the recorder's top-k + # (``VLLM_RECORD_TOPK``, default 1). To record routed_experts, the engine + # must be built with that flag on — the launcher is responsible for that, + # not these CLI args. args = parser.parse_args(cli_args) args.structured_outputs_config.reasoning_parser = reasoning_parser From 052562a27cf2151454e4e76af7a4fb209ea94eae Mon Sep 17 00:00:00 2001 From: "panxuchen.pxc" Date: Fri, 26 Jun 2026 19:09:47 +0800 Subject: [PATCH 27/84] add store --- trinity/buffer/store/__init__.py | 8 ++ trinity/buffer/store/base_store.py | 34 ++++++ trinity/buffer/store/memory_store.py | 126 +++++++++++++++++++++++ trinity/common/models/recording/store.py | 62 ++++++----- 4 files changed, 204 insertions(+), 26 deletions(-) create mode 100644 trinity/buffer/store/__init__.py create mode 100644 trinity/buffer/store/base_store.py create mode 100644 trinity/buffer/store/memory_store.py diff --git a/trinity/buffer/store/__init__.py b/trinity/buffer/store/__init__.py new file mode 100644 index 00000000000..7c8cd86a076 --- /dev/null +++ b/trinity/buffer/store/__init__.py @@ -0,0 +1,8 @@ +from trinity.buffer.store.base_store import BaseStore +from trinity.buffer.store.memory_store import MemoryStore, default_sample_id_getter + +__all__ = [ + "BaseStore", + "MemoryStore", + "default_sample_id_getter", +] diff --git a/trinity/buffer/store/base_store.py b/trinity/buffer/store/base_store.py new file mode 100644 index 00000000000..a55e1b51f82 --- /dev/null +++ b/trinity/buffer/store/base_store.py @@ -0,0 +1,34 @@ +from abc import ABC, abstractmethod +from typing import List + +from trinity.common.experience import Experience + + +class BaseStore(ABC): + """Abstract base class for an in-process experience store. + + The key follows the format ``//`` and each + experience is associated with a unique sample id. + """ + + @abstractmethod + def add(self, key: str, exps: List[Experience]) -> None: + """Add experiences to the store under the given complete key.""" + + @abstractmethod + def overwrite(self, key: str, exps: List[Experience]) -> None: + """Replace all experiences under the given complete key.""" + + @abstractmethod + def update( + self, key: str, reward: float, info: dict | None, sample_ids: List[str] | None + ) -> None: + """Update reward and optional info for selected experiences under a complete key.""" + + @abstractmethod + def get(self, key: str) -> List[Experience]: + """Return experiences for an exact key or prefix without removing them.""" + + @abstractmethod + def remove(self, key: str) -> List[Experience]: + """Remove and return experiences for an exact key or prefix.""" diff --git a/trinity/buffer/store/memory_store.py b/trinity/buffer/store/memory_store.py new file mode 100644 index 00000000000..be9a9b4e0c9 --- /dev/null +++ b/trinity/buffer/store/memory_store.py @@ -0,0 +1,126 @@ +"""In-memory implementation of the experience store interface.""" + +from collections import OrderedDict +from typing import Callable, Iterable, List + +from trinity.buffer.store.base_store import BaseStore +from trinity.common.experience import Experience + +SampleIdGetter = Callable[[Experience], str] + + +def default_sample_id_getter(exp: Experience) -> str: + """Resolve a stable sample id for an experience.""" + info = exp.info or {} + sample_id = info.get("sample_id") + if sample_id is not None: + return str(sample_id) + + request_id = info.get("request_id") + sample_index = info.get("sample_index") + if request_id is not None and sample_index is not None: + return f"{request_id}:{sample_index}" + if request_id is not None: + return str(request_id) + + return exp.eid.uid + + +class MemoryStore(BaseStore): + """A fast in-process store backed by Python dictionaries. + + ``add``, ``overwrite`` and ``update`` require complete keys in the form + ``//``. ``get`` and ``remove`` also accept prefixes + so callers can drain a task or step at once. + """ + + def __init__(self, sample_id_getter: SampleIdGetter | None = None) -> None: + self.sample_id_getter = sample_id_getter or default_sample_id_getter + self._records: dict[str, OrderedDict[str, Experience]] = {} + self._sample_to_key: dict[str, str] = {} + + def __len__(self) -> int: + return sum(len(exps) for exps in self._records.values()) + + def add(self, key: str, exps: List[Experience]) -> None: + self._validate_complete_key(key) + if not exps: + return + + records = self._records.setdefault(key, OrderedDict()) + for exp in exps: + sample_id = self.sample_id_getter(exp) + owner_key = self._sample_to_key.get(sample_id) + if owner_key is not None: + raise ValueError( + f"Duplicate sample_id '{sample_id}' already exists under key '{owner_key}'." + ) + records[sample_id] = exp + self._sample_to_key[sample_id] = key + + def overwrite(self, key: str, exps: List[Experience]) -> None: + self._validate_complete_key(key) + self._drop_key(key) + self.add(key, exps) + + def update( + self, + key: str, + reward: float, + info: dict | None, + sample_ids: List[str] | None, + ) -> None: + self._validate_complete_key(key) + records = self._records.get(key) + if records is None: + raise KeyError(f"Key '{key}' does not exist.") + + target_ids: Iterable[str] = list(records.keys()) if sample_ids is None else sample_ids + for sample_id in target_ids: + if sample_id not in records: + raise KeyError(f"sample_id '{sample_id}' does not exist under key '{key}'.") + exp = records[sample_id] + exp.reward = reward + if info: + if exp.info is None: + exp.info = {} + exp.info.update(info) + + def get(self, key: str) -> List[Experience]: + result: List[Experience] = [] + for matched_key in self._matching_keys(key): + result.extend(self._records[matched_key].values()) + return result + + def remove(self, key: str) -> List[Experience]: + result: List[Experience] = [] + for matched_key in self._matching_keys(key): + result.extend(self._drop_key(matched_key)) + return result + + def keys(self) -> list[str]: + return list(self._records.keys()) + + @staticmethod + def _validate_complete_key(key: str) -> None: + parts = key.split("/") + if len(parts) != 3 or any(part == "" for part in parts): + raise ValueError( + f"Store key must be complete '//', got '{key}'." + ) + + def _matching_keys(self, key: str) -> list[str]: + if key == "": + return list(self._records.keys()) + if key in self._records: + return [key] + prefix = key + "/" + return [record_key for record_key in self._records if record_key.startswith(prefix)] + + def _drop_key(self, key: str) -> list[Experience]: + records = self._records.pop(key, None) + if records is None: + return [] + for sample_id in records: + self._sample_to_key.pop(sample_id, None) + return list(records.values()) diff --git a/trinity/common/models/recording/store.py b/trinity/common/models/recording/store.py index 3851e3d0959..1eaf3e7c936 100644 --- a/trinity/common/models/recording/store.py +++ b/trinity/common/models/recording/store.py @@ -1,9 +1,9 @@ """Pluggable storage backends for recorded experiences.""" import abc -from collections import defaultdict from typing import Optional +from trinity.buffer.store import MemoryStore as BaseMemoryStore from trinity.common.experience import Experience REQUEST_ID_INFO_KEY = "request_id" @@ -46,55 +46,65 @@ async def delete_request_experience(self, record_key: str, request_id: str) -> b """Drop one experience by request id. Return True if one was deleted.""" -class MemoryStore(RecordStore): +class MemoryStore(BaseMemoryStore, RecordStore): """In-process store grouped by recording identity.""" def __init__(self) -> None: - self._records: dict[str, list[Experience]] = defaultdict(list) + super().__init__() @staticmethod def _group_key(exp: Experience) -> str: - record_key = exp.info.get(RECORD_KEY_INFO_KEY) + info = exp.info or {} + record_key = info.get(RECORD_KEY_INFO_KEY) return record_key if record_key else exp.eid.suffix async def append_turn(self, exp: Experience) -> None: - self._records[self._group_key(exp)].append(exp) + self.add(self._group_key(exp), [exp]) async def update_reward_by_record_key( self, record_key: str, reward: float, run: int, task: str ) -> list[Experience]: - exps = self._records.pop(record_key, []) - for exp in exps: - exp.reward = reward - exp.eid.run = run - exp.eid.task = task - return exps + if not self.get(record_key): + return [] + self.update( + key=record_key, + reward=reward, + info={"run": run, "task": task}, + sample_ids=None, + ) + return self.remove(record_key) async def get_record_experiences(self, record_key: str) -> list[Experience]: - return list(self._records.get(record_key, [])) + return self.get(record_key) async def get_request_experience( self, record_key: str, request_id: str ) -> Optional[Experience]: - for exp in self._records.get(record_key, []): - if exp.info.get(REQUEST_ID_INFO_KEY) == request_id: + for exp in self.get(record_key): + info = exp.info or {} + if info.get(REQUEST_ID_INFO_KEY) == request_id: return exp return None async def list_records(self) -> list[str]: - return list(self._records.keys()) + return self.keys() async def delete_record_experiences(self, record_key: str) -> None: - self._records.pop(record_key, None) + self.remove(record_key) async def delete_request_experience(self, record_key: str, request_id: str) -> bool: - exps = self._records.get(record_key) - if not exps: - return False - for index, exp in enumerate(exps): - if exp.info.get(REQUEST_ID_INFO_KEY) == request_id: - del exps[index] - if not exps: - self._records.pop(record_key, None) - return True - return False + kept = [] + deleted = False + for exp in self.get(record_key): + info = exp.info or {} + if info.get(REQUEST_ID_INFO_KEY) == request_id: + deleted = True + else: + kept.append(exp) + + if deleted: + if kept: + self.overwrite(record_key, kept) + else: + await self.delete_record_experiences(record_key) + return deleted From d8f3de633d4596ab59982f311b7f807a611da35d Mon Sep 17 00:00:00 2001 From: "panxuchen.pxc" Date: Fri, 26 Jun 2026 19:23:05 +0800 Subject: [PATCH 28/84] refactor store --- trinity/buffer/store/__init__.py | 14 ++- trinity/buffer/store/base_store.py | 11 +- trinity/buffer/store/memory_store.py | 37 +++++- trinity/common/models/model.py | 4 +- trinity/common/models/recording/__init__.py | 7 +- trinity/common/models/recording/query.py | 51 ++++---- trinity/common/models/recording/recorder.py | 6 +- trinity/common/models/recording/store.py | 110 ------------------ trinity/common/models/sglang_model.py | 6 +- .../models/sglang_patch/recording/__init__.py | 11 +- .../models/sglang_patch/recording/models.py | 5 +- .../models/sglang_patch/recording/recorder.py | 3 +- .../models/sglang_patch/recording/server.py | 3 +- trinity/common/models/tinker_model.py | 2 +- trinity/common/models/vllm_model.py | 7 +- .../models/vllm_patch/recording/__init__.py | 7 +- .../models/vllm_patch/recording/models.py | 13 +-- .../models/vllm_patch/recording/recorder.py | 3 +- trinity/common/workflows/workflow.py | 14 ++- 19 files changed, 131 insertions(+), 183 deletions(-) delete mode 100644 trinity/common/models/recording/store.py diff --git a/trinity/buffer/store/__init__.py b/trinity/buffer/store/__init__.py index 7c8cd86a076..965388b2d65 100644 --- a/trinity/buffer/store/__init__.py +++ b/trinity/buffer/store/__init__.py @@ -1,8 +1,18 @@ -from trinity.buffer.store.base_store import BaseStore -from trinity.buffer.store.memory_store import MemoryStore, default_sample_id_getter +from trinity.buffer.store.base_store import BaseStore, RecordStore +from trinity.buffer.store.memory_store import ( + RECORD_KEY_INFO_KEY, + REQUEST_ID_INFO_KEY, + MemoryStore, + default_sample_id_getter, + get_record_key, +) __all__ = [ "BaseStore", "MemoryStore", + "RECORD_KEY_INFO_KEY", + "REQUEST_ID_INFO_KEY", + "RecordStore", "default_sample_id_getter", + "get_record_key", ] diff --git a/trinity/buffer/store/base_store.py b/trinity/buffer/store/base_store.py index a55e1b51f82..84e5a655f27 100644 --- a/trinity/buffer/store/base_store.py +++ b/trinity/buffer/store/base_store.py @@ -7,7 +7,7 @@ class BaseStore(ABC): """Abstract base class for an in-process experience store. - The key follows the format ``//`` and each + The key follows the format ``//`` and each experience is associated with a unique sample id. """ @@ -23,7 +23,7 @@ def overwrite(self, key: str, exps: List[Experience]) -> None: def update( self, key: str, reward: float, info: dict | None, sample_ids: List[str] | None ) -> None: - """Update reward and optional info for selected experiences under a complete key.""" + """Update reward, EID fields from key, and optional info for selected experiences.""" @abstractmethod def get(self, key: str) -> List[Experience]: @@ -32,3 +32,10 @@ def get(self, key: str) -> List[Experience]: @abstractmethod def remove(self, key: str) -> List[Experience]: """Remove and return experiences for an exact key or prefix.""" + + @abstractmethod + def keys(self) -> list[str]: + """Return complete keys currently stored in insertion order.""" + + +RecordStore = BaseStore diff --git a/trinity/buffer/store/memory_store.py b/trinity/buffer/store/memory_store.py index be9a9b4e0c9..905b4b5739d 100644 --- a/trinity/buffer/store/memory_store.py +++ b/trinity/buffer/store/memory_store.py @@ -6,6 +6,9 @@ from trinity.buffer.store.base_store import BaseStore from trinity.common.experience import Experience +REQUEST_ID_INFO_KEY = "request_id" +RECORD_KEY_INFO_KEY = "record_key" + SampleIdGetter = Callable[[Experience], str] @@ -26,12 +29,23 @@ def default_sample_id_getter(exp: Experience) -> str: return exp.eid.uid +def get_record_key(exp: Experience) -> str: + """Return the complete store key stamped on an experience.""" + info = exp.info or {} + record_key = info.get(RECORD_KEY_INFO_KEY) + if record_key: + return record_key + if exp.eid.batch != "" and exp.eid.task != "": + return exp.eid.rid + return exp.eid.suffix + + class MemoryStore(BaseStore): """A fast in-process store backed by Python dictionaries. ``add``, ``overwrite`` and ``update`` require complete keys in the form - ``//``. ``get`` and ``remove`` also accept prefixes - so callers can drain a task or step at once. + ``//``. ``get`` and ``remove`` also accept prefixes + so callers can drain a batch or task at once. """ def __init__(self, sample_id_getter: SampleIdGetter | None = None) -> None: @@ -75,11 +89,15 @@ def update( if records is None: raise KeyError(f"Key '{key}' does not exist.") + batch, task, run = self._parse_complete_key(key) target_ids: Iterable[str] = list(records.keys()) if sample_ids is None else sample_ids for sample_id in target_ids: if sample_id not in records: raise KeyError(f"sample_id '{sample_id}' does not exist under key '{key}'.") exp = records[sample_id] + exp.eid.batch = batch + exp.eid.task = task + exp.eid.run = run exp.reward = reward if info: if exp.info is None: @@ -103,11 +121,24 @@ def keys(self) -> list[str]: @staticmethod def _validate_complete_key(key: str) -> None: + MemoryStore._parse_complete_key(key) + + @staticmethod + def _parse_complete_key(key: str) -> tuple[str, str, int]: parts = key.split("/") if len(parts) != 3 or any(part == "" for part in parts): raise ValueError( - f"Store key must be complete '//', got '{key}'." + f"Store key must be complete '//', got '{key}'." ) + batch, task, run_text = parts + try: + run = int(run_text) + except ValueError as exc: + raise ValueError( + f"Store key run_id must be an integer in '//', " + f"got '{key}'." + ) from exc + return batch, task, run def _matching_keys(self, key: str) -> list[str]: if key == "": diff --git a/trinity/common/models/model.py b/trinity/common/models/model.py index faf284973e7..4daec467f76 100644 --- a/trinity/common/models/model.py +++ b/trinity/common/models/model.py @@ -155,9 +155,9 @@ async def extract_experience_from_history( if recorder is None: raise ValueError("Recording is not enabled for this model.") await recorder.flush() - exps = await recorder.store.get_record_experiences(record_key) + exps = recorder.store.get(record_key) if clear_history: - await recorder.store.delete_record_experiences(record_key) + recorder.store.remove(record_key) return exps def get_model_config(self) -> InferenceModelConfig: diff --git a/trinity/common/models/recording/__init__.py b/trinity/common/models/recording/__init__.py index 6b585dd2648..1e1e9bd515c 100644 --- a/trinity/common/models/recording/__init__.py +++ b/trinity/common/models/recording/__init__.py @@ -1,5 +1,6 @@ """Engine-agnostic generation recording utilities.""" +from trinity.buffer.store import MemoryStore, RecordStore from trinity.common.models.recording.context import ( RecordingIdentityMiddleware, extract_bearer_token, @@ -12,11 +13,7 @@ TRINITY_RECORDER_ATTR, Recorder, ) -from trinity.common.models.recording.server import ( - add_recording_middleware, - mount_recording_api, -) -from trinity.common.models.recording.store import MemoryStore, RecordStore +from trinity.common.models.recording.server import add_recording_middleware, mount_recording_api __all__ = [ "MemoryStore", diff --git a/trinity/common/models/recording/query.py b/trinity/common/models/recording/query.py index 98dc33dbdba..4131c24b93c 100644 --- a/trinity/common/models/recording/query.py +++ b/trinity/common/models/recording/query.py @@ -5,13 +5,13 @@ from fastapi import APIRouter, HTTPException, Request, Response from pydantic import BaseModel +from trinity.buffer.store import REQUEST_ID_INFO_KEY, RecordStore from trinity.common.experience import Experience from trinity.common.models.recording.recorder import ( TRINITY_RECORD_STORE_ATTR, TRINITY_RECORDER_ATTR, Recorder, ) -from trinity.common.models.recording.store import RecordStore STORE_STATE_ATTR = TRINITY_RECORD_STORE_ATTR RECORDER_STATE_ATTR = TRINITY_RECORDER_ATTR @@ -22,8 +22,6 @@ class _RecordUpdate(BaseModel): record_key: str reward: float - run: int = 0 - task: str = "" class _UpdateRecordRequest(BaseModel): @@ -44,30 +42,31 @@ def _recorder(request: Request) -> Recorder: return rec -async def _get_exp(store: RecordStore, record_key: str, request_id: str) -> Experience: - exp = await store.get_request_experience(record_key, request_id) - if exp is None: - raise HTTPException(status_code=404, detail="experience not found") - return exp +def _get_exp(store: RecordStore, record_key: str, request_id: str) -> Experience: + for exp in store.get(record_key): + info = exp.info or {} + if info.get(REQUEST_ID_INFO_KEY) == request_id: + return exp + raise HTTPException(status_code=404, detail="experience not found") @query_router.get("") async def list_records(request: Request) -> dict: store = _store(request) - return {"record_keys": await store.list_records()} + return {"record_keys": store.keys()} @query_router.get("/{record_key}") async def get_record_experiences(record_key: str, request: Request) -> dict: store = _store(request) - experiences = await store.get_record_experiences(record_key) + experiences = store.get(record_key) return {"record_key": record_key, "experiences": [e.to_dict() for e in experiences]} @query_router.get("/{record_key}/request/{request_id}") async def get_request_experience(record_key: str, request_id: str, request: Request) -> Response: store = _store(request) - exp = await _get_exp(store, record_key, request_id) + exp = _get_exp(store, record_key, request_id) return Response( content=Experience.serialize(exp), media_type="application/octet-stream", @@ -78,16 +77,27 @@ async def get_request_experience(record_key: str, request_id: str, request: Requ @query_router.delete("/{record_key}") async def delete_record_experiences(record_key: str, request: Request) -> dict: store = _store(request) - await store.delete_record_experiences(record_key) + store.remove(record_key) return {"record_key": record_key, "deleted": True} @query_router.delete("/{record_key}/request/{request_id}") async def delete_request_experience(record_key: str, request_id: str, request: Request) -> dict: store = _store(request) - deleted = await store.delete_request_experience(record_key, request_id) + kept = [] + deleted = False + for exp in store.get(record_key): + info = exp.info or {} + if info.get(REQUEST_ID_INFO_KEY) == request_id: + deleted = True + else: + kept.append(exp) if not deleted: raise HTTPException(status_code=404, detail="experience not found") + if kept: + store.overwrite(record_key, kept) + else: + store.remove(record_key) return {"record_key": record_key, "request_id": request_id, "deleted": True} @@ -99,14 +109,15 @@ async def update_record(req: _UpdateRecordRequest, request: Request) -> Response exps: List[Experience] = [] for update in req.updates: - exps.extend( - await store.update_reward_by_record_key( - record_key=update.record_key, - reward=update.reward, - run=update.run, - task=update.task, - ) + if not store.get(update.record_key): + continue + store.update( + key=update.record_key, + reward=update.reward, + info=None, + sample_ids=None, ) + exps.extend(store.remove(update.record_key)) return Response( content=Experience.serialize_many(exps), media_type="application/octet-stream", diff --git a/trinity/common/models/recording/recorder.py b/trinity/common/models/recording/recorder.py index fae8caa378e..d9bbbc37ae7 100644 --- a/trinity/common/models/recording/recorder.py +++ b/trinity/common/models/recording/recorder.py @@ -6,9 +6,9 @@ from datetime import datetime, timezone from typing import Any, Optional +from trinity.buffer.store import RecordStore, get_record_key from trinity.common.experience import Experience from trinity.common.models.recording.context import skip_recording_ctx -from trinity.common.models.recording.store import RecordStore MODEL_VERSION_ATTR = "trinity_model_version" TRINITY_RECORDER_ATTR = "trinity_recorder" @@ -97,9 +97,9 @@ async def _flush_loop(self) -> None: async def _safe_append(self, exp: Experience) -> None: try: - await self.store.append_turn(exp) + self.store.add(get_record_key(exp), [exp]) except Exception: logging.getLogger(__name__).exception( - "recording store.append_turn failed for request %s", + "recording store.add failed for request %s", exp.info.get("request_id"), ) diff --git a/trinity/common/models/recording/store.py b/trinity/common/models/recording/store.py deleted file mode 100644 index 1eaf3e7c936..00000000000 --- a/trinity/common/models/recording/store.py +++ /dev/null @@ -1,110 +0,0 @@ -"""Pluggable storage backends for recorded experiences.""" - -import abc -from typing import Optional - -from trinity.buffer.store import MemoryStore as BaseMemoryStore -from trinity.common.experience import Experience - -REQUEST_ID_INFO_KEY = "request_id" -RECORD_KEY_INFO_KEY = "record_key" - - -class RecordStore(abc.ABC): - """Abstract persistence interface for recorded experiences.""" - - @abc.abstractmethod - async def append_turn(self, exp: Experience) -> None: - """Persist one completed experience.""" - - @abc.abstractmethod - async def update_reward_by_record_key( - self, record_key: str, reward: float, run: int, task: str - ) -> list[Experience]: - """Set reward/run/task on every experience in the group, pop and return it.""" - - @abc.abstractmethod - async def get_record_experiences(self, record_key: str) -> list[Experience]: - """Return all experiences for a record key, in insertion order.""" - - @abc.abstractmethod - async def get_request_experience( - self, record_key: str, request_id: str - ) -> Optional[Experience]: - """Return a single experience, or None if not found.""" - - @abc.abstractmethod - async def list_records(self) -> list[str]: - """Return all known record keys.""" - - @abc.abstractmethod - async def delete_record_experiences(self, record_key: str) -> None: - """Drop all experiences for a record key.""" - - @abc.abstractmethod - async def delete_request_experience(self, record_key: str, request_id: str) -> bool: - """Drop one experience by request id. Return True if one was deleted.""" - - -class MemoryStore(BaseMemoryStore, RecordStore): - """In-process store grouped by recording identity.""" - - def __init__(self) -> None: - super().__init__() - - @staticmethod - def _group_key(exp: Experience) -> str: - info = exp.info or {} - record_key = info.get(RECORD_KEY_INFO_KEY) - return record_key if record_key else exp.eid.suffix - - async def append_turn(self, exp: Experience) -> None: - self.add(self._group_key(exp), [exp]) - - async def update_reward_by_record_key( - self, record_key: str, reward: float, run: int, task: str - ) -> list[Experience]: - if not self.get(record_key): - return [] - self.update( - key=record_key, - reward=reward, - info={"run": run, "task": task}, - sample_ids=None, - ) - return self.remove(record_key) - - async def get_record_experiences(self, record_key: str) -> list[Experience]: - return self.get(record_key) - - async def get_request_experience( - self, record_key: str, request_id: str - ) -> Optional[Experience]: - for exp in self.get(record_key): - info = exp.info or {} - if info.get(REQUEST_ID_INFO_KEY) == request_id: - return exp - return None - - async def list_records(self) -> list[str]: - return self.keys() - - async def delete_record_experiences(self, record_key: str) -> None: - self.remove(record_key) - - async def delete_request_experience(self, record_key: str, request_id: str) -> bool: - kept = [] - deleted = False - for exp in self.get(record_key): - info = exp.info or {} - if info.get(REQUEST_ID_INFO_KEY) == request_id: - deleted = True - else: - kept.append(exp) - - if deleted: - if kept: - self.overwrite(record_key, kept) - else: - await self.delete_record_experiences(record_key) - return deleted diff --git a/trinity/common/models/sglang_model.py b/trinity/common/models/sglang_model.py index 45e213504fa..03027dac02f 100644 --- a/trinity/common/models/sglang_model.py +++ b/trinity/common/models/sglang_model.py @@ -568,11 +568,9 @@ async def run_api_server(self) -> bool: recorder = None routed_experts_layout = None if self.config.enable_history: + from trinity.buffer.store import MemoryStore from trinity.common.models.recording.recorder import Recorder - from trinity.common.models.recording.store import MemoryStore - from trinity.common.models.sglang_patch.recording.models import ( - build_sglang_experience, - ) + from trinity.common.models.sglang_patch.recording.models import build_sglang_experience record_store = MemoryStore() recorder = Recorder( diff --git a/trinity/common/models/sglang_patch/recording/__init__.py b/trinity/common/models/sglang_patch/recording/__init__.py index 260314194c1..062e1bcdcf9 100644 --- a/trinity/common/models/sglang_patch/recording/__init__.py +++ b/trinity/common/models/sglang_patch/recording/__init__.py @@ -5,6 +5,8 @@ core symbols shared with the vLLM recording path (``trinity.common.models.recording``). """ + +from trinity.buffer.store import MemoryStore, RecordStore # noqa: F401 from trinity.common.models.recording.context import ( # noqa: F401 RecordingIdentityMiddleware, get_recording_record_key, @@ -13,17 +15,14 @@ ) from trinity.common.models.recording.query import query_router # noqa: F401 from trinity.common.models.recording.recorder import Recorder # noqa: F401 -from trinity.common.models.recording.store import MemoryStore, RecordStore # noqa: F401 -from trinity.common.models.sglang_patch.recording.models import ( # noqa: F401 +from trinity.common.models.sglang_patch.recording.models import ( build_sglang_experience, -) +) # noqa: F401 from trinity.common.models.sglang_patch.recording.recorder import ( # noqa: F401 create_sglang_recorder, patch_tokenizer_manager_for_recording, ) -from trinity.common.models.sglang_patch.recording.server import ( # noqa: F401 - setup_sglang_recording, -) +from trinity.common.models.sglang_patch.recording.server import setup_sglang_recording # noqa: F401 __all__ = [ "MemoryStore", diff --git a/trinity/common/models/sglang_patch/recording/models.py b/trinity/common/models/sglang_patch/recording/models.py index 8478a48941e..88093f192ff 100644 --- a/trinity/common/models/sglang_patch/recording/models.py +++ b/trinity/common/models/sglang_patch/recording/models.py @@ -12,8 +12,8 @@ from the request. Field mapping (SGLang ``ret`` -> ``Experience``): - meta_info.id -> eid.suffix (traceability; task/run/reward assigned by - ``MemoryStore.update_reward_by_record_key`` at consume) + meta_info.id -> eid.suffix (traceability; batch/task/run and reward + are assigned from record key by ``MemoryStore.update``) record_key -> info["record_key"] (the MemoryStore group key) sample index -> info["sample_index"] (position within the n set) prompt_token_ids -> tokens (prompt) + prompt_length @@ -24,6 +24,7 @@ the model's ``(num_layers, topk)`` layout when base64-str) meta_info.weight_version -> info["model_version"] """ + from typing import Any, List, Optional, Tuple import torch diff --git a/trinity/common/models/sglang_patch/recording/recorder.py b/trinity/common/models/sglang_patch/recording/recorder.py index 23253e3bc80..747fcffb82f 100644 --- a/trinity/common/models/sglang_patch/recording/recorder.py +++ b/trinity/common/models/sglang_patch/recording/recorder.py @@ -23,17 +23,18 @@ ``prompt_token_ids`` on the separate ``return_token_ids`` flag) and only adds an ignored field to ``/generate`` JSON responses. """ + import functools import logging from typing import Any, List, Optional, Tuple +from trinity.buffer.store import MemoryStore, RecordStore from trinity.common.models.recording.context import record_key_ctx from trinity.common.models.recording.recorder import ( TRINITY_RECORD_STORE_ATTR, TRINITY_RECORDER_ATTR, Recorder, ) -from trinity.common.models.recording.store import MemoryStore, RecordStore from trinity.common.models.sglang_patch.recording.models import build_sglang_experience #: Guard attribute marking the wrapped generate_request (mirrors vLLM's style). diff --git a/trinity/common/models/sglang_patch/recording/server.py b/trinity/common/models/sglang_patch/recording/server.py index 54a2e711ca7..704d2cd7133 100644 --- a/trinity/common/models/sglang_patch/recording/server.py +++ b/trinity/common/models/sglang_patch/recording/server.py @@ -12,12 +12,13 @@ via ``extract_experience_from_history``; they are also stashed on ``app.state`` for the ``query_router`` HTTP drain path used by the coordinator. """ + import logging from typing import Optional, Tuple +from trinity.buffer.store import RecordStore from trinity.common.models.recording.recorder import Recorder from trinity.common.models.recording.server import mount_recording_api -from trinity.common.models.recording.store import RecordStore from trinity.common.models.sglang_patch.recording.recorder import create_sglang_recorder diff --git a/trinity/common/models/tinker_model.py b/trinity/common/models/tinker_model.py index 2228fa7a6e3..113d1dda91c 100644 --- a/trinity/common/models/tinker_model.py +++ b/trinity/common/models/tinker_model.py @@ -8,12 +8,12 @@ from tinker import types from torch import Tensor +from trinity.buffer.store import MemoryStore from trinity.common.config import InferenceModelConfig from trinity.common.constants import SyncMethod from trinity.common.experience import Experience from trinity.common.models.model import BaseInferenceModel from trinity.common.models.recording.recorder import MODEL_VERSION_ATTR, Recorder -from trinity.common.models.recording.store import MemoryStore from trinity.manager.synchronizer import Synchronizer diff --git a/trinity/common/models/vllm_model.py b/trinity/common/models/vllm_model.py index f79997643ba..58113522142 100644 --- a/trinity/common/models/vllm_model.py +++ b/trinity/common/models/vllm_model.py @@ -10,16 +10,13 @@ from packaging.version import parse as parse_version from transformers import AutoProcessor +from trinity.buffer.store import RECORD_KEY_INFO_KEY, REQUEST_ID_INFO_KEY from trinity.common.config import InferenceModelConfig from trinity.common.constants import SyncMethod from trinity.common.experience import Experience from trinity.common.models.mm_utils import vLLMMultiModalRender from trinity.common.models.model import BaseInferenceModel from trinity.common.models.recording.context import record_key_ctx, skip_recording_ctx -from trinity.common.models.recording.store import ( - RECORD_KEY_INFO_KEY, - REQUEST_ID_INFO_KEY, -) from trinity.common.models.vllm_patch import get_vllm_version from trinity.common.models.vllm_patch.recording.models import build_experience @@ -301,7 +298,7 @@ async def generate( exp.info[REQUEST_ID_INFO_KEY] = exp.eid.suffix exp.info["rank"] = self.recorder.rank exp.info["model_version"] = self.model_version - await self.recorder.store.append_turn(exp) + self.recorder.store.add(record_key, [exp]) return returned_seq prompt = { "prompt_token_ids": returned_seq diff --git a/trinity/common/models/vllm_patch/recording/__init__.py b/trinity/common/models/vllm_patch/recording/__init__.py index 32eef1fc0dd..10aa385f327 100644 --- a/trinity/common/models/vllm_patch/recording/__init__.py +++ b/trinity/common/models/vllm_patch/recording/__init__.py @@ -1,6 +1,8 @@ """Generation-recording patch for the vLLM OpenAI server. Designed for vllm >= 0.23.0. """ + +from trinity.buffer.store import MemoryStore, RecordStore from trinity.common.models.recording.context import ( RecordingIdentityMiddleware, record_key_ctx, @@ -8,15 +10,12 @@ ) from trinity.common.models.recording.query import query_router from trinity.common.models.recording.recorder import Recorder -from trinity.common.models.recording.store import MemoryStore, RecordStore from trinity.common.models.vllm_patch.recording.models import build_experience from trinity.common.models.vllm_patch.recording.recorder import ( create_vllm_recorder, patch_engine_for_recording, ) -from trinity.common.models.vllm_patch.recording.server import ( - run_api_server_with_recording, -) +from trinity.common.models.vllm_patch.recording.server import run_api_server_with_recording __all__ = [ "MemoryStore", diff --git a/trinity/common/models/vllm_patch/recording/models.py b/trinity/common/models/vllm_patch/recording/models.py index 2cad37febcd..fe4a46b365b 100644 --- a/trinity/common/models/vllm_patch/recording/models.py +++ b/trinity/common/models/vllm_patch/recording/models.py @@ -11,14 +11,12 @@ Field mapping (captured ``RequestOutput`` fields -> ``Experience``): request_id -> eid.suffix (``EID(suffix=...)``; the vLLM engine request id == the OpenAI ``response.id``. Kept for traceability; - ``eid.task``/``run``/``reward`` are left default here and - assigned by ``MemoryStore.update_reward_by_record_key`` at - consume time.) + ``eid.batch``/``task``/``run`` and reward are assigned from + record key by ``MemoryStore.update`` at consume time.) API key / record key -> info["record_key"] (the recording identity; **the group key** the MemoryStore batches experiences by, so a whole reward unit's samples/turns are reward-updated and - consumed together. Falls back to ``eid.suffix`` when - absent.) + consumed together.) sample index -> info["sample_index"] (position within the n-completion set; orders samples/turns inside a record-key group) prompt_token_ids -> tokens (prompt portion) + prompt_length @@ -36,6 +34,7 @@ endpoint / model_version) stashed in ``Experience.info`` so it round-trips with the experience through serialize/deserialize. """ + from typing import Any, List, Optional import torch @@ -159,8 +158,8 @@ def build_experience( request had no prompt or no completion with response tokens. """ request_id = output.request_id - # eid.suffix = request_id for traceability; task/run/reward are left - # default and assigned by MemoryStore.update_reward_by_record_key at consume. + # eid.suffix = request_id for traceability; batch/task/run and reward are + # assigned from record_key by MemoryStore.update at consume. prompt_token_ids = list(output.prompt_token_ids or []) if not prompt_token_ids: diff --git a/trinity/common/models/vllm_patch/recording/recorder.py b/trinity/common/models/vllm_patch/recording/recorder.py index 1c14bbe49a4..3c3e8606fa8 100644 --- a/trinity/common/models/vllm_patch/recording/recorder.py +++ b/trinity/common/models/vllm_patch/recording/recorder.py @@ -17,18 +17,19 @@ computation — the client response is unchanged unless the client itself requested logprobs. Recording stays transparent. """ + import functools import logging from types import SimpleNamespace from typing import Optional +from trinity.buffer.store import MemoryStore, RecordStore from trinity.common.models.recording.context import record_key_ctx from trinity.common.models.recording.recorder import ( TRINITY_RECORD_STORE_ATTR, TRINITY_RECORDER_ATTR, Recorder, ) -from trinity.common.models.recording.store import MemoryStore, RecordStore from trinity.common.models.vllm_patch.recording.models import build_experience #: Guard attribute marking the wrapped generate, mirroring api_patch_v17 style. diff --git a/trinity/common/workflows/workflow.py b/trinity/common/workflows/workflow.py index ca4732878f6..b67e2be893c 100644 --- a/trinity/common/workflows/workflow.py +++ b/trinity/common/workflows/workflow.py @@ -85,9 +85,15 @@ class Workflow: auxiliary_models: List of OpenAI clients (sync or async based on is_async) for auxiliary models. """ - can_reset: bool = False # whether the workflow can be reset with a new task. If true, `reset()` must be implemented. - can_repeat: bool = False # whether the workflow can be repeated multiple times. If true, `set_repeat_times()` must be implemented. - is_async: bool = False # whether the workflow runs in async mode. If true, `run_async()` must be implemented, else `run()` must be implemented. + can_reset: bool = ( + False # whether the workflow can be reset with a new task. If true, `reset()` must be implemented. + ) + can_repeat: bool = ( + False # whether the workflow can be repeated multiple times. If true, `set_repeat_times()` must be implemented. + ) + is_async: bool = ( + False # whether the workflow runs in async mode. If true, `run_async()` must be implemented, else `run()` must be implemented. + ) # When True, the workflow cooperates with the in-vLLM recorder: each chat # call carries a per-sample ``record_key`` (the MemoryStore group key) so @@ -232,7 +238,7 @@ def process_messages_to_experience( # captured the authoritative heavy data (real logprobs without an extra # forward, real routed_experts) into the MemoryStore, keyed by the # record_key the chat call carried. Replace this with an in-process - # lookup by record_key (store.get_record_experiences / update), then concatenate the + # lookup by record_key (store.get / update), then concatenate the # session's turns (info["sample_index"] orders them) into one # experience here. Requires threading the per-call record_key down to # this call site. From 8b91f8ae838f5479c73f3755e374e871b9c4fea3 Mon Sep 17 00:00:00 2001 From: "panxuchen.pxc" Date: Mon, 29 Jun 2026 10:09:13 +0800 Subject: [PATCH 29/84] Merge prefix-matched recorded experiences --- tests/common/recording_recorder_test.py | 210 ++++++++++++++++++ trinity/buffer/store/base_store.py | 4 + trinity/buffer/store/memory_store.py | 27 +++ trinity/common/models/model.py | 1 + trinity/common/models/recording/query.py | 9 + trinity/common/models/recording/recorder.py | 234 +++++++++++++++++++- 6 files changed, 482 insertions(+), 3 deletions(-) create mode 100644 tests/common/recording_recorder_test.py diff --git a/tests/common/recording_recorder_test.py b/tests/common/recording_recorder_test.py new file mode 100644 index 00000000000..6561ebb226d --- /dev/null +++ b/tests/common/recording_recorder_test.py @@ -0,0 +1,210 @@ +import unittest + +import torch + +from trinity.buffer.store import MemoryStore +from trinity.common.experience import EID, Experience +from trinity.common.models.recording.recorder import Recorder + + +def make_turn( + *, + request_id: str, + record_key: str, + tokens: list[int], + prompt_length: int, + logprobs: list[float], + sample_id: str | None = None, + sample_index: int = 0, +) -> Experience: + info = { + "record_key": record_key, + "request_id": request_id, + "sample_index": sample_index, + } + if sample_id is not None: + info["sample_id"] = sample_id + return Experience( + eid=EID(suffix=request_id), + tokens=tokens, + prompt_length=prompt_length, + logprobs=logprobs, + info=info, + ) + + +class RecorderPrefixMergeTest(unittest.IsolatedAsyncioTestCase): + async def test_prefix_experiences_merge_and_keep_final_sample_id(self): + store = MemoryStore() + recorder = Recorder( + store=store, + build_experiences=lambda *_args, **_kwargs: [], + enabled=True, + ) + record_key = "0/task_a/1" + first = make_turn( + request_id="req-1", + record_key=record_key, + tokens=[10, 11, 20, 21], + prompt_length=2, + logprobs=[-0.2, -0.3], + sample_id="sample-old", + ) + second = make_turn( + request_id="req-2", + record_key=record_key, + tokens=[10, 11, 20, 21, 12, 13, 30, 31, 32], + prompt_length=6, + logprobs=[-0.4, -0.5, -0.6], + sample_id="sample-final", + ) + + await recorder._safe_append(first) + await recorder._safe_append(second) + + recorded = store.get(record_key) + self.assertEqual(len(recorded), 1) + merged = recorded[0] + self.assertEqual(merged.info["sample_id"], "sample-final") + self.assertEqual(merged.info["request_id"], "req-2") + self.assertEqual(merged.eid.suffix, "req-2") + self.assertEqual(merged.prompt_length, 2) + self.assertTrue(torch.equal(merged.tokens, second.tokens)) + self.assertTrue( + torch.equal( + merged.action_mask, + torch.tensor([True, True, False, False, True, True, True]), + ) + ) + self.assertTrue( + torch.allclose( + merged.logprobs, + torch.tensor([-0.2, -0.3, 0.0, 0.0, -0.4, -0.5, -0.6]), + ) + ) + self.assertEqual(merged.info["merged_request_ids"], ["req-1", "req-2"]) + self.assertEqual(merged.info["merged_sample_ids"], ["sample-old", "sample-final"]) + + store.update(record_key, reward=1.0, info=None, sample_ids=["sample-final"]) + self.assertEqual(store.get(record_key)[0].reward, 1.0) + with self.assertRaises(KeyError): + store.update(record_key, reward=2.0, info=None, sample_ids=["sample-old"]) + + async def test_non_prefix_experiences_do_not_merge(self): + store = MemoryStore() + recorder = Recorder( + store=store, + build_experiences=lambda *_args, **_kwargs: [], + enabled=True, + ) + record_key = "0/task_a/1" + + await recorder._safe_append( + make_turn( + request_id="req-1", + record_key=record_key, + tokens=[10, 11, 20], + prompt_length=2, + logprobs=[-0.2], + ) + ) + await recorder._safe_append( + make_turn( + request_id="req-2", + record_key=record_key, + tokens=[10, 12, 30], + prompt_length=2, + logprobs=[-0.3], + ) + ) + + self.assertEqual(len(store.get(record_key)), 2) + + async def test_merge_head_replaces_only_matching_sample_stream(self): + store = MemoryStore() + recorder = Recorder( + store=store, + build_experiences=lambda *_args, **_kwargs: [], + enabled=True, + ) + record_key = "0/task_a/1" + sample_zero = make_turn( + request_id="req-1", + record_key=record_key, + tokens=[10, 11, 20], + prompt_length=2, + logprobs=[-0.2], + sample_id="sample-zero", + sample_index=0, + ) + sample_one_first = make_turn( + request_id="req-2", + record_key=record_key, + tokens=[10, 11, 21], + prompt_length=2, + logprobs=[-0.3], + sample_id="sample-one-old", + sample_index=1, + ) + sample_one_final = make_turn( + request_id="req-3", + record_key=record_key, + tokens=[10, 11, 21, 12, 31], + prompt_length=4, + logprobs=[-0.4], + sample_id="sample-one-final", + sample_index=1, + ) + + await recorder._safe_append(sample_zero) + await recorder._safe_append(sample_one_first) + await recorder._safe_append(sample_one_final) + + recorded = store.get(record_key) + self.assertEqual(len(recorded), 2) + self.assertEqual(recorded[0].info["sample_id"], "sample-zero") + self.assertEqual(recorded[1].info["sample_id"], "sample-one-final") + self.assertTrue( + torch.equal( + recorded[1].action_mask, + torch.tensor([True, False, True]), + ) + ) + + async def test_stale_merge_head_falls_back_to_append(self): + store = MemoryStore() + recorder = Recorder( + store=store, + build_experiences=lambda *_args, **_kwargs: [], + enabled=True, + ) + record_key = "0/task_a/1" + first = make_turn( + request_id="req-1", + record_key=record_key, + tokens=[10, 11, 20], + prompt_length=2, + logprobs=[-0.2], + sample_id="sample-old", + ) + second = make_turn( + request_id="req-2", + record_key=record_key, + tokens=[10, 11, 20, 12, 30], + prompt_length=4, + logprobs=[-0.3], + sample_id="sample-final", + ) + + await recorder._safe_append(first) + store.remove(record_key) + await recorder._safe_append(second) + + recorded = store.get(record_key) + self.assertEqual(len(recorded), 1) + self.assertEqual(recorded[0].info["sample_id"], "sample-final") + self.assertEqual(recorded[0].prompt_length, 4) + + +if __name__ == "__main__": + unittest.main() diff --git a/trinity/buffer/store/base_store.py b/trinity/buffer/store/base_store.py index 84e5a655f27..570b9416dff 100644 --- a/trinity/buffer/store/base_store.py +++ b/trinity/buffer/store/base_store.py @@ -19,6 +19,10 @@ def add(self, key: str, exps: List[Experience]) -> None: def overwrite(self, key: str, exps: List[Experience]) -> None: """Replace all experiences under the given complete key.""" + @abstractmethod + def replace(self, key: str, old_sample_id: str, exp: Experience) -> None: + """Replace one experience under the given complete key.""" + @abstractmethod def update( self, key: str, reward: float, info: dict | None, sample_ids: List[str] | None diff --git a/trinity/buffer/store/memory_store.py b/trinity/buffer/store/memory_store.py index 905b4b5739d..fe137c92343 100644 --- a/trinity/buffer/store/memory_store.py +++ b/trinity/buffer/store/memory_store.py @@ -77,6 +77,33 @@ def overwrite(self, key: str, exps: List[Experience]) -> None: self._drop_key(key) self.add(key, exps) + def replace(self, key: str, old_sample_id: str, exp: Experience) -> None: + self._validate_complete_key(key) + records = self._records.get(key) + if records is None: + raise KeyError(f"Key '{key}' does not exist.") + if old_sample_id not in records: + raise KeyError(f"sample_id '{old_sample_id}' does not exist under key '{key}'.") + + new_sample_id = self.sample_id_getter(exp) + owner_key = self._sample_to_key.get(new_sample_id) + if owner_key is not None and (owner_key != key or new_sample_id != old_sample_id): + raise ValueError( + f"Duplicate sample_id '{new_sample_id}' already exists under key '{owner_key}'." + ) + + items = [] + for sample_id, record in records.items(): + if sample_id == old_sample_id: + items.append((new_sample_id, exp)) + else: + items.append((sample_id, record)) + + records.clear() + records.update(items) + self._sample_to_key.pop(old_sample_id, None) + self._sample_to_key[new_sample_id] = key + def update( self, key: str, diff --git a/trinity/common/models/model.py b/trinity/common/models/model.py index 4daec467f76..78b9e668857 100644 --- a/trinity/common/models/model.py +++ b/trinity/common/models/model.py @@ -158,6 +158,7 @@ async def extract_experience_from_history( exps = recorder.store.get(record_key) if clear_history: recorder.store.remove(record_key) + recorder.forget_record(record_key) return exps def get_model_config(self) -> InferenceModelConfig: diff --git a/trinity/common/models/recording/query.py b/trinity/common/models/recording/query.py index 4131c24b93c..373658e1ede 100644 --- a/trinity/common/models/recording/query.py +++ b/trinity/common/models/recording/query.py @@ -50,6 +50,12 @@ def _get_exp(store: RecordStore, record_key: str, request_id: str) -> Experience raise HTTPException(status_code=404, detail="experience not found") +def _forget_record(request: Request, record_key: str) -> None: + rec = getattr(request.app.state, RECORDER_STATE_ATTR, None) + if rec is not None: + rec.forget_record(record_key) + + @query_router.get("") async def list_records(request: Request) -> dict: store = _store(request) @@ -78,6 +84,7 @@ async def get_request_experience(record_key: str, request_id: str, request: Requ async def delete_record_experiences(record_key: str, request: Request) -> dict: store = _store(request) store.remove(record_key) + _forget_record(request, record_key) return {"record_key": record_key, "deleted": True} @@ -98,6 +105,7 @@ async def delete_request_experience(record_key: str, request_id: str, request: R store.overwrite(record_key, kept) else: store.remove(record_key) + _forget_record(request, record_key) return {"record_key": record_key, "request_id": request_id, "deleted": True} @@ -118,6 +126,7 @@ async def update_record(req: _UpdateRecordRequest, request: Request) -> Response sample_ids=None, ) exps.extend(store.remove(update.record_key)) + recorder.forget_record(update.record_key) return Response( content=Experience.serialize_many(exps), media_type="application/octet-stream", diff --git a/trinity/common/models/recording/recorder.py b/trinity/common/models/recording/recorder.py index d9bbbc37ae7..35dd1f175d7 100644 --- a/trinity/common/models/recording/recorder.py +++ b/trinity/common/models/recording/recorder.py @@ -6,7 +6,9 @@ from datetime import datetime, timezone from typing import Any, Optional -from trinity.buffer.store import RecordStore, get_record_key +import torch + +from trinity.buffer.store import RecordStore, default_sample_id_getter, get_record_key from trinity.common.experience import Experience from trinity.common.models.recording.context import skip_recording_ctx @@ -32,15 +34,18 @@ def __init__( enabled: bool, rank: int = 0, engine_client: Any = None, + merge_prefix_experiences: bool = True, ) -> None: self.store = store self.enabled = enabled self.rank = rank self.engine_client = engine_client + self.merge_prefix_experiences = merge_prefix_experiences self._build_experiences = build_experiences self._queue: "asyncio.Queue[Optional[Experience]]" = asyncio.Queue() self._flusher: Optional[asyncio.Task] = None self._pending: "set[asyncio.Task]" = set() + self._merge_heads: dict[str, dict[tuple[str, Any], Experience]] = {} def start(self) -> None: """Start the background flusher. Idempotent.""" @@ -97,9 +102,232 @@ async def _flush_loop(self) -> None: async def _safe_append(self, exp: Experience) -> None: try: - self.store.add(get_record_key(exp), [exp]) + record_key = get_record_key(exp) + if self.merge_prefix_experiences and self._merge_or_append(record_key, exp): + return + self.store.add(record_key, [exp]) + if self.merge_prefix_experiences: + self._remember_merge_head(record_key, exp) except Exception: logging.getLogger(__name__).exception( - "recording store.add failed for request %s", + "recording store write failed for request %s", exp.info.get("request_id"), ) + + def _merge_or_append(self, record_key: str, exp: Experience) -> bool: + stream_key = _sample_stream_key(exp) + heads = self._merge_heads.setdefault(record_key, {}) + candidate = heads.get(stream_key) + if candidate is None: + candidate = _find_longest_prefix_experience(self.store.get(record_key), exp) + elif not _is_strict_token_prefix(candidate.tokens, exp.tokens): + return False + if candidate is None: + return False + + old_sample_id = default_sample_id_getter(candidate) + merged = _merge_prefix_experiences(candidate, exp) + try: + self.store.replace(record_key, old_sample_id, merged) + except KeyError: + heads.pop(stream_key, None) + return False + heads[stream_key] = merged + return True + + def _remember_merge_head(self, record_key: str, exp: Experience) -> None: + self._merge_heads.setdefault(record_key, {})[_sample_stream_key(exp)] = exp + + def forget_record(self, record_key: str) -> None: + self._merge_heads.pop(record_key, None) + + +def _find_longest_prefix_experience( + existing: Sequence[Experience], + exp: Experience, +) -> Optional[Experience]: + best_candidate = None + best_length = -1 + for candidate in existing: + candidate_length = len(candidate.tokens) + if candidate_length <= best_length: + continue + if not _same_sample_stream(candidate, exp): + continue + if _is_strict_token_prefix(candidate.tokens, exp.tokens): + best_candidate = candidate + best_length = candidate_length + return best_candidate + + +def _same_sample_stream(left: Experience, right: Experience) -> bool: + return _sample_stream_key(left) == _sample_stream_key(right) + + +def _sample_stream_key(exp: Experience) -> tuple[str, Any]: + info = exp.info or {} + sample_index = info.get("sample_index") + if sample_index is not None: + return ("sample_index", sample_index) + + sample_id = info.get("sample_id") + if sample_id is not None: + return ("sample_id", sample_id) + + request_id = info.get("request_id") + if isinstance(request_id, str): + _, sep, suffix = request_id.rpartition(":") + if sep and suffix.isdigit(): + return ("request_id_sample_index", int(suffix)) + + return ("default", 0) + + +def _is_strict_token_prefix(prefix: torch.Tensor, tokens: torch.Tensor) -> bool: + prefix_len = len(prefix) + if prefix_len == 0 or prefix_len >= len(tokens): + return False + if prefix.device == tokens.device: + return bool(torch.equal(prefix.detach(), tokens[:prefix_len].detach())) + return bool(torch.equal(prefix.detach().cpu(), tokens[:prefix_len].detach().cpu())) + + +def _merge_prefix_experiences(prefix_exp: Experience, final_exp: Experience) -> Experience: + prefix_len = len(prefix_exp.tokens) + final_prompt_length = final_exp.prompt_length + if final_prompt_length < prefix_len: + final_prompt_length = prefix_len + gap_len = final_prompt_length - prefix_len + final_response_len = len(final_exp.tokens) - final_prompt_length + + prefix_action_mask = _response_action_mask(prefix_exp) + final_source_mask = _response_action_mask(final_exp) + final_action_mask = ( + final_source_mask[-final_response_len:] if final_response_len else final_source_mask[:0] + ) + if gap_len: + action_mask = torch.cat( + [ + prefix_action_mask, + torch.zeros(gap_len, dtype=torch.bool, device=prefix_action_mask.device), + final_action_mask, + ] + ) + else: + action_mask = torch.cat([prefix_action_mask, final_action_mask]) + + logprobs = _merge_logprobs(prefix_exp, final_exp, gap_len, final_response_len) + routed_experts = _merge_routed_experts(prefix_exp, final_exp, gap_len, final_response_len) + info = _merge_info(prefix_exp, final_exp) + + return Experience( + eid=final_exp.eid, + tokens=final_exp.tokens, + logprobs=logprobs, + reward=final_exp.reward, + token_level_reward=final_exp.token_level_reward, + advantages=final_exp.advantages, + returns=final_exp.returns, + truncate_status=final_exp.truncate_status or prefix_exp.truncate_status, + info=info, + metrics=final_exp.metrics, + prompt_length=prefix_exp.prompt_length, + response_text=final_exp.response_text, + prompt_text=prefix_exp.prompt_text, + action_mask=action_mask, + messages=final_exp.messages or prefix_exp.messages, + tools=final_exp.tools or prefix_exp.tools, + multi_modal_inputs=final_exp.multi_modal_inputs, + teacher_logprobs=final_exp.teacher_logprobs, + routed_experts=routed_experts, + custom_fields=final_exp.custom_fields, + ) + + +def _response_action_mask(exp: Experience) -> torch.Tensor: + response_len = len(exp.tokens) - exp.prompt_length + if exp.action_mask is None: + return torch.ones(response_len, dtype=torch.bool) + return exp.action_mask.to(dtype=torch.bool) + + +def _merge_logprobs( + prefix_exp: Experience, + final_exp: Experience, + gap_len: int, + final_response_len: int, +) -> Optional[torch.Tensor]: + if prefix_exp.logprobs is None or final_exp.logprobs is None: + return None + parts = [prefix_exp.logprobs] + if gap_len: + parts.append( + torch.zeros( + gap_len, + dtype=prefix_exp.logprobs.dtype, + device=prefix_exp.logprobs.device, + ) + ) + parts.append( + final_exp.logprobs[-final_response_len:] if final_response_len else final_exp.logprobs[:0] + ) + return torch.cat(parts) + + +def _merge_routed_experts( + prefix_exp: Experience, + final_exp: Experience, + gap_len: int, + final_response_len: int, +) -> Optional[torch.Tensor]: + prefix_routed = _response_routed_experts(prefix_exp) + final_routed = _response_routed_experts(final_exp) + if prefix_routed is None or final_routed is None: + return None + parts = [prefix_routed] + if gap_len: + parts.append( + torch.zeros( + (gap_len, *prefix_routed.shape[1:]), + dtype=prefix_routed.dtype, + device=prefix_routed.device, + ) + ) + parts.append(final_routed[-final_response_len:] if final_response_len else final_routed[:0]) + return torch.cat(parts, dim=0) + + +def _response_routed_experts(exp: Experience) -> Optional[torch.Tensor]: + routed = exp.routed_experts + if routed is None: + return None + response_len = len(exp.tokens) - exp.prompt_length + if len(routed) == response_len: + return routed + # Full-sequence routing is aligned to next-token predictions: + # token i uses routing row i - 1, so response tokens start at prompt_length - 1. + if len(routed) == len(exp.tokens) - 1: + return routed[exp.prompt_length - 1 :] + return None + + +def _merge_info(prefix_exp: Experience, final_exp: Experience) -> dict: + info = dict(final_exp.info or {}) + + merged_request_ids = list((prefix_exp.info or {}).get("merged_request_ids") or []) + prefix_request_id = (prefix_exp.info or {}).get("request_id") + if prefix_request_id is not None and prefix_request_id not in merged_request_ids: + merged_request_ids.append(prefix_request_id) + final_request_id = (final_exp.info or {}).get("request_id") + if final_request_id is not None and final_request_id not in merged_request_ids: + merged_request_ids.append(final_request_id) + if merged_request_ids: + info["merged_request_ids"] = merged_request_ids + + merged_sample_ids = list((prefix_exp.info or {}).get("merged_sample_ids") or []) + for sample_id in (default_sample_id_getter(prefix_exp), default_sample_id_getter(final_exp)): + if sample_id not in merged_sample_ids: + merged_sample_ids.append(sample_id) + info["merged_sample_ids"] = merged_sample_ids + info["merged_turn_count"] = int((prefix_exp.info or {}).get("merged_turn_count") or 1) + 1 + return info From fcca4a6fb388592f4191c283ae7d022be992aea4 Mon Sep 17 00:00:00 2001 From: "panxuchen.pxc" Date: Mon, 29 Jun 2026 10:18:56 +0800 Subject: [PATCH 30/84] finish merger --- trinity/common/models/recording/__init__.py | 5 +- trinity/common/models/recording/merger.py | 235 ++++++++++++++++++ trinity/common/models/recording/recorder.py | 228 +---------------- trinity/common/models/sglang_model.py | 4 +- .../models/sglang_patch/recording/__init__.py | 8 +- .../models/vllm_patch/recording/__init__.py | 4 +- trinity/common/workflows/workflow.py | 12 +- 7 files changed, 259 insertions(+), 237 deletions(-) create mode 100644 trinity/common/models/recording/merger.py diff --git a/trinity/common/models/recording/__init__.py b/trinity/common/models/recording/__init__.py index 1e1e9bd515c..2d849090c23 100644 --- a/trinity/common/models/recording/__init__.py +++ b/trinity/common/models/recording/__init__.py @@ -13,7 +13,10 @@ TRINITY_RECORDER_ATTR, Recorder, ) -from trinity.common.models.recording.server import add_recording_middleware, mount_recording_api +from trinity.common.models.recording.server import ( + add_recording_middleware, + mount_recording_api, +) __all__ = [ "MemoryStore", diff --git a/trinity/common/models/recording/merger.py b/trinity/common/models/recording/merger.py new file mode 100644 index 00000000000..53feb0a65d6 --- /dev/null +++ b/trinity/common/models/recording/merger.py @@ -0,0 +1,235 @@ +"""Prefix-based merging for recorded multi-turn experiences.""" + +from collections.abc import Sequence +from typing import Any, Optional + +import torch + +from trinity.buffer.store import RecordStore, default_sample_id_getter +from trinity.common.experience import Experience + + +class PrefixExperienceMerger: + """Merge same-record experiences whose tokens form a strict prefix chain.""" + + def __init__(self, store: RecordStore) -> None: + self.store = store + self._heads: dict[str, dict[tuple[str, Any], Experience]] = {} + + def try_merge(self, record_key: str, exp: Experience) -> bool: + stream_key = _sample_stream_key(exp) + heads = self._heads.setdefault(record_key, {}) + candidate = heads.get(stream_key) + if candidate is None: + candidate = _find_longest_prefix_experience(self.store.get(record_key), exp) + elif not _is_strict_token_prefix(candidate.tokens, exp.tokens): + return False + if candidate is None: + return False + + old_sample_id = default_sample_id_getter(candidate) + merged = _merge_prefix_experiences(candidate, exp) + try: + self.store.replace(record_key, old_sample_id, merged) + except KeyError: + heads.pop(stream_key, None) + return False + heads[stream_key] = merged + return True + + def remember(self, record_key: str, exp: Experience) -> None: + self._heads.setdefault(record_key, {})[_sample_stream_key(exp)] = exp + + def forget_record(self, record_key: str) -> None: + self._heads.pop(record_key, None) + + +def _find_longest_prefix_experience( + existing: Sequence[Experience], + exp: Experience, +) -> Optional[Experience]: + best_candidate = None + best_length = -1 + for candidate in existing: + candidate_length = len(candidate.tokens) + if candidate_length <= best_length: + continue + if not _same_sample_stream(candidate, exp): + continue + if _is_strict_token_prefix(candidate.tokens, exp.tokens): + best_candidate = candidate + best_length = candidate_length + return best_candidate + + +def _same_sample_stream(left: Experience, right: Experience) -> bool: + return _sample_stream_key(left) == _sample_stream_key(right) + + +def _sample_stream_key(exp: Experience) -> tuple[str, Any]: + info = exp.info or {} + sample_index = info.get("sample_index") + if sample_index is not None: + return ("sample_index", sample_index) + + sample_id = info.get("sample_id") + if sample_id is not None: + return ("sample_id", sample_id) + + request_id = info.get("request_id") + if isinstance(request_id, str): + _, sep, suffix = request_id.rpartition(":") + if sep and suffix.isdigit(): + return ("request_id_sample_index", int(suffix)) + + return ("default", 0) + + +def _is_strict_token_prefix(prefix: torch.Tensor, tokens: torch.Tensor) -> bool: + prefix_len = len(prefix) + if prefix_len == 0 or prefix_len >= len(tokens): + return False + if prefix.device == tokens.device: + return bool(torch.equal(prefix.detach(), tokens[:prefix_len].detach())) + return bool(torch.equal(prefix.detach().cpu(), tokens[:prefix_len].detach().cpu())) + + +def _merge_prefix_experiences(prefix_exp: Experience, final_exp: Experience) -> Experience: + prefix_len = len(prefix_exp.tokens) + final_prompt_length = final_exp.prompt_length + if final_prompt_length < prefix_len: + final_prompt_length = prefix_len + gap_len = final_prompt_length - prefix_len + final_response_len = len(final_exp.tokens) - final_prompt_length + + prefix_action_mask = _response_action_mask(prefix_exp) + final_source_mask = _response_action_mask(final_exp) + final_action_mask = ( + final_source_mask[-final_response_len:] if final_response_len else final_source_mask[:0] + ) + if gap_len: + action_mask = torch.cat( + [ + prefix_action_mask, + torch.zeros(gap_len, dtype=torch.bool, device=prefix_action_mask.device), + final_action_mask, + ] + ) + else: + action_mask = torch.cat([prefix_action_mask, final_action_mask]) + + logprobs = _merge_logprobs(prefix_exp, final_exp, gap_len, final_response_len) + routed_experts = _merge_routed_experts(prefix_exp, final_exp, gap_len, final_response_len) + info = _merge_info(prefix_exp, final_exp) + + return Experience( + eid=final_exp.eid, + tokens=final_exp.tokens, + logprobs=logprobs, + reward=final_exp.reward, + token_level_reward=final_exp.token_level_reward, + advantages=final_exp.advantages, + returns=final_exp.returns, + truncate_status=final_exp.truncate_status or prefix_exp.truncate_status, + info=info, + metrics=final_exp.metrics, + prompt_length=prefix_exp.prompt_length, + response_text=final_exp.response_text, + prompt_text=prefix_exp.prompt_text, + action_mask=action_mask, + messages=final_exp.messages or prefix_exp.messages, + tools=final_exp.tools or prefix_exp.tools, + multi_modal_inputs=final_exp.multi_modal_inputs, + teacher_logprobs=final_exp.teacher_logprobs, + routed_experts=routed_experts, + custom_fields=final_exp.custom_fields, + ) + + +def _response_action_mask(exp: Experience) -> torch.Tensor: + response_len = len(exp.tokens) - exp.prompt_length + if exp.action_mask is None: + return torch.ones(response_len, dtype=torch.bool) + return exp.action_mask.to(dtype=torch.bool) + + +def _merge_logprobs( + prefix_exp: Experience, + final_exp: Experience, + gap_len: int, + final_response_len: int, +) -> Optional[torch.Tensor]: + if prefix_exp.logprobs is None or final_exp.logprobs is None: + return None + parts = [prefix_exp.logprobs] + if gap_len: + parts.append( + torch.zeros( + gap_len, + dtype=prefix_exp.logprobs.dtype, + device=prefix_exp.logprobs.device, + ) + ) + parts.append( + final_exp.logprobs[-final_response_len:] if final_response_len else final_exp.logprobs[:0] + ) + return torch.cat(parts) + + +def _merge_routed_experts( + prefix_exp: Experience, + final_exp: Experience, + gap_len: int, + final_response_len: int, +) -> Optional[torch.Tensor]: + prefix_routed = _response_routed_experts(prefix_exp) + final_routed = _response_routed_experts(final_exp) + if prefix_routed is None or final_routed is None: + return None + parts = [prefix_routed] + if gap_len: + parts.append( + torch.zeros( + (gap_len, *prefix_routed.shape[1:]), + dtype=prefix_routed.dtype, + device=prefix_routed.device, + ) + ) + parts.append(final_routed[-final_response_len:] if final_response_len else final_routed[:0]) + return torch.cat(parts, dim=0) + + +def _response_routed_experts(exp: Experience) -> Optional[torch.Tensor]: + routed = exp.routed_experts + if routed is None: + return None + response_len = len(exp.tokens) - exp.prompt_length + if len(routed) == response_len: + return routed + # Full-sequence routing is aligned to next-token predictions: + # token i uses routing row i - 1, so response tokens start at prompt_length - 1. + if len(routed) == len(exp.tokens) - 1: + return routed[exp.prompt_length - 1 :] + return None + + +def _merge_info(prefix_exp: Experience, final_exp: Experience) -> dict: + info = dict(final_exp.info or {}) + + merged_request_ids = list((prefix_exp.info or {}).get("merged_request_ids") or []) + prefix_request_id = (prefix_exp.info or {}).get("request_id") + if prefix_request_id is not None and prefix_request_id not in merged_request_ids: + merged_request_ids.append(prefix_request_id) + final_request_id = (final_exp.info or {}).get("request_id") + if final_request_id is not None and final_request_id not in merged_request_ids: + merged_request_ids.append(final_request_id) + if merged_request_ids: + info["merged_request_ids"] = merged_request_ids + + merged_sample_ids = list((prefix_exp.info or {}).get("merged_sample_ids") or []) + for sample_id in (default_sample_id_getter(prefix_exp), default_sample_id_getter(final_exp)): + if sample_id not in merged_sample_ids: + merged_sample_ids.append(sample_id) + info["merged_sample_ids"] = merged_sample_ids + info["merged_turn_count"] = int((prefix_exp.info or {}).get("merged_turn_count") or 1) + 1 + return info diff --git a/trinity/common/models/recording/recorder.py b/trinity/common/models/recording/recorder.py index 35dd1f175d7..2c9e02cab9f 100644 --- a/trinity/common/models/recording/recorder.py +++ b/trinity/common/models/recording/recorder.py @@ -6,11 +6,10 @@ from datetime import datetime, timezone from typing import Any, Optional -import torch - -from trinity.buffer.store import RecordStore, default_sample_id_getter, get_record_key +from trinity.buffer.store import RecordStore, get_record_key from trinity.common.experience import Experience from trinity.common.models.recording.context import skip_recording_ctx +from trinity.common.models.recording.merger import PrefixExperienceMerger MODEL_VERSION_ATTR = "trinity_model_version" TRINITY_RECORDER_ATTR = "trinity_recorder" @@ -45,7 +44,7 @@ def __init__( self._queue: "asyncio.Queue[Optional[Experience]]" = asyncio.Queue() self._flusher: Optional[asyncio.Task] = None self._pending: "set[asyncio.Task]" = set() - self._merge_heads: dict[str, dict[tuple[str, Any], Experience]] = {} + self._prefix_merger = PrefixExperienceMerger(store) def start(self) -> None: """Start the background flusher. Idempotent.""" @@ -103,231 +102,16 @@ async def _flush_loop(self) -> None: async def _safe_append(self, exp: Experience) -> None: try: record_key = get_record_key(exp) - if self.merge_prefix_experiences and self._merge_or_append(record_key, exp): + if self.merge_prefix_experiences and self._prefix_merger.try_merge(record_key, exp): return self.store.add(record_key, [exp]) if self.merge_prefix_experiences: - self._remember_merge_head(record_key, exp) + self._prefix_merger.remember(record_key, exp) except Exception: logging.getLogger(__name__).exception( "recording store write failed for request %s", exp.info.get("request_id"), ) - def _merge_or_append(self, record_key: str, exp: Experience) -> bool: - stream_key = _sample_stream_key(exp) - heads = self._merge_heads.setdefault(record_key, {}) - candidate = heads.get(stream_key) - if candidate is None: - candidate = _find_longest_prefix_experience(self.store.get(record_key), exp) - elif not _is_strict_token_prefix(candidate.tokens, exp.tokens): - return False - if candidate is None: - return False - - old_sample_id = default_sample_id_getter(candidate) - merged = _merge_prefix_experiences(candidate, exp) - try: - self.store.replace(record_key, old_sample_id, merged) - except KeyError: - heads.pop(stream_key, None) - return False - heads[stream_key] = merged - return True - - def _remember_merge_head(self, record_key: str, exp: Experience) -> None: - self._merge_heads.setdefault(record_key, {})[_sample_stream_key(exp)] = exp - def forget_record(self, record_key: str) -> None: - self._merge_heads.pop(record_key, None) - - -def _find_longest_prefix_experience( - existing: Sequence[Experience], - exp: Experience, -) -> Optional[Experience]: - best_candidate = None - best_length = -1 - for candidate in existing: - candidate_length = len(candidate.tokens) - if candidate_length <= best_length: - continue - if not _same_sample_stream(candidate, exp): - continue - if _is_strict_token_prefix(candidate.tokens, exp.tokens): - best_candidate = candidate - best_length = candidate_length - return best_candidate - - -def _same_sample_stream(left: Experience, right: Experience) -> bool: - return _sample_stream_key(left) == _sample_stream_key(right) - - -def _sample_stream_key(exp: Experience) -> tuple[str, Any]: - info = exp.info or {} - sample_index = info.get("sample_index") - if sample_index is not None: - return ("sample_index", sample_index) - - sample_id = info.get("sample_id") - if sample_id is not None: - return ("sample_id", sample_id) - - request_id = info.get("request_id") - if isinstance(request_id, str): - _, sep, suffix = request_id.rpartition(":") - if sep and suffix.isdigit(): - return ("request_id_sample_index", int(suffix)) - - return ("default", 0) - - -def _is_strict_token_prefix(prefix: torch.Tensor, tokens: torch.Tensor) -> bool: - prefix_len = len(prefix) - if prefix_len == 0 or prefix_len >= len(tokens): - return False - if prefix.device == tokens.device: - return bool(torch.equal(prefix.detach(), tokens[:prefix_len].detach())) - return bool(torch.equal(prefix.detach().cpu(), tokens[:prefix_len].detach().cpu())) - - -def _merge_prefix_experiences(prefix_exp: Experience, final_exp: Experience) -> Experience: - prefix_len = len(prefix_exp.tokens) - final_prompt_length = final_exp.prompt_length - if final_prompt_length < prefix_len: - final_prompt_length = prefix_len - gap_len = final_prompt_length - prefix_len - final_response_len = len(final_exp.tokens) - final_prompt_length - - prefix_action_mask = _response_action_mask(prefix_exp) - final_source_mask = _response_action_mask(final_exp) - final_action_mask = ( - final_source_mask[-final_response_len:] if final_response_len else final_source_mask[:0] - ) - if gap_len: - action_mask = torch.cat( - [ - prefix_action_mask, - torch.zeros(gap_len, dtype=torch.bool, device=prefix_action_mask.device), - final_action_mask, - ] - ) - else: - action_mask = torch.cat([prefix_action_mask, final_action_mask]) - - logprobs = _merge_logprobs(prefix_exp, final_exp, gap_len, final_response_len) - routed_experts = _merge_routed_experts(prefix_exp, final_exp, gap_len, final_response_len) - info = _merge_info(prefix_exp, final_exp) - - return Experience( - eid=final_exp.eid, - tokens=final_exp.tokens, - logprobs=logprobs, - reward=final_exp.reward, - token_level_reward=final_exp.token_level_reward, - advantages=final_exp.advantages, - returns=final_exp.returns, - truncate_status=final_exp.truncate_status or prefix_exp.truncate_status, - info=info, - metrics=final_exp.metrics, - prompt_length=prefix_exp.prompt_length, - response_text=final_exp.response_text, - prompt_text=prefix_exp.prompt_text, - action_mask=action_mask, - messages=final_exp.messages or prefix_exp.messages, - tools=final_exp.tools or prefix_exp.tools, - multi_modal_inputs=final_exp.multi_modal_inputs, - teacher_logprobs=final_exp.teacher_logprobs, - routed_experts=routed_experts, - custom_fields=final_exp.custom_fields, - ) - - -def _response_action_mask(exp: Experience) -> torch.Tensor: - response_len = len(exp.tokens) - exp.prompt_length - if exp.action_mask is None: - return torch.ones(response_len, dtype=torch.bool) - return exp.action_mask.to(dtype=torch.bool) - - -def _merge_logprobs( - prefix_exp: Experience, - final_exp: Experience, - gap_len: int, - final_response_len: int, -) -> Optional[torch.Tensor]: - if prefix_exp.logprobs is None or final_exp.logprobs is None: - return None - parts = [prefix_exp.logprobs] - if gap_len: - parts.append( - torch.zeros( - gap_len, - dtype=prefix_exp.logprobs.dtype, - device=prefix_exp.logprobs.device, - ) - ) - parts.append( - final_exp.logprobs[-final_response_len:] if final_response_len else final_exp.logprobs[:0] - ) - return torch.cat(parts) - - -def _merge_routed_experts( - prefix_exp: Experience, - final_exp: Experience, - gap_len: int, - final_response_len: int, -) -> Optional[torch.Tensor]: - prefix_routed = _response_routed_experts(prefix_exp) - final_routed = _response_routed_experts(final_exp) - if prefix_routed is None or final_routed is None: - return None - parts = [prefix_routed] - if gap_len: - parts.append( - torch.zeros( - (gap_len, *prefix_routed.shape[1:]), - dtype=prefix_routed.dtype, - device=prefix_routed.device, - ) - ) - parts.append(final_routed[-final_response_len:] if final_response_len else final_routed[:0]) - return torch.cat(parts, dim=0) - - -def _response_routed_experts(exp: Experience) -> Optional[torch.Tensor]: - routed = exp.routed_experts - if routed is None: - return None - response_len = len(exp.tokens) - exp.prompt_length - if len(routed) == response_len: - return routed - # Full-sequence routing is aligned to next-token predictions: - # token i uses routing row i - 1, so response tokens start at prompt_length - 1. - if len(routed) == len(exp.tokens) - 1: - return routed[exp.prompt_length - 1 :] - return None - - -def _merge_info(prefix_exp: Experience, final_exp: Experience) -> dict: - info = dict(final_exp.info or {}) - - merged_request_ids = list((prefix_exp.info or {}).get("merged_request_ids") or []) - prefix_request_id = (prefix_exp.info or {}).get("request_id") - if prefix_request_id is not None and prefix_request_id not in merged_request_ids: - merged_request_ids.append(prefix_request_id) - final_request_id = (final_exp.info or {}).get("request_id") - if final_request_id is not None and final_request_id not in merged_request_ids: - merged_request_ids.append(final_request_id) - if merged_request_ids: - info["merged_request_ids"] = merged_request_ids - - merged_sample_ids = list((prefix_exp.info or {}).get("merged_sample_ids") or []) - for sample_id in (default_sample_id_getter(prefix_exp), default_sample_id_getter(final_exp)): - if sample_id not in merged_sample_ids: - merged_sample_ids.append(sample_id) - info["merged_sample_ids"] = merged_sample_ids - info["merged_turn_count"] = int((prefix_exp.info or {}).get("merged_turn_count") or 1) + 1 - return info + self._prefix_merger.forget_record(record_key) diff --git a/trinity/common/models/sglang_model.py b/trinity/common/models/sglang_model.py index 03027dac02f..463a43e5b1f 100644 --- a/trinity/common/models/sglang_model.py +++ b/trinity/common/models/sglang_model.py @@ -570,7 +570,9 @@ async def run_api_server(self) -> bool: if self.config.enable_history: from trinity.buffer.store import MemoryStore from trinity.common.models.recording.recorder import Recorder - from trinity.common.models.sglang_patch.recording.models import build_sglang_experience + from trinity.common.models.sglang_patch.recording.models import ( + build_sglang_experience, + ) record_store = MemoryStore() recorder = Recorder( diff --git a/trinity/common/models/sglang_patch/recording/__init__.py b/trinity/common/models/sglang_patch/recording/__init__.py index 062e1bcdcf9..6ddd3f980f0 100644 --- a/trinity/common/models/sglang_patch/recording/__init__.py +++ b/trinity/common/models/sglang_patch/recording/__init__.py @@ -15,14 +15,16 @@ ) from trinity.common.models.recording.query import query_router # noqa: F401 from trinity.common.models.recording.recorder import Recorder # noqa: F401 -from trinity.common.models.sglang_patch.recording.models import ( +from trinity.common.models.sglang_patch.recording.models import ( # noqa: F401 build_sglang_experience, -) # noqa: F401 +) from trinity.common.models.sglang_patch.recording.recorder import ( # noqa: F401 create_sglang_recorder, patch_tokenizer_manager_for_recording, ) -from trinity.common.models.sglang_patch.recording.server import setup_sglang_recording # noqa: F401 +from trinity.common.models.sglang_patch.recording.server import ( # noqa: F401 + setup_sglang_recording, +) __all__ = [ "MemoryStore", diff --git a/trinity/common/models/vllm_patch/recording/__init__.py b/trinity/common/models/vllm_patch/recording/__init__.py index 10aa385f327..971db1ce81a 100644 --- a/trinity/common/models/vllm_patch/recording/__init__.py +++ b/trinity/common/models/vllm_patch/recording/__init__.py @@ -15,7 +15,9 @@ create_vllm_recorder, patch_engine_for_recording, ) -from trinity.common.models.vllm_patch.recording.server import run_api_server_with_recording +from trinity.common.models.vllm_patch.recording.server import ( + run_api_server_with_recording, +) __all__ = [ "MemoryStore", diff --git a/trinity/common/workflows/workflow.py b/trinity/common/workflows/workflow.py index b67e2be893c..e4cd706ac97 100644 --- a/trinity/common/workflows/workflow.py +++ b/trinity/common/workflows/workflow.py @@ -85,15 +85,9 @@ class Workflow: auxiliary_models: List of OpenAI clients (sync or async based on is_async) for auxiliary models. """ - can_reset: bool = ( - False # whether the workflow can be reset with a new task. If true, `reset()` must be implemented. - ) - can_repeat: bool = ( - False # whether the workflow can be repeated multiple times. If true, `set_repeat_times()` must be implemented. - ) - is_async: bool = ( - False # whether the workflow runs in async mode. If true, `run_async()` must be implemented, else `run()` must be implemented. - ) + can_reset: bool = False # whether the workflow can be reset with a new task. If true, `reset()` must be implemented. + can_repeat: bool = False # whether the workflow can be repeated multiple times. If true, `set_repeat_times()` must be implemented. + is_async: bool = False # whether the workflow runs in async mode. If true, `run_async()` must be implemented, else `run()` must be implemented. # When True, the workflow cooperates with the in-vLLM recorder: each chat # call carries a per-sample ``record_key`` (the MemoryStore group key) so From e283fcfcd199f264ad6ce279f53a8e6a381084c8 Mon Sep 17 00:00:00 2001 From: "panxuchen.pxc" Date: Mon, 29 Jun 2026 10:36:48 +0800 Subject: [PATCH 31/84] add docstring --- trinity/common/models/recording/merger.py | 22 +++++++++++++++++++++- 1 file changed, 21 insertions(+), 1 deletion(-) diff --git a/trinity/common/models/recording/merger.py b/trinity/common/models/recording/merger.py index 53feb0a65d6..19b1ca89761 100644 --- a/trinity/common/models/recording/merger.py +++ b/trinity/common/models/recording/merger.py @@ -10,7 +10,27 @@ class PrefixExperienceMerger: - """Merge same-record experiences whose tokens form a strict prefix chain.""" + """Merge same-record experiences whose tokens form a strict prefix chain. + + Strategy: + * Experiences are grouped by record key and a best-effort sample stream + key (sample_index, explicit sample_id, request_id suffix, then default). + * Each stream tracks one latest/longest head experience. A new experience + merges only when the head tokens are a strict prefix of the new tokens. + * If no head exists yet, the store is scanned once to find the longest + prefix-matching experience for that stream. + + Limitation: + This assumes each record/sample stream is a single linear conversation + branch. If one task has concurrent writers sharing the same record/sample + stream (for example, multi-agent rollouts under one record key), the latest + head may belong to a different branch, so the prefix hit can be missed or + become ambiguous. + + TODO(recording): support branching/concurrent histories by tracking multiple + heads per record/sample stream, keyed by a stable conversation/thread id or + by token-prefix fingerprints. + """ def __init__(self, store: RecordStore) -> None: self.store = store From 3607157bd2b57e07c638153db934564131127c43 Mon Sep 17 00:00:00 2001 From: "panxuchen.pxc" Date: Mon, 29 Jun 2026 12:20:52 +0800 Subject: [PATCH 32/84] simplify recorder_key sample id --- tests/buffer/memory_store_test.py | 107 ++++++++++++++++++ tests/common/recording_recorder_test.py | 14 +-- tests/common/recording_store_test.py | 99 ++++++++++++++++ tests/common/sglang_test.py | 17 ++- tests/common/vllm_test.py | 13 ++- trinity/buffer/store/__init__.py | 6 +- trinity/buffer/store/memory_store.py | 62 +++++----- trinity/common/models/recording/merger.py | 22 +--- trinity/common/models/recording/query.py | 8 +- trinity/common/models/recording/recorder.py | 2 +- .../models/sglang_patch/recording/models.py | 26 +++-- trinity/common/models/tinker_model.py | 9 +- trinity/common/models/vllm_model.py | 10 +- .../models/vllm_patch/recording/models.py | 35 +++--- trinity/common/workflows/workflow.py | 8 +- trinity/explorer/workflow_runner.py | 3 +- 16 files changed, 323 insertions(+), 118 deletions(-) create mode 100644 tests/buffer/memory_store_test.py create mode 100644 tests/common/recording_store_test.py diff --git a/tests/buffer/memory_store_test.py b/tests/buffer/memory_store_test.py new file mode 100644 index 00000000000..8470c210ce1 --- /dev/null +++ b/tests/buffer/memory_store_test.py @@ -0,0 +1,107 @@ +import unittest +import uuid + +import torch + +from trinity.buffer.store import MemoryStore +from trinity.common.experience import EID, Experience + + +def get_dummy_experience(num: int, request_id: str | None = None): + request_id = request_id or uuid.uuid4().hex[:6] + return [ + Experience( + eid=EID(suffix=request_id), + tokens=torch.zeros(5), + prompt_length=2, + info={ + "sample_index": i, + "model_version": 0, + }, + ) + for i in range(num) + ] + + +class MemoryStoreTest(unittest.TestCase): + def test_add_update_get_remove(self): + store = MemoryStore() + key = "0/task_a/1" + experiences = get_dummy_experience(3, request_id="req_a") + + store.add(key, experiences) + self.assertEqual(len(store), 3) + + store.update(key, reward=1.0, info={"source": "reward_model"}, sample_ids=None) + result = store.get(key) + self.assertEqual(len(result), 3) + for exp in result: + self.assertEqual(exp.reward, 1.0) + self.assertEqual(exp.info["source"], "reward_model") + self.assertEqual(exp.eid.batch, "0") + self.assertEqual(exp.eid.task, "task_a") + self.assertEqual(exp.eid.run, 1) + + removed = store.remove(key) + self.assertEqual(len(removed), 3) + self.assertEqual(store.get(key), []) + self.assertEqual(len(store), 0) + + def test_update_subset_by_sample_ids(self): + store = MemoryStore() + key = "0/task_a/1" + experiences = get_dummy_experience(2, request_id="req_b") + + store.add(key, experiences) + store.update(key, reward=2.0, info=None, sample_ids=["req_b:1"]) + + result = store.get(key) + self.assertIsNone(result[0].reward) + self.assertEqual(result[1].reward, 2.0) + self.assertEqual(result[1].eid.batch, "0") + self.assertEqual(result[1].eid.task, "task_a") + self.assertEqual(result[1].eid.run, 1) + + def test_overwrite_replaces_existing_records(self): + store = MemoryStore() + key = "0/task_a/1" + + store.add(key, get_dummy_experience(2, request_id="old")) + store.overwrite(key, get_dummy_experience(1, request_id="new")) + + result = store.get(key) + self.assertEqual(len(result), 1) + self.assertEqual(result[0].eid.suffix, "new") + + def test_prefix_get_and_remove(self): + store = MemoryStore() + store.add("0/task_a/0", get_dummy_experience(1, request_id="a0")) + store.add("0/task_a/1", get_dummy_experience(2, request_id="a1")) + store.add("0/task_b/0", get_dummy_experience(1, request_id="b0")) + + self.assertEqual(len(store.get("0/task_a")), 3) + self.assertEqual(len(store.remove("0/task_a")), 3) + self.assertEqual(len(store.get("0")), 1) + self.assertEqual(store.keys(), ["0/task_b/0"]) + + def test_complete_key_required_for_mutations(self): + store = MemoryStore() + with self.assertRaises(ValueError): + store.add("0/task_a", get_dummy_experience(1)) + with self.assertRaises(ValueError): + store.overwrite("0/task_a", get_dummy_experience(1)) + with self.assertRaises(ValueError): + store.update("0/task_a", reward=1.0, info=None, sample_ids=None) + with self.assertRaises(ValueError): + store.add("0/task_a/not_int", get_dummy_experience(1)) + + def test_duplicate_sample_id_is_rejected(self): + store = MemoryStore() + exp = get_dummy_experience(1, request_id="dup") + store.add("0/task_a/0", exp) + with self.assertRaises(ValueError): + store.add("0/task_a/1", exp) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/common/recording_recorder_test.py b/tests/common/recording_recorder_test.py index 6561ebb226d..362d4c2a5d1 100644 --- a/tests/common/recording_recorder_test.py +++ b/tests/common/recording_recorder_test.py @@ -2,7 +2,7 @@ import torch -from trinity.buffer.store import MemoryStore +from trinity.buffer.store import MemoryStore, parse_record_key from trinity.common.experience import EID, Experience from trinity.common.models.recording.recorder import Recorder @@ -17,15 +17,12 @@ def make_turn( sample_id: str | None = None, sample_index: int = 0, ) -> Experience: - info = { - "record_key": record_key, - "request_id": request_id, - "sample_index": sample_index, - } + batch, task, run = parse_record_key(record_key) + info = {"sample_index": sample_index} if sample_id is not None: info["sample_id"] = sample_id return Experience( - eid=EID(suffix=request_id), + eid=EID(batch=batch, task=task, run=run, suffix=request_id), tokens=tokens, prompt_length=prompt_length, logprobs=logprobs, @@ -66,7 +63,6 @@ async def test_prefix_experiences_merge_and_keep_final_sample_id(self): self.assertEqual(len(recorded), 1) merged = recorded[0] self.assertEqual(merged.info["sample_id"], "sample-final") - self.assertEqual(merged.info["request_id"], "req-2") self.assertEqual(merged.eid.suffix, "req-2") self.assertEqual(merged.prompt_length, 2) self.assertTrue(torch.equal(merged.tokens, second.tokens)) @@ -82,7 +78,7 @@ async def test_prefix_experiences_merge_and_keep_final_sample_id(self): torch.tensor([-0.2, -0.3, 0.0, 0.0, -0.4, -0.5, -0.6]), ) ) - self.assertEqual(merged.info["merged_request_ids"], ["req-1", "req-2"]) + self.assertEqual(merged.info["merged_eid_suffixes"], ["req-1", "req-2"]) self.assertEqual(merged.info["merged_sample_ids"], ["sample-old", "sample-final"]) store.update(record_key, reward=1.0, info=None, sample_ids=["sample-final"]) diff --git a/tests/common/recording_store_test.py b/tests/common/recording_store_test.py new file mode 100644 index 00000000000..facd41cf897 --- /dev/null +++ b/tests/common/recording_store_test.py @@ -0,0 +1,99 @@ +import unittest + +import torch + +from trinity.buffer.store import MemoryStore, get_record_key, parse_record_key +from trinity.common.experience import EID, Experience + + +def make_exp(request_id: str, record_key: str | None = None) -> Experience: + info = {"sample_index": 0} + eid = EID(suffix=request_id) + if record_key is not None: + batch, task, run = parse_record_key(record_key) + eid.batch = batch + eid.task = task + eid.run = run + return Experience( + eid=eid, + tokens=torch.zeros(5), + prompt_length=2, + info=info, + ) + + +class MemoryStoreTest(unittest.IsolatedAsyncioTestCase): + async def test_update_reward_sets_eid_from_record_key(self): + store = MemoryStore() + record_key = "0/task_a/1" + exp = make_exp("req_a", record_key) + + store.add(get_record_key(exp), [exp]) + store.update(record_key, reward=1.5, info={"source": "reward_model"}, sample_ids=None) + updated = store.remove(record_key) + + self.assertEqual(len(updated), 1) + self.assertEqual(updated[0].reward, 1.5) + self.assertEqual(updated[0].info["source"], "reward_model") + self.assertNotIn("run", updated[0].info) + self.assertNotIn("task", updated[0].info) + self.assertEqual(updated[0].eid.batch, "0") + self.assertEqual(updated[0].eid.task, "task_a") + self.assertEqual(updated[0].eid.run, 1) + self.assertEqual(store.get(record_key), []) + + async def test_complete_record_key_request_lookup_and_delete(self): + store = MemoryStore() + record_key = "0/task_a/1" + exp = make_exp("req_a", record_key) + + store.add(get_record_key(exp), [exp]) + + self.assertEqual(store.keys(), [record_key]) + self.assertIs(_find_request(store, record_key, "req_a"), exp) + + deleted = _delete_request(store, record_key, "req_a") + self.assertTrue(deleted) + self.assertEqual(store.keys(), []) + + async def test_delete_request_experience_keeps_other_experiences(self): + store = MemoryStore() + record_key = "0/task_a/1" + exp_a = make_exp("req_a", record_key) + exp_b = make_exp("req_b", record_key) + + store.add(get_record_key(exp_a), [exp_a]) + store.add(get_record_key(exp_b), [exp_b]) + + deleted = _delete_request(store, record_key, "req_a") + + self.assertTrue(deleted) + remaining = store.get(record_key) + self.assertEqual(remaining, [exp_b]) + + +def _find_request(store: MemoryStore, record_key: str, request_id: str) -> Experience | None: + for exp in store.get(record_key): + if exp.eid.suffix == request_id: + return exp + return None + + +def _delete_request(store: MemoryStore, record_key: str, request_id: str) -> bool: + kept = [] + deleted = False + for exp in store.get(record_key): + if exp.eid.suffix == request_id: + deleted = True + else: + kept.append(exp) + if deleted: + if kept: + store.overwrite(record_key, kept) + else: + store.remove(record_key) + return deleted + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/common/sglang_test.py b/tests/common/sglang_test.py index a22cdeb46b7..289e5598303 100644 --- a/tests/common/sglang_test.py +++ b/tests/common/sglang_test.py @@ -13,6 +13,7 @@ get_moe_model_path, get_template_config, ) +from trinity.buffer.store import get_record_key from trinity.common.experience import Experience from trinity.common.models.allocator import Allocator @@ -484,8 +485,8 @@ async def _get_model_id(self, client: openai.AsyncOpenAI) -> str: # -- per-recorded-experience invariants ----------------------------------- def _assert_recorded_experience(self, exp: Experience, record_key: str): - self.assertEqual(exp.info.get("record_key"), record_key) - self.assertIsNotNone(exp.info.get("request_id")) + self.assertEqual(get_record_key(exp), record_key) + self.assertTrue(exp.eid.suffix) self.assertEqual(exp.info.get("rank"), 0) # SGLang stamps meta_info.weight_version ("default" until a weight sync); # unlike vLLM it is a server-tracked string, not the model_version int. @@ -525,7 +526,7 @@ async def test_record(self): # noqa: C901 task = await self._get_record_experiences(rk_gen) self.assertEqual(len(task["experiences"]), 1) # blob endpoint round-trips a full experience - request_id = task["experiences"][0]["info"]["request_id"] + request_id = task["experiences"][0]["eid"]["suffix"] blob_exp = await self._get_request_experience(rk_gen, request_id) self._assert_recorded_experience(blob_exp, rk_gen) self._assert_recorded_routed_experts(blob_exp) @@ -547,13 +548,11 @@ async def test_record(self): # noqa: C901 await self._flush() task = await self._get_record_experiences(rk_chat) self.assertEqual(len(task["experiences"]), 2) - # SGLang expands n=2 parallel sampling into two scheduler requests, - # each with its own rid -> two distinct request_ids. The list position - # becomes sample_index (0, 1) to order the two samples within the - # record-key group (the GRPO-style grouping the store is built for). - # (vLLM differs: n=2 shares one request_id with sample_index 0/1.) + # SGLang expands n=2 parallel sampling into two scheduler requests. + # The list position becomes sample_index (0, 1) to order the two + # samples within the record-key group. self.assertEqual(sorted(t["info"]["sample_index"] for t in task["experiences"]), [0, 1]) - self.assertEqual(len({t["info"]["request_id"] for t in task["experiences"]}), 2) + self.assertEqual(len({t["eid"]["suffix"] for t in task["experiences"]}), 2) consumed = await self._consume(rk_chat, reward=0.8, run=2, task="t_chat") self.assertEqual(len(consumed), 2) for exp in consumed: diff --git a/tests/common/vllm_test.py b/tests/common/vllm_test.py index f39be3e2a00..7a44c67c784 100644 --- a/tests/common/vllm_test.py +++ b/tests/common/vllm_test.py @@ -24,6 +24,7 @@ get_moe_model_path, get_template_config, ) +from trinity.buffer.store import get_record_key from trinity.common.config import Config from trinity.common.constants import ROLLOUT_WEIGHT_SYNC_GROUP_NAME, SyncMethod from trinity.common.experience import Experience @@ -1928,8 +1929,8 @@ async def _get_model_id(self, client: openai.AsyncOpenAI) -> str: # -- per-recorded-experience invariants ----------------------------------- def _assert_recorded_experience(self, exp: Experience, record_key: str): - self.assertEqual(exp.info.get("record_key"), record_key) - self.assertIsNotNone(exp.info.get("request_id")) + self.assertEqual(get_record_key(exp), record_key) + self.assertTrue(exp.eid.suffix) self.assertEqual(exp.info.get("rank"), 0) self.assertEqual(exp.info.get("model_version"), self.expected_model_version) self.assertGreater(len(exp.tokens), exp.prompt_length) # type: ignore [arg-type] @@ -1973,7 +1974,7 @@ async def test_record(self): # noqa: C901 task = await self._get_record_experiences(rk_gen) self.assertEqual(len(task["experiences"]), 1) # blob endpoint round-trips a full experience - request_id = task["experiences"][0]["info"]["request_id"] + request_id = task["experiences"][0]["eid"]["suffix"] blob_exp = await self._get_request_experience(rk_gen, request_id) self._assert_recorded_experience(blob_exp, rk_gen) self._assert_recorded_routed_experts(blob_exp) @@ -1995,10 +1996,10 @@ async def test_record(self): # noqa: C901 await self._flush() task = await self._get_record_experiences(rk_chat) self.assertEqual(len(task["experiences"]), 2) - # n=2 of one engine request -> two completions sharing one request_id, - # distinguished by sample_index. + # n=2 of one engine request -> two completions distinguished by + # sample_index and a sample-qualified EID suffix. self.assertEqual(sorted(t["info"]["sample_index"] for t in task["experiences"]), [0, 1]) - self.assertEqual(len({t["info"]["request_id"] for t in task["experiences"]}), 1) + self.assertEqual(len({t["eid"]["suffix"] for t in task["experiences"]}), 2) consumed = await self._consume(rk_chat, reward=0.8, run=2, task="t_chat") self.assertEqual(len(consumed), 2) for exp in consumed: diff --git a/trinity/buffer/store/__init__.py b/trinity/buffer/store/__init__.py index 965388b2d65..a96e83a8edc 100644 --- a/trinity/buffer/store/__init__.py +++ b/trinity/buffer/store/__init__.py @@ -1,18 +1,16 @@ from trinity.buffer.store.base_store import BaseStore, RecordStore from trinity.buffer.store.memory_store import ( - RECORD_KEY_INFO_KEY, - REQUEST_ID_INFO_KEY, MemoryStore, default_sample_id_getter, get_record_key, + parse_record_key, ) __all__ = [ "BaseStore", "MemoryStore", - "RECORD_KEY_INFO_KEY", - "REQUEST_ID_INFO_KEY", "RecordStore", "default_sample_id_getter", "get_record_key", + "parse_record_key", ] diff --git a/trinity/buffer/store/memory_store.py b/trinity/buffer/store/memory_store.py index fe137c92343..028de161499 100644 --- a/trinity/buffer/store/memory_store.py +++ b/trinity/buffer/store/memory_store.py @@ -6,12 +6,27 @@ from trinity.buffer.store.base_store import BaseStore from trinity.common.experience import Experience -REQUEST_ID_INFO_KEY = "request_id" -RECORD_KEY_INFO_KEY = "record_key" - SampleIdGetter = Callable[[Experience], str] +def parse_record_key(key: str) -> tuple[str, str, int]: + """Parse a complete ``//`` store key.""" + parts = key.split("/") + if len(parts) != 3 or any(part == "" for part in parts): + raise ValueError( + f"Store key must be complete '//', got '{key}'." + ) + batch, task, run_text = parts + try: + run = int(run_text) + except ValueError as exc: + raise ValueError( + f"Store key run_id must be an integer in '//', " + f"got '{key}'." + ) from exc + return batch, task, run + + def default_sample_id_getter(exp: Experience) -> str: """Resolve a stable sample id for an experience.""" info = exp.info or {} @@ -19,22 +34,15 @@ def default_sample_id_getter(exp: Experience) -> str: if sample_id is not None: return str(sample_id) - request_id = info.get("request_id") sample_index = info.get("sample_index") - if request_id is not None and sample_index is not None: - return f"{request_id}:{sample_index}" - if request_id is not None: - return str(request_id) + if sample_index is not None: + return f"{exp.eid.suffix}:{sample_index}" return exp.eid.uid def get_record_key(exp: Experience) -> str: """Return the complete store key stamped on an experience.""" - info = exp.info or {} - record_key = info.get(RECORD_KEY_INFO_KEY) - if record_key: - return record_key if exp.eid.batch != "" and exp.eid.task != "": return exp.eid.rid return exp.eid.suffix @@ -57,7 +65,7 @@ def __len__(self) -> int: return sum(len(exps) for exps in self._records.values()) def add(self, key: str, exps: List[Experience]) -> None: - self._validate_complete_key(key) + self._parse_complete_key(key) # validate key format if not exps: return @@ -73,12 +81,12 @@ def add(self, key: str, exps: List[Experience]) -> None: self._sample_to_key[sample_id] = key def overwrite(self, key: str, exps: List[Experience]) -> None: - self._validate_complete_key(key) + self._parse_complete_key(key) # validate key format self._drop_key(key) self.add(key, exps) def replace(self, key: str, old_sample_id: str, exp: Experience) -> None: - self._validate_complete_key(key) + self._parse_complete_key(key) # validate key format records = self._records.get(key) if records is None: raise KeyError(f"Key '{key}' does not exist.") @@ -111,12 +119,10 @@ def update( info: dict | None, sample_ids: List[str] | None, ) -> None: - self._validate_complete_key(key) + batch, task, run = self._parse_complete_key(key) # validate key format records = self._records.get(key) if records is None: raise KeyError(f"Key '{key}' does not exist.") - - batch, task, run = self._parse_complete_key(key) target_ids: Iterable[str] = list(records.keys()) if sample_ids is None else sample_ids for sample_id in target_ids: if sample_id not in records: @@ -146,26 +152,10 @@ def remove(self, key: str) -> List[Experience]: def keys(self) -> list[str]: return list(self._records.keys()) - @staticmethod - def _validate_complete_key(key: str) -> None: - MemoryStore._parse_complete_key(key) - @staticmethod def _parse_complete_key(key: str) -> tuple[str, str, int]: - parts = key.split("/") - if len(parts) != 3 or any(part == "" for part in parts): - raise ValueError( - f"Store key must be complete '//', got '{key}'." - ) - batch, task, run_text = parts - try: - run = int(run_text) - except ValueError as exc: - raise ValueError( - f"Store key run_id must be an integer in '//', " - f"got '{key}'." - ) from exc - return batch, task, run + """Parse a complete store key; also usable as a key-format validator.""" + return parse_record_key(key) def _matching_keys(self, key: str) -> list[str]: if key == "": diff --git a/trinity/common/models/recording/merger.py b/trinity/common/models/recording/merger.py index 19b1ca89761..569dcdc03f0 100644 --- a/trinity/common/models/recording/merger.py +++ b/trinity/common/models/recording/merger.py @@ -14,7 +14,7 @@ class PrefixExperienceMerger: Strategy: * Experiences are grouped by record key and a best-effort sample stream - key (sample_index, explicit sample_id, request_id suffix, then default). + key (sample_index, explicit sample_id, then default). * Each stream tracks one latest/longest head experience. A new experience merges only when the head tokens are a strict prefix of the new tokens. * If no head exists yet, the store is scanned once to find the longest @@ -96,12 +96,6 @@ def _sample_stream_key(exp: Experience) -> tuple[str, Any]: if sample_id is not None: return ("sample_id", sample_id) - request_id = info.get("request_id") - if isinstance(request_id, str): - _, sep, suffix = request_id.rpartition(":") - if sep and suffix.isdigit(): - return ("request_id_sample_index", int(suffix)) - return ("default", 0) @@ -236,15 +230,11 @@ def _response_routed_experts(exp: Experience) -> Optional[torch.Tensor]: def _merge_info(prefix_exp: Experience, final_exp: Experience) -> dict: info = dict(final_exp.info or {}) - merged_request_ids = list((prefix_exp.info or {}).get("merged_request_ids") or []) - prefix_request_id = (prefix_exp.info or {}).get("request_id") - if prefix_request_id is not None and prefix_request_id not in merged_request_ids: - merged_request_ids.append(prefix_request_id) - final_request_id = (final_exp.info or {}).get("request_id") - if final_request_id is not None and final_request_id not in merged_request_ids: - merged_request_ids.append(final_request_id) - if merged_request_ids: - info["merged_request_ids"] = merged_request_ids + merged_eid_suffixes = list((prefix_exp.info or {}).get("merged_eid_suffixes") or []) + for suffix in (prefix_exp.eid.suffix, final_exp.eid.suffix): + if suffix not in merged_eid_suffixes: + merged_eid_suffixes.append(suffix) + info["merged_eid_suffixes"] = merged_eid_suffixes merged_sample_ids = list((prefix_exp.info or {}).get("merged_sample_ids") or []) for sample_id in (default_sample_id_getter(prefix_exp), default_sample_id_getter(final_exp)): diff --git a/trinity/common/models/recording/query.py b/trinity/common/models/recording/query.py index 373658e1ede..af2c437e83d 100644 --- a/trinity/common/models/recording/query.py +++ b/trinity/common/models/recording/query.py @@ -5,7 +5,7 @@ from fastapi import APIRouter, HTTPException, Request, Response from pydantic import BaseModel -from trinity.buffer.store import REQUEST_ID_INFO_KEY, RecordStore +from trinity.buffer.store import RecordStore from trinity.common.experience import Experience from trinity.common.models.recording.recorder import ( TRINITY_RECORD_STORE_ATTR, @@ -44,8 +44,7 @@ def _recorder(request: Request) -> Recorder: def _get_exp(store: RecordStore, record_key: str, request_id: str) -> Experience: for exp in store.get(record_key): - info = exp.info or {} - if info.get(REQUEST_ID_INFO_KEY) == request_id: + if exp.eid.suffix == request_id: return exp raise HTTPException(status_code=404, detail="experience not found") @@ -94,8 +93,7 @@ async def delete_request_experience(record_key: str, request_id: str, request: R kept = [] deleted = False for exp in store.get(record_key): - info = exp.info or {} - if info.get(REQUEST_ID_INFO_KEY) == request_id: + if exp.eid.suffix == request_id: deleted = True else: kept.append(exp) diff --git a/trinity/common/models/recording/recorder.py b/trinity/common/models/recording/recorder.py index 2c9e02cab9f..b9eacdb780f 100644 --- a/trinity/common/models/recording/recorder.py +++ b/trinity/common/models/recording/recorder.py @@ -110,7 +110,7 @@ async def _safe_append(self, exp: Experience) -> None: except Exception: logging.getLogger(__name__).exception( "recording store write failed for request %s", - exp.info.get("request_id"), + exp.eid.suffix, ) def forget_record(self, record_key: str) -> None: diff --git a/trinity/common/models/sglang_patch/recording/models.py b/trinity/common/models/sglang_patch/recording/models.py index 88093f192ff..ce63c4f9ce1 100644 --- a/trinity/common/models/sglang_patch/recording/models.py +++ b/trinity/common/models/sglang_patch/recording/models.py @@ -12,9 +12,8 @@ from the request. Field mapping (SGLang ``ret`` -> ``Experience``): - meta_info.id -> eid.suffix (traceability; batch/task/run and reward - are assigned from record key by ``MemoryStore.update``) - record_key -> info["record_key"] (the MemoryStore group key) + meta_info.id -> eid.suffix (traceability) + record_key -> eid.batch/task/run (the MemoryStore group key) sample index -> info["sample_index"] (position within the n set) prompt_token_ids -> tokens (prompt) + prompt_length output_ids -> tokens (response) @@ -29,6 +28,7 @@ import torch +from trinity.buffer.store import parse_record_key from trinity.common.experience import EID, Experience from trinity.common.models.sglang_model import decode_sglang_routed_experts @@ -43,6 +43,12 @@ def _extract_output_logprobs(meta_info: dict) -> List[float]: return [float(logprob) for logprob, *_ in output_token_logprobs] +def _sample_suffix(request_id: str, sample_index: int, num_samples: int) -> str: + if num_samples <= 1: + return request_id + return f"{request_id}:{sample_index}" + + def _extract_routed_experts( routed_experts_value: Any, total_tokens: int, @@ -74,8 +80,9 @@ def build_sglang_experience( """Build Trinity ``Experience`` objects from a finished SGLang ``ret``. One experience per output (``n > 1`` / batch is captured in full). Each - shares ``eid.suffix = meta_info.id`` and ``info["record_key"] = record_key``; - ``info["sample_index"]`` distinguishes samples within the group. + carries ``record_key`` in ``eid.batch/task/run`` and shares + ``eid.suffix = meta_info.id``; ``info["sample_index"]`` distinguishes + samples within the group. Args: ret: A finished SGLang result — a dict, or a list of dicts for ``n > 1`` @@ -145,10 +152,13 @@ def build_sglang_experience( if resolved_model_version is None: resolved_model_version = model_version - eid = EID(suffix=request_id) + suffix = _sample_suffix(request_id, sample_index, len(ret_list)) + if record_key is None: + eid = EID(suffix=suffix) + else: + batch, task, run = parse_record_key(record_key) + eid = EID(batch=batch, task=task, run=run, suffix=suffix) info = { - "request_id": request_id, - "record_key": record_key, "sample_index": sample_index, "rank": rank, "timestamp": timestamp, diff --git a/trinity/common/models/tinker_model.py b/trinity/common/models/tinker_model.py index 113d1dda91c..5ca375b90f4 100644 --- a/trinity/common/models/tinker_model.py +++ b/trinity/common/models/tinker_model.py @@ -8,7 +8,7 @@ from tinker import types from torch import Tensor -from trinity.buffer.store import MemoryStore +from trinity.buffer.store import MemoryStore, parse_record_key from trinity.common.config import InferenceModelConfig from trinity.common.constants import SyncMethod from trinity.common.experience import Experience @@ -26,11 +26,14 @@ def _build_tinker_experiences( model_version: Optional[int] = None, request_id: str, ) -> Sequence[Experience]: + batch, task, run = parse_record_key(record_key) for index, exp in enumerate(experiences): + exp.eid.batch = batch + exp.eid.task = task + exp.eid.run = run + exp.eid.suffix = f"{request_id}:{index}" if exp.info is None: exp.info = {} - exp.info["record_key"] = record_key - exp.info["request_id"] = f"{request_id}:{index}" exp.info["rank"] = rank exp.info["timestamp"] = timestamp if model_version is not None: diff --git a/trinity/common/models/vllm_model.py b/trinity/common/models/vllm_model.py index 58113522142..3386778694c 100644 --- a/trinity/common/models/vllm_model.py +++ b/trinity/common/models/vllm_model.py @@ -10,7 +10,7 @@ from packaging.version import parse as parse_version from transformers import AutoProcessor -from trinity.buffer.store import RECORD_KEY_INFO_KEY, REQUEST_ID_INFO_KEY +from trinity.buffer.store import parse_record_key from trinity.common.config import InferenceModelConfig from trinity.common.constants import SyncMethod from trinity.common.experience import Experience @@ -230,7 +230,7 @@ async def chat( record_key (Optional[str]): Recording identity for the in-vLLM recorder (the MemoryStore group key). Propagated to ``generate`` via ``record_key_ctx`` so the recorder stamps it - into ``info["record_key"]`` without an HTTP hop. None skips + into ``Experience.eid`` without an HTTP hop. None skips recording. kwargs (dict): A dictionary of sampling parameters. @@ -293,9 +293,11 @@ async def generate( # experiences must still be tracked for history extraction and # the buffer/trainer (they are popped by record_key on consume). if self.recorder is not None and record_key is not None: + batch, task, run = parse_record_key(record_key) for exp in returned_seq: - exp.info[RECORD_KEY_INFO_KEY] = record_key - exp.info[REQUEST_ID_INFO_KEY] = exp.eid.suffix + exp.eid.batch = batch + exp.eid.task = task + exp.eid.run = run exp.info["rank"] = self.recorder.rank exp.info["model_version"] = self.model_version self.recorder.store.add(record_key, [exp]) diff --git a/trinity/common/models/vllm_patch/recording/models.py b/trinity/common/models/vllm_patch/recording/models.py index fe4a46b365b..a5b7ace1fda 100644 --- a/trinity/common/models/vllm_patch/recording/models.py +++ b/trinity/common/models/vllm_patch/recording/models.py @@ -13,7 +13,7 @@ id == the OpenAI ``response.id``. Kept for traceability; ``eid.batch``/``task``/``run`` and reward are assigned from record key by ``MemoryStore.update`` at consume time.) - API key / record key -> info["record_key"] (the recording identity; **the + API key / record key -> eid.batch/task/run (the recording identity; **the group key** the MemoryStore batches experiences by, so a whole reward unit's samples/turns are reward-updated and consumed together.) @@ -30,8 +30,8 @@ model_version -> info["model_version"] (which checkpoint policy served the turn; read in-actor by the recorder's provider) -Plus bookkeeping (request_id / record_key / sample_index / rank / timestamp / -endpoint / model_version) stashed in ``Experience.info`` so it round-trips +Plus bookkeeping (sample_index / rank / timestamp / endpoint / model_version) +stashed in ``Experience.info`` so it round-trips with the experience through serialize/deserialize. """ @@ -39,6 +39,7 @@ import torch +from trinity.buffer.store import parse_record_key from trinity.common.experience import EID, Experience from trinity.common.models.mm_utils import combine_output_token_ids @@ -81,6 +82,12 @@ def _extract_chosen_logprobs( return [float(sample_logprobs[i][tid].logprob) for i, tid in enumerate(response_token_ids)] +def _sample_suffix(request_id: str, sample_index: int, num_samples: int) -> str: + if num_samples <= 1: + return request_id + return f"{request_id}:{sample_index}" + + def _extract_routed_experts( output: Any, completion: Any, @@ -125,15 +132,14 @@ def build_experience( """Build Trinity ``Experience`` objects from a finished ``RequestOutput``. One experience per completion (``output.outputs``), so ``n > 1`` sampling - is captured in full. Each experience shares ``eid.suffix = request_id`` and - ``info["record_key"] = record_key`` (the group key); ``info["sample_index"]`` - distinguishes samples within the group. + is captured in full. Each experience carries ``record_key`` in + ``eid.batch/task/run`` and shares ``eid.suffix = request_id``; + ``info["sample_index"]`` distinguishes samples within the group. Args: output: A ``RequestOutput`` with ``finished == True``. record_key: The recording identity (API key / Ray-injected record key); - stored in ``info["record_key"]`` and used as the MemoryStore group - key. + stored in ``eid.batch/task/run`` and used as the MemoryStore group key. rank: Data-parallel serving rank. timestamp: UTC ISO-8601 string (caller-stamped to keep this pure). endpoint: Which OpenAI endpoint served the turn (best-effort). @@ -158,8 +164,8 @@ def build_experience( request had no prompt or no completion with response tokens. """ request_id = output.request_id - # eid.suffix = request_id for traceability; batch/task/run and reward are - # assigned from record_key by MemoryStore.update at consume. + # eid.suffix = request_id for traceability; batch/task/run are assigned + # from record_key when recording metadata is requested. prompt_token_ids = list(output.prompt_token_ids or []) if not prompt_token_ids: @@ -191,10 +197,13 @@ def build_experience( info = None eid = None if include_recording_info: - eid = EID(suffix=request_id) + suffix = _sample_suffix(request_id, sample_index, len(completions)) + if record_key is None: + eid = EID(suffix=suffix) + else: + batch, task, run = parse_record_key(record_key) + eid = EID(batch=batch, task=task, run=run, suffix=suffix) info = { - "request_id": request_id, - "record_key": record_key, "sample_index": sample_index, "rank": rank, "timestamp": timestamp, diff --git a/trinity/common/workflows/workflow.py b/trinity/common/workflows/workflow.py index e4cd706ac97..eb795114375 100644 --- a/trinity/common/workflows/workflow.py +++ b/trinity/common/workflows/workflow.py @@ -6,6 +6,7 @@ from dataclasses import asdict, dataclass, field from typing import TYPE_CHECKING, List, Optional, Type, Union +from trinity.buffer.store import parse_record_key from trinity.common.config import FormatConfig, GenerationConfig from trinity.common.experience import Experience from trinity.common.rewards.reward_fn import RewardFn @@ -333,10 +334,11 @@ def _attach_rewards(self, responses, *, base: int) -> List[Experience]: @staticmethod def _stamp_record_key(exps: List[Experience], record_key: str) -> None: + batch, task, run = parse_record_key(record_key) for exp in exps: - if exp.info is None: - exp.info = {} - exp.info["record_key"] = record_key + exp.eid.batch = batch + exp.eid.task = task + exp.eid.run = run class SimpleWorkflow(BaseSimpleWorkflow): diff --git a/trinity/explorer/workflow_runner.py b/trinity/explorer/workflow_runner.py index fe8519b7bfa..b00d59cad73 100644 --- a/trinity/explorer/workflow_runner.py +++ b/trinity/explorer/workflow_runner.py @@ -10,6 +10,7 @@ from typing import Dict, List, Optional, Tuple from trinity.buffer import get_buffer_reader, get_buffer_writer +from trinity.buffer.store import get_record_key from trinity.common.config import Config, StorageConfig from trinity.common.constants import LOG_DIR_ENV_VAR, LOG_LEVEL_ENV_VAR from trinity.common.experience import Experience @@ -473,7 +474,7 @@ async def run_task( # by the coordinator at finalize time. updates = [ { - "record_key": exp.info.get("record_key") or exp.eid.suffix, + "record_key": get_record_key(exp), "reward": exp.reward, "run": exp.eid.run, "task": str(task.task_id), From b4d916d88ebde7b7d8aa83e40d9285c0898f6122 Mon Sep 17 00:00:00 2001 From: "panxuchen.pxc" Date: Mon, 29 Jun 2026 14:59:45 +0800 Subject: [PATCH 33/84] optimize memory store --- tests/buffer/memory_store_test.py | 2 +- tests/common/recording_recorder_test.py | 23 +++----- trinity/buffer/store/__init__.py | 7 ++- trinity/buffer/store/base_store.py | 5 +- trinity/buffer/store/memory_store.py | 66 ++++++++++++++--------- trinity/common/models/recording/merger.py | 12 ++--- 6 files changed, 57 insertions(+), 58 deletions(-) diff --git a/tests/buffer/memory_store_test.py b/tests/buffer/memory_store_test.py index 8470c210ce1..d86187b189a 100644 --- a/tests/buffer/memory_store_test.py +++ b/tests/buffer/memory_store_test.py @@ -11,7 +11,7 @@ def get_dummy_experience(num: int, request_id: str | None = None): request_id = request_id or uuid.uuid4().hex[:6] return [ Experience( - eid=EID(suffix=request_id), + eid=EID(suffix=request_id if num == 1 else f"{request_id}:{i}"), tokens=torch.zeros(5), prompt_length=2, info={ diff --git a/tests/common/recording_recorder_test.py b/tests/common/recording_recorder_test.py index 362d4c2a5d1..0572f9bb0e7 100644 --- a/tests/common/recording_recorder_test.py +++ b/tests/common/recording_recorder_test.py @@ -14,13 +14,10 @@ def make_turn( tokens: list[int], prompt_length: int, logprobs: list[float], - sample_id: str | None = None, sample_index: int = 0, ) -> Experience: batch, task, run = parse_record_key(record_key) info = {"sample_index": sample_index} - if sample_id is not None: - info["sample_id"] = sample_id return Experience( eid=EID(batch=batch, task=task, run=run, suffix=request_id), tokens=tokens, @@ -45,7 +42,6 @@ async def test_prefix_experiences_merge_and_keep_final_sample_id(self): tokens=[10, 11, 20, 21], prompt_length=2, logprobs=[-0.2, -0.3], - sample_id="sample-old", ) second = make_turn( request_id="req-2", @@ -53,7 +49,6 @@ async def test_prefix_experiences_merge_and_keep_final_sample_id(self): tokens=[10, 11, 20, 21, 12, 13, 30, 31, 32], prompt_length=6, logprobs=[-0.4, -0.5, -0.6], - sample_id="sample-final", ) await recorder._safe_append(first) @@ -62,7 +57,6 @@ async def test_prefix_experiences_merge_and_keep_final_sample_id(self): recorded = store.get(record_key) self.assertEqual(len(recorded), 1) merged = recorded[0] - self.assertEqual(merged.info["sample_id"], "sample-final") self.assertEqual(merged.eid.suffix, "req-2") self.assertEqual(merged.prompt_length, 2) self.assertTrue(torch.equal(merged.tokens, second.tokens)) @@ -79,12 +73,12 @@ async def test_prefix_experiences_merge_and_keep_final_sample_id(self): ) ) self.assertEqual(merged.info["merged_eid_suffixes"], ["req-1", "req-2"]) - self.assertEqual(merged.info["merged_sample_ids"], ["sample-old", "sample-final"]) + self.assertEqual(merged.info["merged_sample_ids"], ["req-1", "req-2"]) - store.update(record_key, reward=1.0, info=None, sample_ids=["sample-final"]) + store.update(record_key, reward=1.0, info=None, sample_ids=["req-2"]) self.assertEqual(store.get(record_key)[0].reward, 1.0) with self.assertRaises(KeyError): - store.update(record_key, reward=2.0, info=None, sample_ids=["sample-old"]) + store.update(record_key, reward=2.0, info=None, sample_ids=["req-1"]) async def test_non_prefix_experiences_do_not_merge(self): store = MemoryStore() @@ -130,7 +124,6 @@ async def test_merge_head_replaces_only_matching_sample_stream(self): tokens=[10, 11, 20], prompt_length=2, logprobs=[-0.2], - sample_id="sample-zero", sample_index=0, ) sample_one_first = make_turn( @@ -139,7 +132,6 @@ async def test_merge_head_replaces_only_matching_sample_stream(self): tokens=[10, 11, 21], prompt_length=2, logprobs=[-0.3], - sample_id="sample-one-old", sample_index=1, ) sample_one_final = make_turn( @@ -148,7 +140,6 @@ async def test_merge_head_replaces_only_matching_sample_stream(self): tokens=[10, 11, 21, 12, 31], prompt_length=4, logprobs=[-0.4], - sample_id="sample-one-final", sample_index=1, ) @@ -158,8 +149,8 @@ async def test_merge_head_replaces_only_matching_sample_stream(self): recorded = store.get(record_key) self.assertEqual(len(recorded), 2) - self.assertEqual(recorded[0].info["sample_id"], "sample-zero") - self.assertEqual(recorded[1].info["sample_id"], "sample-one-final") + self.assertEqual(recorded[0].eid.suffix, "req-1") + self.assertEqual(recorded[1].eid.suffix, "req-3") self.assertTrue( torch.equal( recorded[1].action_mask, @@ -181,7 +172,6 @@ async def test_stale_merge_head_falls_back_to_append(self): tokens=[10, 11, 20], prompt_length=2, logprobs=[-0.2], - sample_id="sample-old", ) second = make_turn( request_id="req-2", @@ -189,7 +179,6 @@ async def test_stale_merge_head_falls_back_to_append(self): tokens=[10, 11, 20, 12, 30], prompt_length=4, logprobs=[-0.3], - sample_id="sample-final", ) await recorder._safe_append(first) @@ -198,7 +187,7 @@ async def test_stale_merge_head_falls_back_to_append(self): recorded = store.get(record_key) self.assertEqual(len(recorded), 1) - self.assertEqual(recorded[0].info["sample_id"], "sample-final") + self.assertEqual(recorded[0].eid.suffix, "req-2") self.assertEqual(recorded[0].prompt_length, 4) diff --git a/trinity/buffer/store/__init__.py b/trinity/buffer/store/__init__.py index a96e83a8edc..78d60297c1e 100644 --- a/trinity/buffer/store/__init__.py +++ b/trinity/buffer/store/__init__.py @@ -1,16 +1,15 @@ -from trinity.buffer.store.base_store import BaseStore, RecordStore +from trinity.buffer.store.base_store import RecordStore from trinity.buffer.store.memory_store import ( MemoryStore, - default_sample_id_getter, get_record_key, + get_sample_id, parse_record_key, ) __all__ = [ - "BaseStore", "MemoryStore", "RecordStore", - "default_sample_id_getter", "get_record_key", + "get_sample_id", "parse_record_key", ] diff --git a/trinity/buffer/store/base_store.py b/trinity/buffer/store/base_store.py index 570b9416dff..d8fa635ed60 100644 --- a/trinity/buffer/store/base_store.py +++ b/trinity/buffer/store/base_store.py @@ -4,7 +4,7 @@ from trinity.common.experience import Experience -class BaseStore(ABC): +class RecordStore(ABC): """Abstract base class for an in-process experience store. The key follows the format ``//`` and each @@ -40,6 +40,3 @@ def remove(self, key: str) -> List[Experience]: @abstractmethod def keys(self) -> list[str]: """Return complete keys currently stored in insertion order.""" - - -RecordStore = BaseStore diff --git a/trinity/buffer/store/memory_store.py b/trinity/buffer/store/memory_store.py index 028de161499..2835c2ce1f9 100644 --- a/trinity/buffer/store/memory_store.py +++ b/trinity/buffer/store/memory_store.py @@ -1,13 +1,11 @@ """In-memory implementation of the experience store interface.""" from collections import OrderedDict -from typing import Callable, Iterable, List +from typing import Iterable, List -from trinity.buffer.store.base_store import BaseStore +from trinity.buffer.store.base_store import RecordStore from trinity.common.experience import Experience -SampleIdGetter = Callable[[Experience], str] - def parse_record_key(key: str) -> tuple[str, str, int]: """Parse a complete ``//`` store key.""" @@ -27,18 +25,9 @@ def parse_record_key(key: str) -> tuple[str, str, int]: return batch, task, run -def default_sample_id_getter(exp: Experience) -> str: - """Resolve a stable sample id for an experience.""" - info = exp.info or {} - sample_id = info.get("sample_id") - if sample_id is not None: - return str(sample_id) - - sample_index = info.get("sample_index") - if sample_index is not None: - return f"{exp.eid.suffix}:{sample_index}" - - return exp.eid.uid +def get_sample_id(exp: Experience) -> str: + """Return the short sample id used by ``MemoryStore``.""" + return exp.eid.suffix def get_record_key(exp: Experience) -> str: @@ -48,7 +37,7 @@ def get_record_key(exp: Experience) -> str: return exp.eid.suffix -class MemoryStore(BaseStore): +class MemoryStore(RecordStore): """A fast in-process store backed by Python dictionaries. ``add``, ``overwrite`` and ``update`` require complete keys in the form @@ -56,22 +45,26 @@ class MemoryStore(BaseStore): so callers can drain a batch or task at once. """ - def __init__(self, sample_id_getter: SampleIdGetter | None = None) -> None: - self.sample_id_getter = sample_id_getter or default_sample_id_getter + def __init__(self) -> None: + # main storage of experiences, keyed by complete store key and sample_id self._records: dict[str, OrderedDict[str, Experience]] = {} + # extra indices to support prefix-based lookups in get() and remove() + self._batch_keys: dict[str, OrderedDict[str, None]] = {} + self._task_keys: dict[tuple[str, str], OrderedDict[str, None]] = {} self._sample_to_key: dict[str, str] = {} def __len__(self) -> int: return sum(len(exps) for exps in self._records.values()) def add(self, key: str, exps: List[Experience]) -> None: - self._parse_complete_key(key) # validate key format + batch, task, _ = self._parse_complete_key(key) # validate key format if not exps: return records = self._records.setdefault(key, OrderedDict()) + self._index_key(batch, task, key) for exp in exps: - sample_id = self.sample_id_getter(exp) + sample_id = get_sample_id(exp) owner_key = self._sample_to_key.get(sample_id) if owner_key is not None: raise ValueError( @@ -93,7 +86,7 @@ def replace(self, key: str, old_sample_id: str, exp: Experience) -> None: if old_sample_id not in records: raise KeyError(f"sample_id '{old_sample_id}' does not exist under key '{key}'.") - new_sample_id = self.sample_id_getter(exp) + new_sample_id = get_sample_id(exp) owner_key = self._sample_to_key.get(new_sample_id) if owner_key is not None and (owner_key != key or new_sample_id != old_sample_id): raise ValueError( @@ -162,13 +155,38 @@ def _matching_keys(self, key: str) -> list[str]: return list(self._records.keys()) if key in self._records: return [key] - prefix = key + "/" - return [record_key for record_key in self._records if record_key.startswith(prefix)] + + parts = key.split("/") + if len(parts) == 1 and parts[0] != "": + return list(self._batch_keys.get(parts[0], ())) + if len(parts) == 2 and parts[0] != "" and parts[1] != "": + return list(self._task_keys.get((parts[0], parts[1]), ())) + return [] def _drop_key(self, key: str) -> list[Experience]: records = self._records.pop(key, None) if records is None: return [] + batch, task, _ = self._parse_complete_key(key) + self._unindex_key(batch, task, key) for sample_id in records: self._sample_to_key.pop(sample_id, None) return list(records.values()) + + def _index_key(self, batch: str, task: str, key: str) -> None: + self._batch_keys.setdefault(batch, OrderedDict())[key] = None + self._task_keys.setdefault((batch, task), OrderedDict())[key] = None + + def _unindex_key(self, batch: str, task: str, key: str) -> None: + batch_keys = self._batch_keys.get(batch) + if batch_keys is not None: + batch_keys.pop(key, None) + if not batch_keys: + self._batch_keys.pop(batch, None) + + task_key = (batch, task) + task_keys = self._task_keys.get(task_key) + if task_keys is not None: + task_keys.pop(key, None) + if not task_keys: + self._task_keys.pop(task_key, None) diff --git a/trinity/common/models/recording/merger.py b/trinity/common/models/recording/merger.py index 569dcdc03f0..cf3bc4295d8 100644 --- a/trinity/common/models/recording/merger.py +++ b/trinity/common/models/recording/merger.py @@ -5,7 +5,7 @@ import torch -from trinity.buffer.store import RecordStore, default_sample_id_getter +from trinity.buffer.store import RecordStore, get_sample_id from trinity.common.experience import Experience @@ -14,7 +14,7 @@ class PrefixExperienceMerger: Strategy: * Experiences are grouped by record key and a best-effort sample stream - key (sample_index, explicit sample_id, then default). + key (sample_index, then default). * Each stream tracks one latest/longest head experience. A new experience merges only when the head tokens are a strict prefix of the new tokens. * If no head exists yet, the store is scanned once to find the longest @@ -47,7 +47,7 @@ def try_merge(self, record_key: str, exp: Experience) -> bool: if candidate is None: return False - old_sample_id = default_sample_id_getter(candidate) + old_sample_id = get_sample_id(candidate) merged = _merge_prefix_experiences(candidate, exp) try: self.store.replace(record_key, old_sample_id, merged) @@ -92,10 +92,6 @@ def _sample_stream_key(exp: Experience) -> tuple[str, Any]: if sample_index is not None: return ("sample_index", sample_index) - sample_id = info.get("sample_id") - if sample_id is not None: - return ("sample_id", sample_id) - return ("default", 0) @@ -237,7 +233,7 @@ def _merge_info(prefix_exp: Experience, final_exp: Experience) -> dict: info["merged_eid_suffixes"] = merged_eid_suffixes merged_sample_ids = list((prefix_exp.info or {}).get("merged_sample_ids") or []) - for sample_id in (default_sample_id_getter(prefix_exp), default_sample_id_getter(final_exp)): + for sample_id in (get_sample_id(prefix_exp), get_sample_id(final_exp)): if sample_id not in merged_sample_ids: merged_sample_ids.append(sample_id) info["merged_sample_ids"] = merged_sample_ids From 6d47fb418d79cd044e9c6d014383b12c93775378 Mon Sep 17 00:00:00 2001 From: pxc Date: Mon, 29 Jun 2026 15:15:47 +0800 Subject: [PATCH 34/84] simplify doc --- trinity/common/models/vllm_patch/recording/recorder.py | 7 +------ trinity/common/models/vllm_patch/recording/server.py | 6 +----- 2 files changed, 2 insertions(+), 11 deletions(-) diff --git a/trinity/common/models/vllm_patch/recording/recorder.py b/trinity/common/models/vllm_patch/recording/recorder.py index 3c3e8606fa8..7770ff7ff07 100644 --- a/trinity/common/models/vllm_patch/recording/recorder.py +++ b/trinity/common/models/vllm_patch/recording/recorder.py @@ -34,17 +34,12 @@ #: Guard attribute marking the wrapped generate, mirroring api_patch_v17 style. _PATCHED_FLAG = "__patched_engine_recording__" -#: Instance attribute on the AsyncLLM engine_client holding the current serving -#: checkpoint version. Mirrored by ``VLLMModel.sync_model_weights`` (and at -#: engine creation); read live here so each experience is attributed to the -#: right policy without a launch-time parameter. -_MODEL_VERSION_ATTR = "trinity_model_version" #: Force at least this many top-k logprobs per generated token so recording #: captures the chosen token's logprob even when the caller didn't request #: logprobs. We store ONLY the sampled token's logprob, and vLLM force-includes #: the sampled token at ``logprobs=1``, so 1 is the only useful value — no need #: to thread a knob through the launcher. The engine's ``max_logprobs`` cap -#: (default 20, set at engine build) already covers it. +#: (default 1, set at engine build) already covers it. _RECORDER_LOGPROB_WIDTH = 1 TRINITY_MM_RENDER_ATTR = "trinity_mm_render" diff --git a/trinity/common/models/vllm_patch/recording/server.py b/trinity/common/models/vllm_patch/recording/server.py index 4e8f4446b09..890d2be9698 100644 --- a/trinity/common/models/vllm_patch/recording/server.py +++ b/trinity/common/models/vllm_patch/recording/server.py @@ -262,11 +262,7 @@ async def run_api_server_with_recording( # ``--max-logprobs`` here would be inert (init_app_state does not read them). # ``enable_return_routed_experts`` is opt-in via ``enable_router_replay`` # (mirrored in ``config_validator``); it is not implied by recording, so a - # dense model records history with ``routed_experts=None``. The engine's - # default ``max_logprobs=20`` covers the recorder's top-k - # (``VLLM_RECORD_TOPK``, default 1). To record routed_experts, the engine - # must be built with that flag on — the launcher is responsible for that, - # not these CLI args. + # dense model records history with ``routed_experts=None``. args = parser.parse_args(cli_args) args.structured_outputs_config.reasoning_parser = reasoning_parser From 464e566d8c26cfda08631df1447e1f721d3dc4b7 Mon Sep 17 00:00:00 2001 From: "panxuchen.pxc" Date: Mon, 29 Jun 2026 17:02:09 +0800 Subject: [PATCH 35/84] refactor workflow interface --- tests/explorer/scheduler_test.py | 66 +++-- trinity/common/config.py | 2 +- trinity/common/models/model.py | 86 ++++++ trinity/common/models/recording/query.py | 64 +++-- .../models/vllm_patch/recording/server.py | 2 +- trinity/common/workflows/__init__.py | 3 +- trinity/common/workflows/workflow.py | 65 ++++- trinity/explorer/proxy/service.py | 2 +- trinity/explorer/rollout_coordinator.py | 130 +++------ trinity/explorer/scheduler.py | 99 +++++-- trinity/explorer/workflow_runner.py | 257 +++++++----------- 11 files changed, 448 insertions(+), 328 deletions(-) diff --git a/tests/explorer/scheduler_test.py b/tests/explorer/scheduler_test.py index 88d9b3fc74c..0f7846f777d 100644 --- a/tests/explorer/scheduler_test.py +++ b/tests/explorer/scheduler_test.py @@ -381,6 +381,18 @@ def create_role_models(model_config, role, actor_cls) -> None: return actor_handles +def _resolve_rollout_actors(config) -> Dict[int, ray.actor.ActorHandle]: + allocator = Allocator(config.explorer) + rollout_config = config.explorer.rollout_model + return { + engine_id: ray.get_actor( + allocator.get_actor_name("rollout", engine_id, 0), + namespace=rollout_config.ray_namespace, + ) + for engine_id in range(rollout_config.engine_num) + } + + def _cleanup_named_model_actors(actor_handles: Optional[List]) -> None: if not actor_handles: return @@ -500,8 +512,11 @@ def setUp(self): self.config.check_and_update() self.model_actors = _create_named_model_actors(self.config) + def _create_scheduler(self) -> Scheduler: + return Scheduler(self.config, rollout_actors=_resolve_rollout_actors(self.config)) + async def test_get_payload_results(self): - scheduler = Scheduler(self.config) + scheduler = self._create_scheduler() await scheduler.start() tasks = generate_tasks(8) @@ -605,7 +620,7 @@ async def test_get_payload_results(self): await scheduler.stop() async def test_concurrent_operations(self): - scheduler = Scheduler(self.config) + scheduler = self._create_scheduler() await scheduler.start() async def schedule_tasks(batch_id, num_tasks): @@ -628,7 +643,7 @@ async def schedule_tasks(batch_id, num_tasks): await scheduler.stop() async def test_scheduler_restart_after_stop(self): - scheduler = Scheduler(self.config) + scheduler = self._create_scheduler() await scheduler.start() tasks = generate_tasks(2) @@ -649,7 +664,7 @@ async def test_scheduler_restart_after_stop(self): async def test_split_tasks(self): self.config.explorer.max_repeat_times_per_runner = 2 self.config.check_and_update() - scheduler = Scheduler(self.config) + scheduler = self._create_scheduler() await scheduler.start() exp_list = [] @@ -693,7 +708,7 @@ async def test_split_tasks(self): async def test_multi_step_execution(self): self.config.explorer.max_repeat_times_per_runner = 1 self.config.check_and_update() - scheduler = Scheduler(self.config) + scheduler = self._create_scheduler() await scheduler.start() tasks = generate_tasks(2, repeat_times=4) @@ -709,7 +724,7 @@ async def test_multi_step_execution(self): async def test_non_repeatable_workflow(self): self.config.explorer.max_repeat_times_per_runner = 2 self.config.check_and_update() - scheduler = Scheduler(self.config) + scheduler = self._create_scheduler() await scheduler.start() task_num, repeat_times = 5, 4 tasks = generate_tasks(task_num, repeat_times=repeat_times, repeatable=False) @@ -742,7 +757,7 @@ async def test_non_repeatable_workflow(self): async def test_async_workflow(self): self.config.explorer.max_repeat_times_per_runner = 2 self.config.check_and_update() - scheduler = Scheduler(self.config) + scheduler = self._create_scheduler() await scheduler.start() task_num, repeat_times, step_num = 5, 4, 3 tasks = [ @@ -778,7 +793,7 @@ async def test_stepwise_experience_eid(self): self.config.buffer.train_batch_size = task_num * repeat_times * step_num self.config.explorer.max_repeat_times_per_runner = 2 self.config.check_and_update() - scheduler = Scheduler(self.config) + scheduler = self._create_scheduler() await scheduler.start() batch_num = 2 @@ -831,7 +846,7 @@ async def test_stepwise_experience_eid(self): async def test_metric_calculation_with_repeatable_workflow(self, max_repeat_times_per_runner): self.config.explorer.max_repeat_times_per_runner = max_repeat_times_per_runner self.config.check_and_update() - scheduler = Scheduler(self.config) + scheduler = self._create_scheduler() await scheduler.start() tasks = [] tasks.extend(generate_tasks(total_num=1, step_num=1, repeat_times=4, repeatable=True)) @@ -855,7 +870,7 @@ async def test_metric_calculation_with_non_repeatable_workflow( ): self.config.explorer.max_repeat_times_per_runner = max_repeat_times_per_runner self.config.check_and_update() - scheduler = Scheduler(self.config) + scheduler = self._create_scheduler() await scheduler.start() tasks = [] tasks.extend(generate_tasks(total_num=1, step_num=3, repeat_times=4, repeatable=False)) @@ -879,7 +894,7 @@ async def test_over_rollout_min_wait(self): self.config.buffer.batch_size = 4 self.config.synchronizer.sync_style = SyncStyle.EXPLORER_DRIVEN self.config.check_and_update() - scheduler = Scheduler(self.config) + scheduler = self._create_scheduler() await scheduler.start() tasks = [] tasks.extend(generate_tasks(0, timeout_num=2, repeat_times=1, timeout_seconds=1)) @@ -898,7 +913,7 @@ async def test_over_rollout_return_partial_tasks(self): self.config.synchronizer.sync_style = SyncStyle.EXPLORER_DRIVEN self.config.buffer.batch_size = 2 self.config.check_and_update() - scheduler = Scheduler(self.config) + scheduler = self._create_scheduler() await scheduler.start() tasks = [ @@ -999,7 +1014,7 @@ async def test_over_rollout_async_cancelled_runner_accepts_next_batch(self): self.config.synchronizer.sync_style = SyncStyle.EXPLORER_DRIVEN self.config.buffer.batch_size = 2 self.config.check_and_update() - scheduler = Scheduler(self.config) + scheduler = self._create_scheduler() await scheduler.start() tasks = [ @@ -1059,7 +1074,7 @@ async def test_over_rollout_sync_cancel_does_not_imply_immediate_runner_reuse(se self.config.synchronizer.sync_style = SyncStyle.EXPLORER_DRIVEN self.config.buffer.batch_size = 2 self.config.check_and_update() - scheduler = Scheduler(self.config) + scheduler = self._create_scheduler() await scheduler.start() tasks = [ @@ -1124,7 +1139,7 @@ async def test_timeout_cleanup_still_restarts_runner(self): self.config.explorer.max_repeat_times_per_runner = None self.config.synchronizer.sync_style = SyncStyle.EXPLORER_DRIVEN self.config.check_and_update() - scheduler = Scheduler(self.config) + scheduler = self._create_scheduler() await scheduler.start() tasks = generate_tasks(0, timeout_num=2, repeat_times=1, timeout_seconds=10) @@ -1140,7 +1155,7 @@ async def test_timeout_cleanup_still_restarts_runner(self): async def test_unexpected_task_exception_restarts_runner(self): self.config.explorer.runner_per_model = 1 self.config.check_and_update() - scheduler = Scheduler(self.config) + scheduler = self._create_scheduler() await scheduler.start() scheduler.runners[0].run_with_retry = AsyncMock(side_effect=RuntimeError("boom")) @@ -1171,7 +1186,7 @@ async def test_dynamic_timeout(self): self.config.buffer.batch_size = 4 self.config.explorer.max_timeout = 20 self.config.explorer.max_retry_times = 0 # no retry here - scheduler = Scheduler(self.config) + scheduler = self._create_scheduler() await scheduler.start() tasks = [] tasks.extend(generate_tasks(0, timeout_num=4, repeat_times=1, timeout_seconds=1)) @@ -1222,7 +1237,7 @@ async def test_dynamic_timeout_warmup_min_steps_uses_completed_steps(self): self.config.explorer.max_repeat_times_per_runner = 2 self.config.check_and_update() - scheduler = Scheduler(self.config) + scheduler = self._create_scheduler() await scheduler.start() tasks = generate_tasks(0, timeout_num=2, repeat_times=4, timeout_seconds=1) @@ -1248,7 +1263,7 @@ async def test_dynamic_timeout_warmup_min_steps_uses_completed_steps(self): await scheduler.stop() async def test_collect_results_reads_payloads_returned_by_workflow_runner(self): - scheduler = Scheduler(self.config) + scheduler = self._create_scheduler() await scheduler.start() scheduler.schedule(generate_tasks(3, repeat_times=2), batch_id=0) @@ -1260,7 +1275,7 @@ async def test_collect_results_reads_payloads_returned_by_workflow_runner(self): await scheduler.stop() async def test_timeout_cleanup_keeps_completed_payloads_local(self): - scheduler = Scheduler(self.config) + scheduler = self._create_scheduler() await scheduler.start() scheduler.schedule(generate_tasks(1, timeout_num=1, timeout_seconds=10), batch_id=0) @@ -1272,7 +1287,7 @@ async def test_timeout_cleanup_keeps_completed_payloads_local(self): await scheduler.stop() async def test_eval_tasks_do_not_return_training_experiences(self): - scheduler = Scheduler(self.config) + scheduler = self._create_scheduler() await scheduler.start() eval_tasks = generate_tasks(2, repeat_times=2) @@ -1288,7 +1303,7 @@ async def test_eval_tasks_do_not_return_training_experiences(self): await scheduler.stop() async def test_get_statuses_skips_payload_deserialization(self): - scheduler = Scheduler(self.config) + scheduler = self._create_scheduler() await scheduler.start() scheduler.schedule(generate_tasks(2, repeat_times=2), batch_id=0) @@ -1304,7 +1319,7 @@ async def test_get_statuses_skips_payload_deserialization(self): await scheduler.stop() async def test_get_payload_results_keeps_payloads_serialized(self): - scheduler = Scheduler(self.config) + scheduler = self._create_scheduler() await scheduler.start() scheduler.schedule(generate_tasks(2, repeat_times=2), batch_id=0) @@ -1344,6 +1359,9 @@ def setUp(self): self.config.check_and_update() self.model_actors = _create_named_model_actors(self.config) + def _create_scheduler(self) -> Scheduler: + return Scheduler(self.config, rollout_actors=_resolve_rollout_actors(self.config)) + def tearDown(self): try: _cleanup_named_model_actors(getattr(self, "model_actors", None)) @@ -1355,7 +1373,7 @@ def tearDown(self): pass async def test_runner_state_collection(self): - scheduler = Scheduler(self.config) + scheduler = self._create_scheduler() # 4 runner in side the scheduler await scheduler.start() diff --git a/trinity/common/config.py b/trinity/common/config.py index 98701d7fb25..dbf907584b3 100644 --- a/trinity/common/config.py +++ b/trinity/common/config.py @@ -573,7 +573,7 @@ class InferenceModelConfig: # recording identity (``record_key``). This is the single switch for the # recording flow — when on, the explorer also consumes from the store: the # WorkflowRunner ships only a small reward map keyed by ``record_key`` and - # the RolloutCoordinator pulls heavy experiences via ``/records/update_record`` + # the RolloutCoordinator pulls heavy experiences via ``/records/drain`` # at finalize time. When off (default), runners ship serialized experiences # through the scheduler as before (legacy path). When True, the Allocator # forces ``enable_openai_api`` (the recorder runs in the API server). diff --git a/trinity/common/models/model.py b/trinity/common/models/model.py index 78b9e668857..8d966cb1474 100644 --- a/trinity/common/models/model.py +++ b/trinity/common/models/model.py @@ -161,6 +161,49 @@ async def extract_experience_from_history( recorder.forget_record(record_key) return exps + async def update_experience_reward( + self, + record_key: str, + reward: float, + info: Optional[dict] = None, + sample_ids: Optional[List[str]] = None, + ) -> None: + """Update reward and optional info on recorded experiences.""" + recorder = getattr(self, "recorder", None) + if recorder is None: + raise ValueError("Recording is not enabled for this model.") + await recorder.flush() + if not recorder.store.get(record_key): + return + recorder.store.update( + key=record_key, + reward=reward, + info=info, + sample_ids=sample_ids, + ) + + async def drain_experience_records(self, prefix: str) -> List[Experience]: + """Remove and return recorded experiences matching a key or prefix.""" + recorder = getattr(self, "recorder", None) + if recorder is None: + raise ValueError("Recording is not enabled for this model.") + await recorder.flush() + matched_keys = [ + key for key in recorder.store.keys() if key == prefix or key.startswith(f"{prefix}/") + ] + exps = recorder.store.remove(prefix) + for key in matched_keys: + recorder.forget_record(key) + return exps + + async def drain_experience_records_bytes(self, prefix: str) -> bytes: + """Remove matching recorded experiences and return serialized bytes.""" + return Experience.serialize_many(await self.drain_experience_records(prefix)) + + async def delete_experience_records(self, prefix: str) -> int: + """Remove recorded experiences matching a key or prefix.""" + return len(await self.drain_experience_records(prefix)) + def get_model_config(self) -> InferenceModelConfig: """Get the model configuration.""" return self.config @@ -783,6 +826,49 @@ def extract_experience_from_history( self.recording_history_offsets[record_key] = len(exps) return new_exps + async def update_experience_reward_async( + self, + record_key: str, + reward: float, + info: Optional[dict] = None, + sample_ids: Optional[List[str]] = None, + ) -> None: + """Update reward and optional info on recorded experiences.""" + if not self.enable_history: + raise ValueError("History recording is not enabled.") + if self.model is None: + raise ValueError("Recording update requires an inference model actor.") + await self.model.update_experience_reward.remote( + record_key=record_key, + reward=reward, + info=info, + sample_ids=sample_ids, + ) + + async def drain_experience_records_async(self, prefix: str) -> List[Experience]: + """Remove and return recorded experiences matching a key or prefix.""" + if not self.enable_history: + raise ValueError("History recording is not enabled.") + if self.model is None: + raise ValueError("Recording drain requires an inference model actor.") + return await self.model.drain_experience_records.remote(prefix=prefix) + + async def drain_experience_records_bytes_async(self, prefix: str) -> bytes: + """Remove matching recorded experiences and return serialized bytes.""" + if not self.enable_history: + raise ValueError("History recording is not enabled.") + if self.model is None: + raise ValueError("Recording drain requires an inference model actor.") + return await self.model.drain_experience_records_bytes.remote(prefix=prefix) + + async def delete_experience_records_async(self, prefix: str) -> int: + """Remove recorded experiences matching a key or prefix.""" + if not self.enable_history: + raise ValueError("History recording is not enabled.") + if self.model is None: + raise ValueError("Recording delete requires an inference model actor.") + return await self.model.delete_experience_records.remote(prefix=prefix) + # Workflow state management methods async def set_workflow_state(self, state: Dict) -> None: """Set the state of workflow using the model.""" diff --git a/trinity/common/models/recording/query.py b/trinity/common/models/recording/query.py index af2c437e83d..719256f9c49 100644 --- a/trinity/common/models/recording/query.py +++ b/trinity/common/models/recording/query.py @@ -1,4 +1,4 @@ -"""HTTP query/update endpoints over recorded experiences.""" +"""HTTP endpoints over recorded generation experiences.""" from typing import List @@ -19,13 +19,15 @@ query_router = APIRouter(prefix="/records", tags=["trinity-recording"]) -class _RecordUpdate(BaseModel): +class _RewardUpdateRequest(BaseModel): record_key: str reward: float + info: dict | None = None + sample_ids: List[str] | None = None -class _UpdateRecordRequest(BaseModel): - updates: List[_RecordUpdate] +class _PrefixRequest(BaseModel): + prefix: str def _store(request: Request) -> RecordStore: @@ -107,25 +109,51 @@ async def delete_request_experience(record_key: str, request_id: str, request: R return {"record_key": record_key, "request_id": request_id, "deleted": True} -@query_router.post("/update_record") -async def update_record(req: _UpdateRecordRequest, request: Request) -> Response: +@query_router.post("/update_reward") +async def update_reward(req: _RewardUpdateRequest, request: Request) -> dict: store = _store(request) recorder = _recorder(request) await recorder.flush() + if not store.get(req.record_key): + return {"record_key": req.record_key, "updated": 0} + store.update( + key=req.record_key, + reward=req.reward, + info=req.info, + sample_ids=req.sample_ids, + ) + return { + "record_key": req.record_key, + "updated": ( + len(req.sample_ids) if req.sample_ids is not None else len(store.get(req.record_key)) + ), + } + - exps: List[Experience] = [] - for update in req.updates: - if not store.get(update.record_key): - continue - store.update( - key=update.record_key, - reward=update.reward, - info=None, - sample_ids=None, - ) - exps.extend(store.remove(update.record_key)) - recorder.forget_record(update.record_key) +@query_router.post("/drain") +async def drain_records(req: _PrefixRequest, request: Request) -> Response: + store = _store(request) + recorder = _recorder(request) + await recorder.flush() + matched_keys = [ + key for key in store.keys() if key == req.prefix or key.startswith(f"{req.prefix}/") + ] + exps = store.remove(req.prefix) + for key in matched_keys: + recorder.forget_record(key) return Response( content=Experience.serialize_many(exps), media_type="application/octet-stream", ) + + +@query_router.delete("") +async def delete_records(req: _PrefixRequest, request: Request) -> dict: + store = _store(request) + matched_keys = [ + key for key in store.keys() if key == req.prefix or key.startswith(f"{req.prefix}/") + ] + deleted = len(store.remove(req.prefix)) + for key in matched_keys: + _forget_record(request, key) + return {"prefix": req.prefix, "deleted": deleted} diff --git a/trinity/common/models/vllm_patch/recording/server.py b/trinity/common/models/vllm_patch/recording/server.py index 890d2be9698..b8283b9cc40 100644 --- a/trinity/common/models/vllm_patch/recording/server.py +++ b/trinity/common/models/vllm_patch/recording/server.py @@ -109,7 +109,7 @@ def _setup_recording( by VLLMModel at engine creation and in ``sync_model_weights``). The store backend is always the in-process ``MemoryStore``; the coordinator - drains it at finalize time via ``/records/update_record`` (fanned out per + drains it at finalize time via ``/records/drain`` (fanned out per rank), so heavy experience bytes never touch SQL or Ray serialization. Args: diff --git a/trinity/common/workflows/__init__.py b/trinity/common/workflows/__init__.py index 7627cdca296..e16e931e996 100644 --- a/trinity/common/workflows/__init__.py +++ b/trinity/common/workflows/__init__.py @@ -1,6 +1,6 @@ # -*- coding: utf-8 -*- """Workflow module""" -from trinity.common.workflows.workflow import Task, Workflow +from trinity.common.workflows.workflow import Status, Task, Workflow from trinity.utils.registry import Registry WORKFLOWS: Registry = Registry( @@ -57,6 +57,7 @@ __all__ = [ "Task", + "Status", "Workflow", "WORKFLOWS", ] diff --git a/trinity/common/workflows/workflow.py b/trinity/common/workflows/workflow.py index eb795114375..3947859acf3 100644 --- a/trinity/common/workflows/workflow.py +++ b/trinity/common/workflows/workflow.py @@ -4,7 +4,7 @@ from __future__ import annotations from dataclasses import asdict, dataclass, field -from typing import TYPE_CHECKING, List, Optional, Type, Union +from typing import TYPE_CHECKING, Dict, List, Optional, Type, Union from trinity.buffer.store import parse_record_key from trinity.common.config import FormatConfig, GenerationConfig @@ -18,6 +18,21 @@ from trinity.common.models.model import ModelWrapper +@dataclass(frozen=True) +class Status: + """Status of workflow, task, and batch execution.""" + + completed_runs: int + total_runs: int + metrics: List[Dict[str, float]] + successful_run_ids: List[int] = field(default_factory=list) + message: Optional[str] = None + + @property + def ok(self) -> bool: + return self.completed_runs == self.total_runs + + @dataclass class Task(dict): """A Task class that defines a task and its associated reward function / workflow.""" @@ -79,7 +94,10 @@ def to_dict(self) -> dict: class Workflow: """The base workflow class. - A workflow is a runnable object which generates a list of experiences. + A workflow is a runnable object that executes rollout logic and returns a + :class:`Status`. Training experiences are captured by the rollout model's + built-in recording path during generation; workflows should update rewards + on those recorded experiences before returning. Attributes: auxiliary_model_wrappers: List of ModelWrapper instances for auxiliary models. @@ -90,12 +108,12 @@ class Workflow: can_repeat: bool = False # whether the workflow can be repeated multiple times. If true, `set_repeat_times()` must be implemented. is_async: bool = False # whether the workflow runs in async mode. If true, `run_async()` must be implemented, else `run()` must be implemented. - # When True, the workflow cooperates with the in-vLLM recorder: each chat - # call carries a per-sample ``record_key`` (the MemoryStore group key) so - # the runner can later report ``{record_key: reward}`` and the coordinator - # can join reward inside the store. Set by the WorkflowRunner from - # ``config.explorer.rollout_model.enable_history``. See ``SimpleWorkflow`` - # for the per-sample (n=1) loop this triggers. + # Mirrors ``config.explorer.rollout_model.enable_history``. When it is set, + # the rollout model records generated experiences in its model-side store + # under ``record_key`` groups. Workflows use that rollout-model history to + # update rewards before returning Status; the coordinator later drains a + # whole step by prefix and sends the recorded experiences to the + # ExperiencePipeline. enable_history: bool = False def __init__( @@ -146,20 +164,41 @@ def set_repeat_times(self, repeat_times: int, run_id_base: int) -> None: Set the number of times to repeat the workflow. Args: repeat_times (int): number of times to repeat the workflow (if repeatable). - run_id_base (int): base run_id for setting run_id in experiences. + run_id_base (int): base run_id for recording keys and returned statuses. """ raise NotImplementedError( "set_repeat_times() must be implemented for a repeatable workflow." ) - def run(self) -> List[Experience]: - """Run workflow and return a list of experiences.""" + def run(self) -> Status: + """Run workflow and return its execution status.""" raise NotImplementedError - async def run_async(self) -> List[Experience]: - """Run workflow in async and return a list of experiences.""" + async def run_async(self) -> Status: + """Run workflow in async and return its execution status.""" raise NotImplementedError + def _build_record_key(self, run_id: Optional[int] = None) -> str: + run = self.run_id_base if run_id is None else run_id + return f"{self.task.batch_id}/{self.task.task_id}/{run}" + + async def update_reward( + self, + reward: float, + info: Optional[dict] = None, + sample_ids: Optional[List[str]] = None, + run_id: Optional[int] = None, + ) -> None: + """Update recorded experiences for one run with reward and optional info.""" + if not self.enable_history: + return + await self.model.update_experience_reward_async( + record_key=self._build_record_key(run_id), + reward=reward, + info=info, + sample_ids=sample_ids, + ) + class MultiTurnWorkflow(Workflow): """ diff --git a/trinity/explorer/proxy/service.py b/trinity/explorer/proxy/service.py index 89eef711ed7..5261d6ceb6a 100644 --- a/trinity/explorer/proxy/service.py +++ b/trinity/explorer/proxy/service.py @@ -15,7 +15,7 @@ class ExplorerService: The proxy is a request router + model-weight sync coordinator for serve mode. Experience collection used to live here (SQL-mediated ``/feedback``/``/commit``); it has been removed in favor of the in-vLLM - ``MemoryStore`` + ``/records/update_record`` flow driven by the rollout + ``MemoryStore`` + ``/records/drain`` flow driven by the rollout coordinator. Serve-mode external reward reporting is therefore pending (see the recording refactor plan). """ diff --git a/trinity/explorer/rollout_coordinator.py b/trinity/explorer/rollout_coordinator.py index 6667b3facba..ad439b215d8 100644 --- a/trinity/explorer/rollout_coordinator.py +++ b/trinity/explorer/rollout_coordinator.py @@ -1,19 +1,15 @@ """Rollout coordinator for async batch submission and finalize.""" import asyncio -import pickle -from collections import defaultdict from dataclasses import dataclass, field from enum import Enum -from typing import Any, Dict, List, Literal, Optional, Union +from typing import Any, Dict, Literal, Optional, Union -import httpx import ray from ray.actor import ActorHandle from trinity.buffer.pipelines.experience_pipeline import ExperiencePipeline from trinity.common.config import Config -from trinity.common.experience import Experience from trinity.common.workflows import Task from trinity.explorer.scheduler import Scheduler from trinity.utils.log import get_logger @@ -22,10 +18,6 @@ BatchId = Union[int, str] BatchType = Literal["train", "eval"] -#: Default per-rank consume HTTP timeout (seconds). The consume returns heavy -#: experience bytes, so allow generous headroom over the inference timeout. -_CONSUME_TIMEOUT = 300.0 - class BatchLifecycleState(str, Enum): """Lifecycle states for one submitted batch.""" @@ -72,42 +64,39 @@ def __init__( self.pending_batches: Dict[BatchId, BatchState] = {} self.running = False self.detailed_stats = getattr(getattr(config, "monitor", None), "detailed_stats", False) - # Lazily-resolved map of rollout engine_id -> API server URL, for the - # recording path's per-rank /records/update_record fan-out. - self._rank_urls: Optional[Dict[int, str]] = None + # Lazily-resolved map of rollout engine_id -> rollout actor handle, for + # recording residual cleanup. + self._rollout_actors: Dict[int, ActorHandle] = {} def _enable_history_recording(self) -> bool: """Whether the recording-consume path is active for train batches.""" return bool(self.config.explorer.rollout_model.enable_history) - def _resolve_rank_urls(self) -> Dict[int, str]: - """Resolve each rollout engine's API server URL via named Ray actors. + def _resolve_rollout_actors(self) -> Dict[int, ActorHandle]: + """Resolve each rollout engine's actor handle via named Ray actors. Mirrors ``Allocator.get_actor_name`` + ``ray.get_actor``: rollout model actors are named ``f"{explorer.name}_rollout_model_{engine_id}_0"`` - (node_id 0 holds the API server). Cached after first resolution. + (node_id 0 holds the recording store). Cached after first resolution. """ - if self._rank_urls is not None: - return self._rank_urls + if self._rollout_actors: + return self._rollout_actors rollout_cfg = self.config.explorer.rollout_model name = self.config.explorer.name namespace = rollout_cfg.ray_namespace - urls: Dict[int, str] = {} + actors: Dict[int, ActorHandle] = {} for engine_id in range(rollout_cfg.engine_num): actor_name = f"{name}_rollout_model_{engine_id}_0" try: - actor = ray.get_actor(actor_name, namespace=namespace) - except ValueError: - self.logger.warning( - "rollout actor %s not found in namespace %s; skipping rank %d", - actor_name, - namespace, - engine_id, - ) - continue - urls[engine_id] = ray.get(actor.get_api_server_url.remote()) - self._rank_urls = urls - return urls + actors[engine_id] = ray.get_actor(actor_name, namespace=namespace) + except ValueError as exc: + raise RuntimeError( + "Rollout actor %s not found in namespace %s." + " RolloutCoordinator cannot initialize Scheduler without all rollout actors." + % (actor_name, namespace) + ) from exc + self._rollout_actors = actors + return actors async def prepare(self) -> None: """Initialize the owned pipeline and scheduler.""" @@ -142,6 +131,7 @@ async def _init_scheduler(self): return self.scheduler = Scheduler( self.config, + rollout_actors=self._resolve_rollout_actors(), ) await self.scheduler.start() @@ -218,6 +208,7 @@ async def _finalize_eval_batch( if task_id in batch_state.statuses: continue batch_state.statuses[task_id] = status + await self._discard_recorded_experiences(str(batch_state.batch_id)) return self._finish_batch(batch_state, pipeline_metrics={}) async def abort_batch( @@ -242,6 +233,7 @@ async def abort_batch( restart_runners=True, ) scheduler.discard_completed_results(batch_id) + await self._discard_recorded_experiences(str(batch_id)) batch_state.state = BatchLifecycleState.ABORTED batch_state.final_result = self._build_batch_result(batch_state, pipeline_metrics={}) @@ -319,10 +311,8 @@ async def _finalize_train_batch( batch_state.state = BatchLifecycleState.FINALIZING try: - if self._enable_history_recording(): - pipeline_metrics = await self._consume_recorded_experiences(payload_chunks) - else: - pipeline_metrics = await self.process_experiences(payload_chunks) + pipeline_metrics = await self.process_experiences(payload_chunks) + await self._discard_recorded_experiences(str(batch_state.batch_id)) if not is_complete: await self._cleanup_train_batch_runtime(batch_state) except Exception: @@ -331,65 +321,18 @@ async def _finalize_train_batch( return self._finish_batch(batch_state, pipeline_metrics=pipeline_metrics) - async def _consume_recorded_experiences(self, payload_chunks: List[bytes]) -> dict: - """Recording path: pull heavy experiences from each vLLM rank's store. - - ``payload_chunks`` are small pickle reward maps produced by the runners - (``{"engine_id": int, "updates": [{"record_key", "reward", "run", "task"}]}``). - Group updates by engine, fan out ``POST /records/update_record`` to each - rank (which drains its recorder, reward-stamps the matching record-key - groups, pops them, and returns ``serialize_many`` bytes), deserialize, - and feed the assembled experiences straight into the pipeline — no Ray - serialization of heavy tensors, and reward is fused inside the store. - """ - if self.experience_pipeline is None: - raise RuntimeError("Experience pipeline is not initialized.") - per_engine: Dict[int, List[dict]] = defaultdict(list) - for chunk in payload_chunks: - if not chunk: - continue - data = pickle.loads(chunk) - per_engine[int(data["engine_id"])].extend(data["updates"]) - - if not per_engine: - return {} - - rank_urls = self._resolve_rank_urls() - async with httpx.AsyncClient(timeout=_CONSUME_TIMEOUT) as client: - requests = [ - self._post_update_record(client, rank_urls[engine_id], updates) - for engine_id, updates in per_engine.items() - if engine_id in rank_urls - ] - responses = await asyncio.gather(*requests) - - exps: List[Experience] = [] - for resp_bytes in responses: - if resp_bytes: - exps.extend(Experience.deserialize_many(resp_bytes)) - return await self.experience_pipeline.process_experiences(exps) - - async def _post_update_record( - self, client: httpx.AsyncClient, rank_url: str, updates: List[dict] - ) -> bytes: - """POST a batch of record-key reward updates to one rank; return heavy bytes.""" - try: - resp = await client.post( - f"{rank_url}/records/update_record", - json={"updates": updates}, - ) - except (httpx.TimeoutException, httpx.RequestError) as exc: - self.logger.error("update_record to %s failed: %s", rank_url, exc) - return b"" - if resp.status_code != 200: - self.logger.error( - "update_record to %s returned %d: %s", - rank_url, - resp.status_code, - resp.text[:200], - ) - return b"" - return resp.content + async def _discard_recorded_experiences(self, prefix: str) -> None: + """Delete recorded experiences matching a prefix from all rollout ranks.""" + if not self._enable_history_recording(): + return + actors = self._resolve_rollout_actors() + results = await asyncio.gather( + *[actor.delete_experience_records.remote(prefix=prefix) for actor in actors.values()], + return_exceptions=True, + ) + for result in results: + if isinstance(result, Exception): + self.logger.error("records delete from rollout actor failed: %s", result) def _finish_batch( self, @@ -417,6 +360,7 @@ async def _cleanup_train_batch_runtime(self, batch_state: BatchState) -> None: return_partial_tasks=False, restart_runners=True, ) + await self._discard_recorded_experiences(str(batch_state.batch_id)) def _build_batch_result( self, diff --git a/trinity/explorer/scheduler.py b/trinity/explorer/scheduler.py index 57289e6e28d..23a74f47d81 100644 --- a/trinity/explorer/scheduler.py +++ b/trinity/explorer/scheduler.py @@ -9,6 +9,7 @@ from typing import Dict, List, Optional, Tuple, Union import ray +from ray.actor import ActorHandle from trinity.common.config import Config from trinity.common.workflows import Task @@ -31,6 +32,7 @@ class TaskWrapper: completed_runs: int = 0 total_runs: int = 0 # total planned runs for the whole task metrics: List[Dict[str, float]] = field(default_factory=list) + successful_run_ids: List[int] = field(default_factory=list) experience_payloads: List[bytes] = field(default_factory=list) first_error: Optional[str] = None emitted: bool = False @@ -64,6 +66,7 @@ def __init__( rollout_model_id: int, auxiliary_model_ids: List[int], config: Config, + rollout_actor: ActorHandle, ): self.logger = get_logger(__name__) self.runner_id = runner_id @@ -75,6 +78,7 @@ def __init__( self.namespace = config.ray_namespace self.runner = self._create_runner() self.state = {} + self.rollout_actor = rollout_actor def _create_runner(self): return ( @@ -103,6 +107,48 @@ async def update_state(self) -> None: async def prepare(self): await self.runner.prepare.remote() + def _enable_history_recording(self) -> bool: + return bool(self.config.explorer.rollout_model.enable_history) + + def _record_key(self, task: TaskWrapper, run_id: int) -> str: + return f"{task.batch_id}/{task.task.task_id}/{run_id}" + + async def _drain_records(self, prefix: str) -> bytes: + try: + return await self.rollout_actor.drain_experience_records_bytes.remote(prefix=prefix) + except Exception as exc: + self.logger.error("records drain from rollout actor failed: %s", exc) + return b"" + + async def _delete_records(self, prefix: str) -> None: + try: + await self.rollout_actor.delete_experience_records.remote(prefix=prefix) + except Exception as exc: + self.logger.error("records delete from rollout actor failed: %s", exc) + + async def _consume_finished_records(self, task: TaskWrapper, status: Status) -> List[bytes]: + if not self._enable_history_recording(): + return [] + if not status.successful_run_ids: + return [] + + if task.task.is_eval: + await asyncio.gather( + *[ + self._delete_records(self._record_key(task, run_id)) + for run_id in status.successful_run_ids + ] + ) + return [] + + payloads = await asyncio.gather( + *[ + self._drain_records(self._record_key(task, run_id)) + for run_id in status.successful_run_ids + ] + ) + return [payload for payload in payloads if payload] + async def run_with_retry( self, task: TaskWrapper, @@ -110,7 +156,7 @@ async def run_with_retry( run_id_base: int, timeout: float, collect_partial_runs: bool, - ) -> Tuple[Status, bytes, int, float]: + ) -> Tuple[Status, List[bytes], int, float]: """ Args: task (`TaskWrapper`): The task to run. @@ -120,7 +166,7 @@ async def run_with_retry( Returns: `Status`: The return status of the task. - `List`: The experiences generated by the task. + `List[bytes]`: Serialized recorded experiences drained by Scheduler. `int`: The runner_id of current runner. `float`: The time taken to run the task. """ @@ -128,7 +174,6 @@ async def run_with_retry( await self.runner.__ray_ready__.remote() start_time = time.time() status = Status(completed_runs=0, total_runs=repeat_times, metrics=list()) - exp_payload = b"" run_task_ref = None task2run = replace( task.task, @@ -147,7 +192,7 @@ async def run_with_retry( run_id_base=run_id_base, collect_partial_runs=collect_partial_runs, ) - status, exp_payload = await asyncio.wait_for( + status = await asyncio.wait_for( run_task_ref, timeout=timeout, ) @@ -192,7 +237,8 @@ async def run_with_retry( finally: end_time = time.time() status.metrics.append({"time/task_execution": end_time - start_time}) - return status, exp_payload, self.runner_id, end_time - start_time + experience_payloads = await self._consume_finished_records(task, status) + return status, experience_payloads, self.runner_id, end_time - start_time async def restart_runner(self): old_runner = self.runner @@ -228,9 +274,18 @@ class Scheduler: def __init__( self, config: Config, + rollout_actors: Dict[int, ActorHandle], ): self.logger = get_logger(__name__) self.config = config + expected_rollout_actor_ids = set(range(config.explorer.rollout_model.engine_num)) + missing_rollout_actor_ids = expected_rollout_actor_ids.difference(rollout_actors) + if missing_rollout_actor_ids: + raise ValueError( + "Scheduler requires rollout actors for all rollout engines; " + f"missing engine ids: {sorted(missing_rollout_actor_ids)}" + ) + self.rollout_actors = rollout_actors self.namespace = ray.get_runtime_context().namespace self.default_timeout = config.explorer.max_timeout * (config.explorer.max_retry_times + 1) self.max_retry_times = config.explorer.max_retry_times @@ -283,6 +338,9 @@ async def _create_runner( for j in range(len(self.config.explorer.auxiliary_models)) ], config=self.config, + rollout_actor=self.rollout_actors[ + runner_id % self.config.explorer.rollout_model.engine_num + ], ) await runner.prepare() self.runners[runner_id] = runner @@ -377,11 +435,11 @@ def task_done_callback(self, async_task: asyncio.Task): self.logger.error(f"Task {task.task.task_id} failed: {async_task.exception()}") self._schedule_runner_restart(runner_id) else: - status, exp_payload, runner_id, run_time = async_task.result() + status, experience_payloads, runner_id, run_time = async_task.result() if not task.task.is_eval: self.total_running_time += run_time self.total_completed_sub_tasks += 1 - self._accumulate_task_result(task, status, exp_payload) + self._accumulate_task_result(task, status, experience_payloads) self.busy_runners.pop(runner_id, None) self.idle_runners.add(runner_id) # If all sub runs in a task are completed @@ -397,17 +455,17 @@ def task_done_callback(self, async_task: asyncio.Task): del self.running_tasks[task.batch_id] def _accumulate_task_result( - self, task: TaskWrapper, status: Status, experience_payload: bytes + self, task: TaskWrapper, status: Status, experience_payloads: List[bytes] ) -> None: task.finished_sub_task_num += 1 task.completed_runs += status.completed_runs task.metrics.extend(status.metrics) - if experience_payload: - task.experience_payloads.append(experience_payload) + task.successful_run_ids.extend(status.successful_run_ids) + task.experience_payloads.extend(experience_payloads) if not status.ok and task.first_error is None: task.first_error = status.message - def _build_task_result(self, task: TaskWrapper) -> Tuple[Status, List[bytes]]: + def _build_task_result(self, task: TaskWrapper) -> Status: if task.completed_runs < task.total_runs: message = f"{task.completed_runs}/{task.total_runs} runs completed successfully." if task.first_error: @@ -420,20 +478,21 @@ def _build_task_result(self, task: TaskWrapper) -> Tuple[Status, List[bytes]]: completed_runs=task.completed_runs, total_runs=task.total_runs, metrics=[calculate_task_level_metrics(task.metrics, task.task.is_eval)], + successful_run_ids=sorted(task.successful_run_ids), message=message, ) - return status, list(task.experience_payloads) + return status def _emit_task_result(self, task: TaskWrapper) -> None: if task.emitted: return - status, experience_payloads = self._build_task_result(task) + status = self._build_task_result(task) task_id = task.task.task_id completed_result = CompletedTaskResult( batch_id=task.batch_id, task_id=task_id, status=status, - experience_payloads=experience_payloads, + experience_payloads=list(task.experience_payloads), ) self.completed_tasks[task.batch_id][task_id] = completed_result task.emitted = True @@ -677,15 +736,14 @@ def _collect_batch_results(self, batch_id: Union[int, str]) -> Tuple[List[Status completed_results = list(self.completed_tasks.get(batch_id, {}).values()) for result in completed_results: statuses.append(result.status) - if result.experience_payloads: - payload_chunks.extend(result.experience_payloads) + payload_chunks.extend(result.experience_payloads) return statuses, payload_chunks async def drain_batch_payload_results( self, batch_id: Union[int, str] ) -> Tuple[List[Status], List[bytes]]: - """Drain cached completed results for one batch.""" + """Drain cached completed statuses and payload chunks for one batch.""" statuses, payload_chunks = self._collect_batch_results(batch_id) @@ -706,7 +764,7 @@ async def _get_batch_payload_results( clear_timeout_tasks: bool, return_partial_tasks: bool, ) -> Tuple[List[Status], List[bytes]]: - """Wait for one batch and drain its completed payload chunks.""" + """Wait for one batch and drain completed statuses plus payload chunks.""" timeout = timeout or self.default_timeout scheduled_num, min_num = self._resolve_result_target(batch_id, min_num) @@ -747,15 +805,16 @@ async def get_payload_results( clear_timeout_tasks: bool = True, return_partial_tasks: bool = False, ) -> Tuple[List[Status], List[bytes]]: - """Wait for one batch and return task statuses plus serialized payload chunks.""" + """Wait for one batch and return statuses plus Scheduler-drained payload chunks.""" - return await self._get_batch_payload_results( + statuses, payload_chunks = await self._get_batch_payload_results( batch_id=batch_id, min_num=min_num, timeout=timeout, clear_timeout_tasks=clear_timeout_tasks, return_partial_tasks=return_partial_tasks, ) + return statuses, payload_chunks async def get_statuses( self, diff --git a/trinity/explorer/workflow_runner.py b/trinity/explorer/workflow_runner.py index b00d59cad73..db68046b77c 100644 --- a/trinity/explorer/workflow_runner.py +++ b/trinity/explorer/workflow_runner.py @@ -3,49 +3,26 @@ import asyncio import os -import pickle import time import traceback -from dataclasses import dataclass from typing import Dict, List, Optional, Tuple from trinity.buffer import get_buffer_reader, get_buffer_writer -from trinity.buffer.store import get_record_key from trinity.common.config import Config, StorageConfig from trinity.common.constants import LOG_DIR_ENV_VAR, LOG_LEVEL_ENV_VAR -from trinity.common.experience import Experience from trinity.common.models.allocator import Allocator from trinity.common.models.model import ModelWrapper -from trinity.common.workflows import Task, Workflow +from trinity.common.workflows import Status, Task, Workflow from trinity.utils.log import get_logger -from trinity.utils.metrics import aggregate_run_level_metrics - - -@dataclass(frozen=True) -class Status: - """Status of the task running result.""" - - completed_runs: int - total_runs: int - metrics: List[Dict[str, float]] - # A list of metric dictionaries, where each dictionary is from a single run. - message: Optional[str] = None - - @property - def ok(self) -> bool: - return self.completed_runs == self.total_runs - - -@dataclass(frozen=True) -class RunnerExecutionResult: - """Execution result for one runner task.""" - - status: Status - experiences: List[Experience] class WorkflowRunner: - """A Ray remote actor to run the workflow and generate experiences.""" + """A Ray remote actor that runs workflows and returns execution statuses. + + Experience payloads are not returned through the runner. The rollout model + owns experience capture through its recording/history path, and the rollout + coordinator drains those model-side stores at step finalization. + """ def __init__( self, @@ -134,12 +111,17 @@ def _create_workflow_instance(self, task: Task, record_key: Optional[str] = None def _enable_history_recording(self) -> bool: return bool(self.config.explorer.rollout_model.enable_history) - async def _run_workflow(self, workflow_instance: Workflow) -> List[Experience]: + async def _run_workflow(self, workflow_instance: Workflow) -> Status: if workflow_instance.asynchronous: - exps = await workflow_instance.run_async() + status = await workflow_instance.run_async() else: - exps = workflow_instance.run() - return exps + status = workflow_instance.run() + if not isinstance(status, Status): + raise TypeError( + f"{workflow_instance.__class__.__name__}.run must return Status, " + f"got {type(status).__name__}." + ) + return status def _create_isolated_workflow_instance( self, task: Task, record_key: Optional[str] = None @@ -157,14 +139,14 @@ def _create_isolated_workflow_instance( wf.enable_history = self._enable_history_recording() return wf, model_wrapper - def _build_execution_result( + def _build_status( self, total_runs: int, - completed_runs: int, metrics: List[Dict[str, float]], - experiences: List[Experience], + successful_run_ids: List[int], first_error: Optional[str] = None, - ) -> RunnerExecutionResult: + ) -> Status: + completed_runs = len(successful_run_ids) if first_error is None: message = None elif completed_runs > 0: @@ -175,39 +157,35 @@ def _build_execution_result( else: message = first_error - return RunnerExecutionResult( - status=Status( - completed_runs=completed_runs, - total_runs=total_runs, - metrics=list(metrics), - message=message, - ), - experiences=experiences, + return Status( + completed_runs=completed_runs, + total_runs=total_runs, + metrics=list(metrics), + successful_run_ids=list(successful_run_ids), + message=message, ) def _aggregate_run_results( self, total_runs: int, - results: List[Tuple[bool, List[Experience], Optional[Dict[str, float]], Optional[str]]], - ) -> RunnerExecutionResult: - exps = [] + results: List[Status], + ) -> Status: run_metrics = [] + successful_run_ids = [] first_error = None - for ok, new_exps, run_metric, error in results: - if ok: - exps.extend(new_exps) - if run_metric is not None: - run_metrics.append(run_metric) + for status in results: + if status.ok: + run_metrics.extend(status.metrics) + successful_run_ids.extend(status.successful_run_ids) continue if first_error is None: - first_error = error + first_error = status.message - return self._build_execution_result( + return self._build_status( total_runs=total_runs, - completed_runs=len(run_metrics), metrics=run_metrics, - experiences=exps, + successful_run_ids=successful_run_ids, first_error=first_error, ) @@ -218,10 +196,8 @@ async def _run_parallel_runs( run_id_base: int, collect_partial_runs: bool = True, use_threads: bool = False, - ) -> RunnerExecutionResult: - async def run_single( - i: int, - ) -> Tuple[bool, List[Experience], Optional[Dict[str, float]], Optional[str]]: + ) -> Status: + async def run_single(i: int) -> Status: run_index = run_id_base + i record_key = self._build_record_key(task, run_index) workflow, model_wrapper = self._create_isolated_workflow_instance(task, record_key) @@ -267,8 +243,7 @@ async def run_single( future_to_run_index.pop(future) result = future.result() results.append(result) - ok, _, _, _ = result - if not ok: + if not result.ok: should_stop = True if should_stop: for future in pending: @@ -287,32 +262,47 @@ async def _execute_single_run( run_id_base: int, model_wrapper: Optional[ModelWrapper] = None, record_key: Optional[str] = None, - ) -> Tuple[bool, List[Experience], Optional[Dict[str, float]], Optional[str]]: + ) -> Status: st = time.time() model_wrapper = model_wrapper or self.model_wrapper self._set_record_key(model_wrapper, record_key) await model_wrapper.clean_workflow_state() - self.runner_state["workflow_id"] = self._build_record_key(task, run_id_base + run_index) + run_id = run_id_base + run_index + self.runner_state["workflow_id"] = self._build_record_key(task, run_id) self.runner_state["terminate_time"] = None self.runner_state["begin_time"] = st try: - new_exps = await self._run_workflow(workflow) + status = await self._run_workflow(workflow) et = time.time() self.runner_state["terminate_time"] = et - run_metric = aggregate_run_level_metrics( - [exp.metrics for exp in new_exps if exp.metrics] + metrics = [dict(metric) for metric in status.metrics] + if not metrics: + metrics = [{}] + for metric in metrics: + metric["time/run_execution"] = et - st + successful_run_ids = status.successful_run_ids or ( + [run_id] if status.completed_runs > 0 else [] ) - run_metric["time/run_execution"] = et - st - for exp in new_exps: - exp.eid.run = run_id_base + run_index - return True, new_exps, run_metric, None + status = Status( + completed_runs=len(successful_run_ids), + total_runs=status.total_runs, + metrics=metrics, + successful_run_ids=successful_run_ids, + message=status.message, + ) + return status except Exception as exc: self.runner_state["terminate_time"] = time.time() error_trace_back = traceback.format_exc() self.logger.error( "WorkflowRunner single run error: " f"{exc}\nTraceback:\n{error_trace_back}" ) - return False, [], None, error_trace_back.rstrip() + return Status( + completed_runs=0, + total_runs=1, + metrics=[], + message=error_trace_back.rstrip(), + ) async def _run_task( self, @@ -320,7 +310,7 @@ async def _run_task( repeat_times: int, run_id_base: int, collect_partial_runs: bool = True, - ) -> RunnerExecutionResult: + ) -> Status: """Init workflow from the task and run it.""" if task.workflow.can_repeat: record_key = self._build_record_key(task, run_id_base) @@ -332,18 +322,20 @@ async def _run_task( self.runner_state["workflow_id"] = record_key self.runner_state["terminate_time"] = None self.runner_state["begin_time"] = st - exps = await self._run_workflow(workflow_instance) + status = await self._run_workflow(workflow_instance) et = time.time() self.runner_state["terminate_time"] = et - # repeatable workflow cannot calculate run level metrics, we use experience level metrics directly - run_metrics = [exp.metrics for exp in exps if exp.metrics] + run_metrics = [dict(metric) for metric in status.metrics] for metric in run_metrics: metric["time/run_execution"] = et - st - return self._build_execution_result( + successful_run_ids = status.successful_run_ids or list( + range(run_id_base, run_id_base + status.completed_runs) + ) + return self._build_status( total_runs=repeat_times, - completed_runs=repeat_times, metrics=run_metrics, - experiences=exps, + successful_run_ids=successful_run_ids, + first_error=status.message, ) else: return await self.concurrent_run_fn( @@ -359,7 +351,7 @@ async def _sequential_run( repeat_times: int, run_id_base: int, collect_partial_runs: bool = True, - ) -> RunnerExecutionResult: + ) -> Status: results = [] for i in range(repeat_times): run_index = run_id_base + i @@ -375,8 +367,7 @@ async def _sequential_run( results.append(result) if collect_partial_runs: continue - ok, _, _, _ = result - if ok: + if result.ok: continue break return self._aggregate_run_results(repeat_times, results) @@ -387,7 +378,7 @@ async def _asynchronous_run( repeat_times: int, run_id_base: int, collect_partial_runs: bool = True, - ) -> RunnerExecutionResult: + ) -> Status: return await self._run_parallel_runs( task, repeat_times, @@ -401,7 +392,7 @@ async def _multi_threading_run( repeat_times: int, run_id_base: int, collect_partial_runs: bool = True, - ) -> RunnerExecutionResult: + ) -> Status: return await self._run_parallel_runs( task, repeat_times, @@ -423,8 +414,8 @@ async def run_task( repeat_times: int = 1, run_id_base: int = 0, collect_partial_runs: bool = True, - ) -> Tuple[Status, bytes]: - """Run the task and return the states.""" + ) -> Status: + """Run the task and return its execution status.""" st = time.time() try: model_version = await self.model_wrapper.model_version_async @@ -432,74 +423,22 @@ async def run_task( self.logger.info( f"Starting task: step={batch_id}, model_version={model_version}, repeat_times={repeat_times}, run_id_base={run_id_base}" ) - execution_result = await self._run_task( + status = await self._run_task( task, repeat_times, run_id_base, collect_partial_runs=collect_partial_runs, ) - model_version_after = await self.model_wrapper.model_version_async - exps = execution_result.experiences - if execution_result.status.completed_runs > 0: - assert exps is not None and len(exps) > 0, "An empty experience is generated" - # set eid for each experience - for exp in exps: - exp.eid.batch = task.batch_id - # keep exp.eid.task if it has been set before (e.g., in workflow) - if exp.eid.task == "": # "" is the default value - exp.eid.task = task.task_id - if not hasattr(exp, "info") or exp.info is None: - exp.info = {} - exp.info["model_version"] = model_version - exp.info["model_version_drift"] = model_version_after - model_version - exp.info["use_count"] = 0 - exp.info["task_index"] = task.index - - if not hasattr(exp, "metrics") or exp.metrics is None: - exp.metrics = {} - - status = execution_result.status - - if task.is_eval: - # Eval tasks are not written to the training buffer. Under the - # recording path their turns are still in the vLLM MemoryStore; - # the coordinator's eval finalize must drain-and-discard them - # (TODO: wire eval consume-and-discard so eval turns don't leak - # in the store). For now, return no payload. - return status, b"" - elif self._enable_history_recording(): - # Recording path: ship only the small reward map keyed by the - # per-sample record_key the workflow stamped on each exp. The - # heavy experiences live in the vLLM MemoryStore and are pulled - # by the coordinator at finalize time. - updates = [ - { - "record_key": get_record_key(exp), - "reward": exp.reward, - "run": exp.eid.run, - "task": str(task.task_id), - } - for exp in exps - ] - reward_payload = pickle.dumps( - {"engine_id": self.rollout_model_id, "updates": updates} - ) - return status, reward_payload - else: - exp_payload = Experience.serialize_many(exps) - return status, exp_payload + return status except Exception as e: error_trace_back = traceback.format_exc() self.logger.error(f"WorkflowRunner run task error: {e}\nTraceback:\n{error_trace_back}") - return ( - Status( - completed_runs=0, - total_runs=repeat_times, - metrics=[{"time/run_execution": time.time() - st}], - message=error_trace_back.rstrip(), - ), - b"", + return Status( + completed_runs=0, + total_runs=repeat_times, + metrics=[{"time/run_execution": time.time() - st}], + message=error_trace_back.rstrip(), ) @@ -556,22 +495,28 @@ async def debug(self) -> None: task = tasks[0] self.logger.info(f"Start debugging task:\n{task.raw_task}") if not self.enable_profiling: - status, exp_payload = await self.run_task( - task=task, batch_id="debug", repeat_times=1, run_id_base=0 - ) + status = await self.run_task(task=task, batch_id="debug", repeat_times=1, run_id_base=0) else: from viztracer import VizTracer with VizTracer(output_file=self.output_profiling_file): - status, exp_payload = await self.run_task( + status = await self.run_task( task=task, batch_id="debug", repeat_times=1, run_id_base=0 ) - experiences = Experience.deserialize_many(exp_payload) if exp_payload else [] + experiences = [] + if self.config.explorer.rollout_model.enable_history: + try: + experiences = await self.model_wrapper.drain_experience_records_async("debug") + except Exception: + experiences = [] if not status.ok and not experiences: - experiences = self.model_wrapper.extract_experience_from_history() - self.logger.info( - f"Debugging failed, extracting {len(experiences)} experiences from history." - ) + try: + experiences = self.model_wrapper.extract_experience_from_history() + self.logger.info( + f"Debugging failed, extracting {len(experiences)} experiences from history." + ) + except Exception: + experiences = [] await self.sqlite_writer.write(experiences) if status.ok: print(f"Task {task.task_id} completed successfully with metrics:\n{status.metrics}") From 9ad8797b468420b6b06266fd1e833bb0f5fb34ca Mon Sep 17 00:00:00 2001 From: "panxuchen.pxc" Date: Mon, 29 Jun 2026 17:12:03 +0800 Subject: [PATCH 36/84] remove query layer --- tests/common/sglang_test.py | 88 +++------- tests/common/vllm_test.py | 90 +++------- tests/explorer/explorer_test.py | 6 +- tests/trainer/trainer_test.py | 6 +- trinity/common/config.py | 9 +- trinity/common/models/recording/query.py | 159 ------------------ trinity/common/models/recording/server.py | 15 +- .../models/sglang_patch/recording/__init__.py | 2 - .../models/sglang_patch/recording/server.py | 6 +- .../models/sglang_patch/server_patch.py | 4 +- .../models/vllm_patch/recording/__init__.py | 2 - .../models/vllm_patch/recording/server.py | 7 +- trinity/explorer/proxy/service.py | 6 +- 13 files changed, 70 insertions(+), 330 deletions(-) delete mode 100644 trinity/common/models/recording/query.py diff --git a/tests/common/sglang_test.py b/tests/common/sglang_test.py index 289e5598303..ee32a7a8de4 100644 --- a/tests/common/sglang_test.py +++ b/tests/common/sglang_test.py @@ -372,9 +372,8 @@ class TestRecording(RayUnittestBaseAsync): Mirrors ``tests/common/vllm_test.py::TestRecording``. Verifies that every call path lands its finished turn in the in-process ``MemoryStore`` under - the right ``record_key``, and that ``POST /records/update_record`` flushes - the recorder, reward-stamps the whole record-key group, pops it, and returns - it as serialized experiences. + the right ``record_key``, and that actor-side reward update + drain APIs + stamp and return recorded experiences. Paths covered (all async): * Ray-direct ``generate`` / ``chat`` — SGLang's Ray-direct path is over @@ -437,41 +436,11 @@ async def asyncTearDown(self): await self.model_wrapper.shutdown() await super().asyncTearDown() - # -- /records store query/consume helpers --------------------------------- + # -- actor-side recording store helpers ----------------------------------- - async def _flush(self): - """Drain the recorder without popping anything (an empty consume).""" - resp = await self._http.post( - f"{self.api_address}/records/update_record", json={"updates": []} - ) - resp.raise_for_status() - - async def _list_record_keys(self): - resp = await self._http.get(f"{self.api_address}/records") - resp.raise_for_status() - return resp.json()["record_keys"] - - async def _get_record_experiences(self, record_key: str) -> dict: - resp = await self._http.get(f"{self.api_address}/records/{record_key}") - resp.raise_for_status() - return resp.json() - - async def _get_request_experience(self, record_key: str, request_id: str) -> Experience: - resp = await self._http.get(f"{self.api_address}/records/{record_key}/request/{request_id}") - resp.raise_for_status() - return Experience.deserialize(resp.content) - - async def _consume( - self, record_key: str, reward: float, run: int, task: str - ) -> list[Experience]: - resp = await self._http.post( - f"{self.api_address}/records/update_record", - json={ - "updates": [{"record_key": record_key, "reward": reward, "run": run, "task": task}] - }, - ) - resp.raise_for_status() - return Experience.deserialize_many(resp.content) + async def _consume(self, record_key: str, reward: float) -> list[Experience]: + await self.model_wrapper.update_experience_reward_async(record_key, reward=reward) + return await self.model_wrapper.drain_experience_records_async(record_key) async def _openai_client(self, record_key: str) -> openai.AsyncOpenAI: # record_key travels as the Bearer api_key -> RecordingIdentityMiddleware. @@ -517,54 +486,40 @@ async def test_record(self): # noqa: C901 ] # ===== 1. Ray-direct generate (record_key via Authorization bearer) ===== - rk_gen = "trinity_record_generate" + rk_gen = "0/t_gen/1" await self.model_wrapper.generate_async( ["Hello, world!"], n=1, temperature=1.0, max_tokens=16, record_key=rk_gen ) - await self._flush() - self.assertIn(rk_gen, await self._list_record_keys()) - task = await self._get_record_experiences(rk_gen) - self.assertEqual(len(task["experiences"]), 1) - # blob endpoint round-trips a full experience - request_id = task["experiences"][0]["eid"]["suffix"] - blob_exp = await self._get_request_experience(rk_gen, request_id) - self._assert_recorded_experience(blob_exp, rk_gen) - self._assert_recorded_routed_experts(blob_exp) - consumed = await self._consume(rk_gen, reward=0.5, run=1, task="t_gen") + consumed = await self._consume(rk_gen, reward=0.5) self.assertEqual(len(consumed), 1) self.assertEqual(consumed[0].reward, 0.5) self.assertEqual(consumed[0].eid.run, 1) self.assertEqual(consumed[0].eid.task, "t_gen") self._assert_recorded_experience(consumed[0], rk_gen) self._assert_recorded_routed_experts(consumed[0]) - self.assertNotIn(rk_gen, await self._list_record_keys()) # popped # ===== 2. Ray-direct chat, n=2 (one record-key group, two samples) ===== - rk_chat = "trinity_record_chat" + rk_chat = "0/t_chat/2" chat_exps = await self.model_wrapper.chat_async( messages, n=2, temperature=1.0, max_tokens=16, record_key=rk_chat ) self.assertEqual(len(chat_exps), 2) - await self._flush() - task = await self._get_record_experiences(rk_chat) - self.assertEqual(len(task["experiences"]), 2) + consumed = await self._consume(rk_chat, reward=0.8) + self.assertEqual(len(consumed), 2) # SGLang expands n=2 parallel sampling into two scheduler requests. # The list position becomes sample_index (0, 1) to order the two # samples within the record-key group. - self.assertEqual(sorted(t["info"]["sample_index"] for t in task["experiences"]), [0, 1]) - self.assertEqual(len({t["eid"]["suffix"] for t in task["experiences"]}), 2) - consumed = await self._consume(rk_chat, reward=0.8, run=2, task="t_chat") - self.assertEqual(len(consumed), 2) + self.assertEqual(sorted(exp.info["sample_index"] for exp in consumed), [0, 1]) + self.assertEqual(len({exp.eid.suffix for exp in consumed}), 2) for exp in consumed: self.assertEqual(exp.reward, 0.8) self.assertEqual(exp.eid.run, 2) self.assertEqual(exp.eid.task, "t_chat") self._assert_recorded_experience(exp, rk_chat) self._assert_recorded_routed_experts(exp) - self.assertNotIn(rk_chat, await self._list_record_keys()) # ===== 3. OpenAI regular (HTTP; record_key = Bearer api_key) ===== - rk_oai = "trinity_record_openai" + rk_oai = "0/t_oai/3" client = await self._openai_client(rk_oai) model_id = await self._get_model_id(client) resp = await client.chat.completions.create( @@ -574,16 +529,15 @@ async def test_record(self): # noqa: C901 temperature=0.7, max_tokens=32, ) - consumed = await self._consume(rk_oai, reward=0.3, run=3, task="t_oai") + consumed = await self._consume(rk_oai, reward=0.3) self.assertEqual(len(consumed), 1) self._assert_recorded_experience(consumed[0], rk_oai) self._assert_recorded_routed_experts(consumed[0]) # No reasoning_parser is configured, so message.content == ret.text. self.assertEqual(consumed[0].response_text, resp.choices[0].message.content) - self.assertNotIn(rk_oai, await self._list_record_keys()) # ===== 4. OpenAI streaming (HTTP) ===== - rk_str = "trinity_record_stream" + rk_str = "0/t_str/4" sclient = await self._openai_client(rk_str) stream = await sclient.chat.completions.create( model=model_id, @@ -599,7 +553,7 @@ async def test_record(self): # noqa: C901 if delta: content += delta self.assertGreater(len(content), 0) - consumed = await self._consume(rk_str, reward=0.1, run=4, task="t_str") + consumed = await self._consume(rk_str, reward=0.1) self.assertEqual(len(consumed), 1) self._assert_recorded_experience(consumed[0], rk_str) self._assert_recorded_routed_experts(consumed[0]) @@ -607,10 +561,9 @@ async def test_record(self): # noqa: C901 decoded_content = self.tokenizer.decode(response_token_ids, skip_special_tokens=True) self.assertEqual(decoded_content, content) self.assertEqual(consumed[0].response_text, content) - self.assertNotIn(rk_str, await self._list_record_keys()) # ===== 5. OpenAI tool-call parsing (HTTP) ===== - rk_tool = "trinity_record_tool" + rk_tool = "0/t_tool/5" tclient = await self._openai_client(rk_tool) tools = [ { @@ -641,7 +594,7 @@ async def test_record(self): # noqa: C901 max_tokens=64, extra_body=no_think, ) - consumed = await self._consume(rk_tool, reward=1.0, run=5, task="t_tool") + consumed = await self._consume(rk_tool, reward=1.0) self.assertEqual(len(consumed), 1) self._assert_recorded_experience(consumed[0], rk_tool) self._assert_recorded_routed_experts(consumed[0]) @@ -657,7 +610,6 @@ async def test_record(self): # noqa: C901 if choice.finish_reason == "tool_calls" and choice.message.tool_calls: for tc in choice.message.tool_calls: self.assertIn(tc.function.name, consumed[0].response_text) - self.assertNotIn(rk_tool, await self._list_record_keys()) # ===== global: every group consumed -> store is drained ===== - self.assertEqual(await self._list_record_keys(), []) + self.assertEqual(await self.model_wrapper.delete_experience_records_async("0"), 0) diff --git a/tests/common/vllm_test.py b/tests/common/vllm_test.py index 7a44c67c784..73ccb121d09 100644 --- a/tests/common/vllm_test.py +++ b/tests/common/vllm_test.py @@ -1816,9 +1816,8 @@ class TestRecording(VLLMTestBase): """Correctness of the in-vLLM generation recording flow (``enable_history``). Verifies that every call path lands its finished turn in the in-process - ``MemoryStore`` under the right ``record_key``, and that - ``POST /records/update_record`` flushes the recorder, reward-stamps the - whole record-key group, pops it, and returns it as serialized experiences. + ``MemoryStore`` under the right ``record_key``, and that actor-side + reward update + drain APIs stamp and return recorded experiences. Paths covered (all async): * Ray-direct ``generate`` / ``chat`` — record_key propagated via @@ -1881,41 +1880,11 @@ async def asyncTearDown(self): await self._http.aclose() await super().asyncTearDown() - # -- /records store query/consume helpers --------------------------------- - - async def _flush(self): - """Drain the recorder without popping anything (an empty consume).""" - resp = await self._http.post( - f"{self.api_address}/records/update_record", json={"updates": []} - ) - resp.raise_for_status() - - async def _list_record_keys(self): - resp = await self._http.get(f"{self.api_address}/records") - resp.raise_for_status() - return resp.json()["record_keys"] - - async def _get_record_experiences(self, record_key: str) -> dict: - resp = await self._http.get(f"{self.api_address}/records/{record_key}") - resp.raise_for_status() - return resp.json() - - async def _get_request_experience(self, record_key: str, request_id: str) -> Experience: - resp = await self._http.get(f"{self.api_address}/records/{record_key}/request/{request_id}") - resp.raise_for_status() - return Experience.deserialize(resp.content) - - async def _consume( - self, record_key: str, reward: float, run: int, task: str - ) -> list[Experience]: - resp = await self._http.post( - f"{self.api_address}/records/update_record", - json={ - "updates": [{"record_key": record_key, "reward": reward, "run": run, "task": task}] - }, - ) - resp.raise_for_status() - return Experience.deserialize_many(resp.content) + # -- actor-side recording store helpers ----------------------------------- + + async def _consume(self, record_key: str, reward: float) -> list[Experience]: + await self.model_wrapper.update_experience_reward_async(record_key, reward=reward) + return await self.model_wrapper.drain_experience_records_async(record_key) async def _openai_client(self, record_key: str) -> openai.AsyncOpenAI: # record_key travels as the Bearer api_key -> RecordingIdentityMiddleware. @@ -1965,53 +1934,39 @@ async def test_record(self): # noqa: C901 no_think = {"chat_template_kwargs": {"enable_thinking": False}} # ===== 1. Ray-direct generate (record_key via record_key_ctx) ===== - rk_gen = "trinity_record_generate" + rk_gen = "0/t_gen/1" await self.model_wrapper.generate_async( ["Hello, world!"], n=1, temperature=1.0, max_tokens=16, record_key=rk_gen ) - await self._flush() - self.assertIn(rk_gen, await self._list_record_keys()) - task = await self._get_record_experiences(rk_gen) - self.assertEqual(len(task["experiences"]), 1) - # blob endpoint round-trips a full experience - request_id = task["experiences"][0]["eid"]["suffix"] - blob_exp = await self._get_request_experience(rk_gen, request_id) - self._assert_recorded_experience(blob_exp, rk_gen) - self._assert_recorded_routed_experts(blob_exp) - consumed = await self._consume(rk_gen, reward=0.5, run=1, task="t_gen") + consumed = await self._consume(rk_gen, reward=0.5) self.assertEqual(len(consumed), 1) self.assertEqual(consumed[0].reward, 0.5) self.assertEqual(consumed[0].eid.run, 1) self.assertEqual(consumed[0].eid.task, "t_gen") self._assert_recorded_experience(consumed[0], rk_gen) self._assert_recorded_routed_experts(consumed[0]) - self.assertNotIn(rk_gen, await self._list_record_keys()) # popped # ===== 2. Ray-direct chat, n=2 (one record-key group, two samples) ===== - rk_chat = "trinity_record_chat" + rk_chat = "0/t_chat/2" chat_exps = await self.model_wrapper.chat_async( messages, n=2, temperature=1.0, max_tokens=16, record_key=rk_chat ) self.assertEqual(len(chat_exps), 2) - await self._flush() - task = await self._get_record_experiences(rk_chat) - self.assertEqual(len(task["experiences"]), 2) + consumed = await self._consume(rk_chat, reward=0.8) + self.assertEqual(len(consumed), 2) # n=2 of one engine request -> two completions distinguished by # sample_index and a sample-qualified EID suffix. - self.assertEqual(sorted(t["info"]["sample_index"] for t in task["experiences"]), [0, 1]) - self.assertEqual(len({t["eid"]["suffix"] for t in task["experiences"]}), 2) - consumed = await self._consume(rk_chat, reward=0.8, run=2, task="t_chat") - self.assertEqual(len(consumed), 2) + self.assertEqual(sorted(exp.info["sample_index"] for exp in consumed), [0, 1]) + self.assertEqual(len({exp.eid.suffix for exp in consumed}), 2) for exp in consumed: self.assertEqual(exp.reward, 0.8) self.assertEqual(exp.eid.run, 2) self.assertEqual(exp.eid.task, "t_chat") self._assert_recorded_experience(exp, rk_chat) self._assert_recorded_routed_experts(exp) - self.assertNotIn(rk_chat, await self._list_record_keys()) # ===== 3. OpenAI regular (HTTP; record_key = Bearer api_key) ===== - rk_oai = "trinity_record_openai" + rk_oai = "0/t_oai/3" client = await self._openai_client(rk_oai) model_id = await self._get_model_id(client) resp = await client.chat.completions.create( @@ -2022,15 +1977,14 @@ async def test_record(self): # noqa: C901 max_tokens=32, extra_body=no_think, ) - consumed = await self._consume(rk_oai, reward=0.3, run=3, task="t_oai") + consumed = await self._consume(rk_oai, reward=0.3) self.assertEqual(len(consumed), 1) self._assert_recorded_experience(consumed[0], rk_oai) self._assert_recorded_routed_experts(consumed[0]) self.assertEqual(consumed[0].response_text, resp.choices[0].message.content) - self.assertNotIn(rk_oai, await self._list_record_keys()) # ===== 4. OpenAI streaming (HTTP) ===== - rk_str = "trinity_record_stream" + rk_str = "0/t_str/4" sclient = await self._openai_client(rk_str) stream = await sclient.chat.completions.create( model=model_id, @@ -2047,7 +2001,7 @@ async def test_record(self): # noqa: C901 if delta: content += delta self.assertGreater(len(content), 0) - consumed = await self._consume(rk_str, reward=0.1, run=4, task="t_str") + consumed = await self._consume(rk_str, reward=0.1) self.assertEqual(len(consumed), 1) self._assert_recorded_experience(consumed[0], rk_str) self._assert_recorded_routed_experts(consumed[0]) @@ -2055,10 +2009,9 @@ async def test_record(self): # noqa: C901 decoded_content = self.tokenizer.decode(response_token_ids, skip_special_tokens=True) self.assertEqual(decoded_content, content) self.assertEqual(consumed[0].response_text, content) - self.assertNotIn(rk_str, await self._list_record_keys()) # ===== 5. OpenAI tool usage (HTTP) ===== - rk_tool = "trinity_record_tool" + rk_tool = "0/t_tool/5" tclient = await self._openai_client(rk_tool) tools = [ { @@ -2088,7 +2041,7 @@ async def test_record(self): # noqa: C901 max_tokens=64, extra_body=no_think, ) - consumed = await self._consume(rk_tool, reward=1.0, run=5, task="t_tool") + consumed = await self._consume(rk_tool, reward=1.0) self.assertEqual(len(consumed), 1) self._assert_recorded_experience(consumed[0], rk_tool) self._assert_recorded_routed_experts(consumed[0]) @@ -2101,10 +2054,9 @@ async def test_record(self): # noqa: C901 if choice.finish_reason == "tool_calls" and choice.message.tool_calls: for tc in choice.message.tool_calls: self.assertIn(tc.function.name, consumed[0].response_text) - self.assertNotIn(rk_tool, await self._list_record_keys()) # ===== global: every group consumed -> store is drained ===== - self.assertEqual(await self._list_record_keys(), []) + self.assertEqual(await self.model_wrapper.delete_experience_records_async("0"), 0) class TestSuperLongGeneration(VLLMTestBase): diff --git a/tests/explorer/explorer_test.py b/tests/explorer/explorer_test.py index 93eaec0292b..a3806eeea7f 100644 --- a/tests/explorer/explorer_test.py +++ b/tests/explorer/explorer_test.py @@ -439,9 +439,9 @@ def run_agent(proxy_url, model_path: str, stream: bool): @unittest.skip( - "serve-mode experience collection moved to the in-vLLM MemoryStore + " - "/records/update_record flow; the proxy /feedback//commit path and external " - "reward reporting are being redesigned (see recording refactor plan)." + "serve-mode experience collection moved to rollout model-side recording stores; " + "the proxy /feedback//commit path and external reward reporting are being " + "redesigned (see recording refactor plan)." ) class ServeTest(RayUnittestBaseAsync): def setUp(self): diff --git a/tests/trainer/trainer_test.py b/tests/trainer/trainer_test.py index 6865aa275ec..8aaf5bea5eb 100644 --- a/tests/trainer/trainer_test.py +++ b/tests/trainer/trainer_test.py @@ -1099,9 +1099,9 @@ async def run_math_workflow(serve_url: str, task: dict): @unittest.skip( - "serve-mode experience collection moved to the in-vLLM MemoryStore + " - "/records/update_record flow; the proxy /feedback//commit path and external " - "reward reporting are being redesigned (see recording refactor plan)." + "serve-mode experience collection moved to rollout model-side recording stores; " + "the proxy /feedback//commit path and external reward reporting are being " + "redesigned (see recording refactor plan)." ) class TestServeWithTrainer(RayUnittestBaseAsync): def setUp(self): diff --git a/trinity/common/config.py b/trinity/common/config.py index dbf907584b3..751e9532d89 100644 --- a/trinity/common/config.py +++ b/trinity/common/config.py @@ -572,11 +572,10 @@ class InferenceModelConfig: # Trinity ``Experience`` to the in-process ``MemoryStore``, keyed by the # recording identity (``record_key``). This is the single switch for the # recording flow — when on, the explorer also consumes from the store: the - # WorkflowRunner ships only a small reward map keyed by ``record_key`` and - # the RolloutCoordinator pulls heavy experiences via ``/records/drain`` - # at finalize time. When off (default), runners ship serialized experiences - # through the scheduler as before (legacy path). When True, the Allocator - # forces ``enable_openai_api`` (the recorder runs in the API server). + # Workflow updates rewards by ``record_key`` and the Scheduler drains + # completed task experiences through rollout model actor methods. When off + # (default), recording APIs are disabled. When True, the Allocator forces + # ``enable_openai_api`` for engine-side request capture. # VLLMModel mirrors the recording config onto the engine instance for the # recorder to read. The capture width (top-k logprobs) reuses ``logprobs`` # below (default 1). Requires ``enable_openai_api=True`` (the recording diff --git a/trinity/common/models/recording/query.py b/trinity/common/models/recording/query.py deleted file mode 100644 index 719256f9c49..00000000000 --- a/trinity/common/models/recording/query.py +++ /dev/null @@ -1,159 +0,0 @@ -"""HTTP endpoints over recorded generation experiences.""" - -from typing import List - -from fastapi import APIRouter, HTTPException, Request, Response -from pydantic import BaseModel - -from trinity.buffer.store import RecordStore -from trinity.common.experience import Experience -from trinity.common.models.recording.recorder import ( - TRINITY_RECORD_STORE_ATTR, - TRINITY_RECORDER_ATTR, - Recorder, -) - -STORE_STATE_ATTR = TRINITY_RECORD_STORE_ATTR -RECORDER_STATE_ATTR = TRINITY_RECORDER_ATTR - -query_router = APIRouter(prefix="/records", tags=["trinity-recording"]) - - -class _RewardUpdateRequest(BaseModel): - record_key: str - reward: float - info: dict | None = None - sample_ids: List[str] | None = None - - -class _PrefixRequest(BaseModel): - prefix: str - - -def _store(request: Request) -> RecordStore: - store = getattr(request.app.state, STORE_STATE_ATTR, None) - if store is None: - raise HTTPException(status_code=503, detail="recording store not configured") - return store - - -def _recorder(request: Request) -> Recorder: - rec = getattr(request.app.state, RECORDER_STATE_ATTR, None) - if rec is None: - raise HTTPException(status_code=503, detail="recorder not configured") - return rec - - -def _get_exp(store: RecordStore, record_key: str, request_id: str) -> Experience: - for exp in store.get(record_key): - if exp.eid.suffix == request_id: - return exp - raise HTTPException(status_code=404, detail="experience not found") - - -def _forget_record(request: Request, record_key: str) -> None: - rec = getattr(request.app.state, RECORDER_STATE_ATTR, None) - if rec is not None: - rec.forget_record(record_key) - - -@query_router.get("") -async def list_records(request: Request) -> dict: - store = _store(request) - return {"record_keys": store.keys()} - - -@query_router.get("/{record_key}") -async def get_record_experiences(record_key: str, request: Request) -> dict: - store = _store(request) - experiences = store.get(record_key) - return {"record_key": record_key, "experiences": [e.to_dict() for e in experiences]} - - -@query_router.get("/{record_key}/request/{request_id}") -async def get_request_experience(record_key: str, request_id: str, request: Request) -> Response: - store = _store(request) - exp = _get_exp(store, record_key, request_id) - return Response( - content=Experience.serialize(exp), - media_type="application/octet-stream", - headers={"Content-Disposition": f'attachment; filename="{request_id}.bin"'}, - ) - - -@query_router.delete("/{record_key}") -async def delete_record_experiences(record_key: str, request: Request) -> dict: - store = _store(request) - store.remove(record_key) - _forget_record(request, record_key) - return {"record_key": record_key, "deleted": True} - - -@query_router.delete("/{record_key}/request/{request_id}") -async def delete_request_experience(record_key: str, request_id: str, request: Request) -> dict: - store = _store(request) - kept = [] - deleted = False - for exp in store.get(record_key): - if exp.eid.suffix == request_id: - deleted = True - else: - kept.append(exp) - if not deleted: - raise HTTPException(status_code=404, detail="experience not found") - if kept: - store.overwrite(record_key, kept) - else: - store.remove(record_key) - _forget_record(request, record_key) - return {"record_key": record_key, "request_id": request_id, "deleted": True} - - -@query_router.post("/update_reward") -async def update_reward(req: _RewardUpdateRequest, request: Request) -> dict: - store = _store(request) - recorder = _recorder(request) - await recorder.flush() - if not store.get(req.record_key): - return {"record_key": req.record_key, "updated": 0} - store.update( - key=req.record_key, - reward=req.reward, - info=req.info, - sample_ids=req.sample_ids, - ) - return { - "record_key": req.record_key, - "updated": ( - len(req.sample_ids) if req.sample_ids is not None else len(store.get(req.record_key)) - ), - } - - -@query_router.post("/drain") -async def drain_records(req: _PrefixRequest, request: Request) -> Response: - store = _store(request) - recorder = _recorder(request) - await recorder.flush() - matched_keys = [ - key for key in store.keys() if key == req.prefix or key.startswith(f"{req.prefix}/") - ] - exps = store.remove(req.prefix) - for key in matched_keys: - recorder.forget_record(key) - return Response( - content=Experience.serialize_many(exps), - media_type="application/octet-stream", - ) - - -@query_router.delete("") -async def delete_records(req: _PrefixRequest, request: Request) -> dict: - store = _store(request) - matched_keys = [ - key for key in store.keys() if key == req.prefix or key.startswith(f"{req.prefix}/") - ] - deleted = len(store.remove(req.prefix)) - for key in matched_keys: - _forget_record(request, key) - return {"prefix": req.prefix, "deleted": deleted} diff --git a/trinity/common/models/recording/server.py b/trinity/common/models/recording/server.py index 42de0e909cd..ef55c02da24 100644 --- a/trinity/common/models/recording/server.py +++ b/trinity/common/models/recording/server.py @@ -3,12 +3,14 @@ import logging from trinity.common.models.recording.context import RecordingIdentityMiddleware -from trinity.common.models.recording.query import ( - RECORDER_STATE_ATTR, - STORE_STATE_ATTR, - query_router, +from trinity.common.models.recording.recorder import ( + TRINITY_RECORD_STORE_ATTR, + TRINITY_RECORDER_ATTR, + Recorder, ) -from trinity.common.models.recording.recorder import Recorder + +STORE_STATE_ATTR = TRINITY_RECORD_STORE_ATTR +RECORDER_STATE_ATTR = TRINITY_RECORDER_ATTR def add_recording_middleware(app) -> None: @@ -31,9 +33,8 @@ def mount_recording_api( engine_name: str, start_recorder: bool = False, ) -> None: - """Mount recording middleware/query routes and expose state to handlers.""" + """Mount recording middleware and expose state to the server process.""" add_recording_middleware(app) - app.include_router(query_router) setattr(app.state, STORE_STATE_ATTR, recorder.store) setattr(app.state, RECORDER_STATE_ATTR, recorder) diff --git a/trinity/common/models/sglang_patch/recording/__init__.py b/trinity/common/models/sglang_patch/recording/__init__.py index 6ddd3f980f0..4811f6e518c 100644 --- a/trinity/common/models/sglang_patch/recording/__init__.py +++ b/trinity/common/models/sglang_patch/recording/__init__.py @@ -13,7 +13,6 @@ record_key_ctx, skip_recording_ctx, ) -from trinity.common.models.recording.query import query_router # noqa: F401 from trinity.common.models.recording.recorder import Recorder # noqa: F401 from trinity.common.models.sglang_patch.recording.models import ( # noqa: F401 build_sglang_experience, @@ -35,7 +34,6 @@ "create_sglang_recorder", "get_recording_record_key", "patch_tokenizer_manager_for_recording", - "query_router", "record_key_ctx", "setup_sglang_recording", "skip_recording_ctx", diff --git a/trinity/common/models/sglang_patch/recording/server.py b/trinity/common/models/sglang_patch/recording/server.py index 704d2cd7133..b82ce9a3985 100644 --- a/trinity/common/models/sglang_patch/recording/server.py +++ b/trinity/common/models/sglang_patch/recording/server.py @@ -3,14 +3,14 @@ Mirrors ``trinity/common/models/vllm_patch/recording/server.py:_setup_recording``: (1) the engine wrap (``create_sglang_recorder``), (2) ``RecordingIdentityMiddleware`` — an in-process ASGI middleware reading ``Authorization: Bearer `` -into a contextvar, and (3) ``query_router`` — ``/records/*`` endpoints. +into a contextvar, and (3) actor-side recording APIs over the model-owned store. Called from ``sglang_patch.server_patch.get_api_server`` after the ``tokenizer_manager`` is created and **before** the uvicorn task starts serving, so the middleware/router are mounted on ``app`` in time. The recorder and store are owned by ``SGLangRolloutModel`` (passed in) so it can drain them in-process -via ``extract_experience_from_history``; they are also stashed on ``app.state`` -for the ``query_router`` HTTP drain path used by the coordinator. +via actor methods; they are also stashed on ``app.state`` for server-local +recording lifecycle management. """ import logging diff --git a/trinity/common/models/sglang_patch/server_patch.py b/trinity/common/models/sglang_patch/server_patch.py index 46f1c1ad5f3..8829823d80b 100644 --- a/trinity/common/models/sglang_patch/server_patch.py +++ b/trinity/common/models/sglang_patch/server_patch.py @@ -383,8 +383,8 @@ def get_api_server( # Wire generation recording before the uvicorn task starts serving. The # recorder/store are owned by ``SGLangRolloutModel``; this installs the - # engine wrap on ``tokenizer_manager``, ``RecordingIdentityMiddleware`` and - # ``query_router`` on ``app``, and stashes store/recorder on ``app.state``. + # engine wrap on ``tokenizer_manager`` and ``RecordingIdentityMiddleware`` + # on ``app``, and stashes store/recorder on ``app.state``. if enable_history: from trinity.common.models.sglang_patch.recording import setup_sglang_recording diff --git a/trinity/common/models/vllm_patch/recording/__init__.py b/trinity/common/models/vllm_patch/recording/__init__.py index 971db1ce81a..857d865bb54 100644 --- a/trinity/common/models/vllm_patch/recording/__init__.py +++ b/trinity/common/models/vllm_patch/recording/__init__.py @@ -8,7 +8,6 @@ record_key_ctx, skip_recording_ctx, ) -from trinity.common.models.recording.query import query_router from trinity.common.models.recording.recorder import Recorder from trinity.common.models.vllm_patch.recording.models import build_experience from trinity.common.models.vllm_patch.recording.recorder import ( @@ -27,7 +26,6 @@ "build_experience", "create_vllm_recorder", "patch_engine_for_recording", - "query_router", "record_key_ctx", "run_api_server_with_recording", "skip_recording_ctx", diff --git a/trinity/common/models/vllm_patch/recording/server.py b/trinity/common/models/vllm_patch/recording/server.py index b8283b9cc40..317e5d8336e 100644 --- a/trinity/common/models/vllm_patch/recording/server.py +++ b/trinity/common/models/vllm_patch/recording/server.py @@ -11,7 +11,7 @@ 1. ``vLLMRolloutModel`` owns the recorder and attaches it to ``async_llm``. 2. ``RecordingIdentityMiddleware`` — in-process ASGI middleware reading ``Authorization: Bearer `` into a contextvar. - 3. ``query_router`` — ``/records/*`` endpoints for later analysis. + 3. Actor-side recording APIs drain/update the model-owned store. Only for vllm versions >= 0.17.0. """ @@ -108,9 +108,8 @@ def _setup_recording( version is read live off ``engine_client.trinity_model_version`` (mirrored by VLLMModel at engine creation and in ``sync_model_weights``). - The store backend is always the in-process ``MemoryStore``; the coordinator - drains it at finalize time via ``/records/drain`` (fanned out per - rank), so heavy experience bytes never touch SQL or Ray serialization. + The store backend is always the in-process ``MemoryStore``; the scheduler + drains completed task records through rollout model actor methods. Args: engine_client: AsyncLLM instance with ``trinity_recorder`` already set diff --git a/trinity/explorer/proxy/service.py b/trinity/explorer/proxy/service.py index 5261d6ceb6a..4385a7adfb1 100644 --- a/trinity/explorer/proxy/service.py +++ b/trinity/explorer/proxy/service.py @@ -14,9 +14,9 @@ class ExplorerService: The proxy is a request router + model-weight sync coordinator for serve mode. Experience collection used to live here (SQL-mediated - ``/feedback``/``/commit``); it has been removed in favor of the in-vLLM - ``MemoryStore`` + ``/records/drain`` flow driven by the rollout - coordinator. Serve-mode external reward reporting is therefore pending + ``/feedback``/``/commit``); it has been removed in favor of rollout + model-side recording stores drained through actor methods. Serve-mode + external reward reporting is therefore pending (see the recording refactor plan). """ From d42c27d93fc6306475f1bcd5812dc91a91adb159 Mon Sep 17 00:00:00 2001 From: "panxuchen.pxc" Date: Mon, 29 Jun 2026 18:27:27 +0800 Subject: [PATCH 37/84] clean workflow interface --- tests/common/sglang_test.py | 5 +- tests/common/vllm_test.py | 5 +- tests/explorer/scheduler_test.py | 13 +- tests/explorer/workflow_test.py | 21 ++- tests/utils/plugins/main.py | 11 +- tests/utils/plugins/my_workflow.py | 11 +- trinity/common/models/model.py | 20 +-- trinity/common/workflows/__init__.py | 12 +- .../envs/alfworld/RAFT_alfworld_workflow.py | 5 +- .../envs/sciworld/sciworld_workflow.py | 5 +- .../workflows/on_policy_distill_workflow.py | 5 +- trinity/common/workflows/workflow.py | 147 +++++++++++------- trinity/explorer/workflow_runner.py | 25 ++- 13 files changed, 151 insertions(+), 134 deletions(-) diff --git a/tests/common/sglang_test.py b/tests/common/sglang_test.py index ee32a7a8de4..a7aa337225e 100644 --- a/tests/common/sglang_test.py +++ b/tests/common/sglang_test.py @@ -440,7 +440,8 @@ async def asyncTearDown(self): async def _consume(self, record_key: str, reward: float) -> list[Experience]: await self.model_wrapper.update_experience_reward_async(record_key, reward=reward) - return await self.model_wrapper.drain_experience_records_async(record_key) + payload = await self.model_wrapper.drain_experience_records_bytes_async(record_key) + return Experience.deserialize_many(payload) async def _openai_client(self, record_key: str) -> openai.AsyncOpenAI: # record_key travels as the Bearer api_key -> RecordingIdentityMiddleware. @@ -612,4 +613,4 @@ async def test_record(self): # noqa: C901 self.assertIn(tc.function.name, consumed[0].response_text) # ===== global: every group consumed -> store is drained ===== - self.assertEqual(await self.model_wrapper.delete_experience_records_async("0"), 0) + await self.model_wrapper.delete_experience_records_async("0") diff --git a/tests/common/vllm_test.py b/tests/common/vllm_test.py index 73ccb121d09..06d760c377f 100644 --- a/tests/common/vllm_test.py +++ b/tests/common/vllm_test.py @@ -1884,7 +1884,8 @@ async def asyncTearDown(self): async def _consume(self, record_key: str, reward: float) -> list[Experience]: await self.model_wrapper.update_experience_reward_async(record_key, reward=reward) - return await self.model_wrapper.drain_experience_records_async(record_key) + payload = await self.model_wrapper.drain_experience_records_bytes_async(record_key) + return Experience.deserialize_many(payload) async def _openai_client(self, record_key: str) -> openai.AsyncOpenAI: # record_key travels as the Bearer api_key -> RecordingIdentityMiddleware. @@ -2056,7 +2057,7 @@ async def test_record(self): # noqa: C901 self.assertIn(tc.function.name, consumed[0].response_text) # ===== global: every group consumed -> store is drained ===== - self.assertEqual(await self.model_wrapper.delete_experience_records_async("0"), 0) + await self.model_wrapper.delete_experience_records_async("0") class TestSuperLongGeneration(VLLMTestBase): diff --git a/tests/explorer/scheduler_test.py b/tests/explorer/scheduler_test.py index 0f7846f777d..0d4f1ace6b5 100644 --- a/tests/explorer/scheduler_test.py +++ b/tests/explorer/scheduler_test.py @@ -16,14 +16,12 @@ from trinity.common.experience import EID, Experience from trinity.common.models.allocator import Allocator from trinity.common.models.model import InferenceModel, ModelWrapper -from trinity.common.workflows import WORKFLOWS, Task, Workflow +from trinity.common.workflows import WORKFLOWS, RepeatableWorkflow, Task, Workflow from trinity.explorer.scheduler import Scheduler @WORKFLOWS.register_module("dummy_workflow") -class DummyWorkflow(Workflow): - can_repeat: bool = True - +class DummyWorkflow(RepeatableWorkflow): def __init__(self, *, task, model, auxiliary_models): super().__init__(task=task, model=model, auxiliary_models=auxiliary_models) self.step_num = task.workflow_args.get("step_num", 1) @@ -187,8 +185,7 @@ def run(self): @WORKFLOWS.register_module("dummy_async_workflow") -class DummyAsyncWorkflow(Workflow): - can_repeat: bool = True +class DummyAsyncWorkflow(RepeatableWorkflow): is_async: bool = True def __init__(self, *, task, model, auxiliary_models): @@ -222,8 +219,7 @@ def run(self): @WORKFLOWS.register_module("dummy_workflow_with_state") -class DummyWorkflowWithState(Workflow): - can_repeat: bool = True +class DummyWorkflowWithState(RepeatableWorkflow): is_async: bool = True def __init__(self, *, task, model: ModelWrapper, auxiliary_models): @@ -258,7 +254,6 @@ async def run_async(self) -> List[Experience]: @WORKFLOWS.register_module("dummy_concurrent_workflow") class DummyConcurrentWorkflow(Workflow): - can_repeat: bool = False is_async: bool = True def __init__(self, *, task, model, auxiliary_models): diff --git a/tests/explorer/workflow_test.py b/tests/explorer/workflow_test.py index 782824bc9a9..520f5a3d620 100644 --- a/tests/explorer/workflow_test.py +++ b/tests/explorer/workflow_test.py @@ -29,10 +29,14 @@ from trinity.common.experience import EID, Experience from trinity.common.models.allocator import Allocator from trinity.common.models.model import ModelWrapper -from trinity.common.workflows import WORKFLOWS, Workflow +from trinity.common.workflows import WORKFLOWS, RepeatableWorkflow, Workflow from trinity.common.workflows.customized_math_workflows import MathBoxedWorkflow from trinity.common.workflows.eval_workflow import MathEvalWorkflow -from trinity.common.workflows.workflow import MathWorkflow, MultiTurnWorkflow, Task +from trinity.common.workflows.workflow import ( + MathWorkflow, + RepeatableMultiTurnWorkflow, + Task, +) from trinity.explorer.workflow_runner import WorkflowRunner @@ -63,9 +67,8 @@ class MockResponse: action_mask: Optional[Tensor] = None -class DummyWorkflow(Workflow): +class DummyWorkflow(RepeatableWorkflow): can_reset: bool = True - can_repeat: bool = True def __init__(self, model, task: Task, auxiliary_models=None): super().__init__(task=task, model=model, auxiliary_models=auxiliary_models) @@ -107,9 +110,8 @@ def run(self): raise ValueError("Invalid output format") -class DummyAsyncWorkflow(Workflow): +class DummyAsyncWorkflow(RepeatableWorkflow): can_reset: bool = True - can_repeat: bool = True is_async: bool = True def __init__(self, model, task: Task, auxiliary_models=None): @@ -153,9 +155,7 @@ async def run_async(self): raise ValueError("Invalid output format") -class DummyMultiTurnWorkflow(MultiTurnWorkflow): - can_repeat: bool = True - +class DummyMultiTurnWorkflow(RepeatableMultiTurnWorkflow): def __init__(self, model, task: Task, auxiliary_models=None): super().__init__(task=task, model=model, auxiliary_models=auxiliary_models) self.contents = task.raw_task["contents"] # type: ignore @@ -171,9 +171,8 @@ def run(self): return experience_list -class DummyAsyncMultiTurnWorkflow(MultiTurnWorkflow): +class DummyAsyncMultiTurnWorkflow(RepeatableMultiTurnWorkflow): is_async: bool = True - can_repeat: bool = True def __init__(self, model, task: Task, auxiliary_models=None): super().__init__(task=task, model=model, auxiliary_models=auxiliary_models) diff --git a/tests/utils/plugins/main.py b/tests/utils/plugins/main.py index fc07712658a..92129097295 100644 --- a/tests/utils/plugins/main.py +++ b/tests/utils/plugins/main.py @@ -1,17 +1,10 @@ from tests.utils.plugins.dependencies import DEPENDENCY_VALUE, dependency_func -from trinity.common.workflows.workflow import Workflow +from trinity.common.workflows.workflow import RepeatableWorkflow -class MainDummyWorkflow(Workflow): +class MainDummyWorkflow(RepeatableWorkflow): def __init__(self, *, task, model, auxiliary_models=None): super().__init__(task=task, model=model, auxiliary_models=auxiliary_models) - @property - def repeatable(self): - return True - - def set_repeat_times(self, repeat_times, run_id_base): - pass - def run(self) -> list: return [DEPENDENCY_VALUE, dependency_func()] diff --git a/tests/utils/plugins/my_workflow.py b/tests/utils/plugins/my_workflow.py index 471b2371cff..505f624271e 100644 --- a/tests/utils/plugins/my_workflow.py +++ b/tests/utils/plugins/my_workflow.py @@ -1,19 +1,12 @@ from typing import List -from trinity.common.workflows import WORKFLOWS, Workflow +from trinity.common.workflows import WORKFLOWS, RepeatableWorkflow @WORKFLOWS.register_module("my_workflow") -class MyWorkflow(Workflow): +class MyWorkflow(RepeatableWorkflow): def __init__(self, *, task, model, auxiliary_models=None): super().__init__(task=task, model=model, auxiliary_models=auxiliary_models) - @property - def repeatable(self): - return True - - def set_repeat_times(self, repeat_times, run_id_base): - pass - def run(self) -> List: return ["Hello world", "Hi"] diff --git a/trinity/common/models/model.py b/trinity/common/models/model.py index 8d966cb1474..3fd02341a5e 100644 --- a/trinity/common/models/model.py +++ b/trinity/common/models/model.py @@ -182,7 +182,7 @@ async def update_experience_reward( sample_ids=sample_ids, ) - async def drain_experience_records(self, prefix: str) -> List[Experience]: + async def _drain_experience_records(self, prefix: str) -> List[Experience]: """Remove and return recorded experiences matching a key or prefix.""" recorder = getattr(self, "recorder", None) if recorder is None: @@ -198,11 +198,11 @@ async def drain_experience_records(self, prefix: str) -> List[Experience]: async def drain_experience_records_bytes(self, prefix: str) -> bytes: """Remove matching recorded experiences and return serialized bytes.""" - return Experience.serialize_many(await self.drain_experience_records(prefix)) + return Experience.serialize_many(await self._drain_experience_records(prefix)) - async def delete_experience_records(self, prefix: str) -> int: + async def delete_experience_records(self, prefix: str) -> None: """Remove recorded experiences matching a key or prefix.""" - return len(await self.drain_experience_records(prefix)) + await self._drain_experience_records(prefix) def get_model_config(self) -> InferenceModelConfig: """Get the model configuration.""" @@ -845,14 +845,6 @@ async def update_experience_reward_async( sample_ids=sample_ids, ) - async def drain_experience_records_async(self, prefix: str) -> List[Experience]: - """Remove and return recorded experiences matching a key or prefix.""" - if not self.enable_history: - raise ValueError("History recording is not enabled.") - if self.model is None: - raise ValueError("Recording drain requires an inference model actor.") - return await self.model.drain_experience_records.remote(prefix=prefix) - async def drain_experience_records_bytes_async(self, prefix: str) -> bytes: """Remove matching recorded experiences and return serialized bytes.""" if not self.enable_history: @@ -861,13 +853,13 @@ async def drain_experience_records_bytes_async(self, prefix: str) -> bytes: raise ValueError("Recording drain requires an inference model actor.") return await self.model.drain_experience_records_bytes.remote(prefix=prefix) - async def delete_experience_records_async(self, prefix: str) -> int: + async def delete_experience_records_async(self, prefix: str) -> None: """Remove recorded experiences matching a key or prefix.""" if not self.enable_history: raise ValueError("History recording is not enabled.") if self.model is None: raise ValueError("Recording delete requires an inference model actor.") - return await self.model.delete_experience_records.remote(prefix=prefix) + await self.model.delete_experience_records.remote(prefix=prefix) # Workflow state management methods async def set_workflow_state(self, state: Dict) -> None: diff --git a/trinity/common/workflows/__init__.py b/trinity/common/workflows/__init__.py index e16e931e996..11f099d1e1e 100644 --- a/trinity/common/workflows/__init__.py +++ b/trinity/common/workflows/__init__.py @@ -1,6 +1,13 @@ # -*- coding: utf-8 -*- """Workflow module""" -from trinity.common.workflows.workflow import Status, Task, Workflow +from trinity.common.workflows.workflow import ( + Metrics, + RepeatableMultiTurnWorkflow, + RepeatableWorkflow, + Status, + Task, + Workflow, +) from trinity.utils.registry import Registry WORKFLOWS: Registry = Registry( @@ -58,6 +65,9 @@ __all__ = [ "Task", "Status", + "Metrics", + "RepeatableMultiTurnWorkflow", + "RepeatableWorkflow", "Workflow", "WORKFLOWS", ] diff --git a/trinity/common/workflows/envs/alfworld/RAFT_alfworld_workflow.py b/trinity/common/workflows/envs/alfworld/RAFT_alfworld_workflow.py index 4bc34833cb8..deb13de4990 100644 --- a/trinity/common/workflows/envs/alfworld/RAFT_alfworld_workflow.py +++ b/trinity/common/workflows/envs/alfworld/RAFT_alfworld_workflow.py @@ -13,10 +13,10 @@ process_messages_to_experience_async, validate_trajectory_format, ) -from trinity.common.workflows.workflow import Task, Workflow +from trinity.common.workflows.workflow import RepeatableWorkflow, Task -class RAFTAlfworldWorkflow(Workflow): +class RAFTAlfworldWorkflow(RepeatableWorkflow): """ RAFT workflow for alfworld using trajectory context. @@ -26,7 +26,6 @@ class RAFTAlfworldWorkflow(Workflow): """ can_reset: bool = True - can_repeat: bool = True is_async: bool = True def __init__( diff --git a/trinity/common/workflows/envs/sciworld/sciworld_workflow.py b/trinity/common/workflows/envs/sciworld/sciworld_workflow.py index d5aeb1f2178..fb2bd6ecfa7 100644 --- a/trinity/common/workflows/envs/sciworld/sciworld_workflow.py +++ b/trinity/common/workflows/envs/sciworld/sciworld_workflow.py @@ -4,7 +4,7 @@ from trinity.common.experience import Experience from trinity.common.models.model import ModelWrapper -from trinity.common.workflows.workflow import MultiTurnWorkflow, Task +from trinity.common.workflows.workflow import RepeatableMultiTurnWorkflow, Task SCIWORLD_SYSTEM_PROMPT = """ You are an agent, your job is to do some scientific experiment in a virtual text-based environment. @@ -55,11 +55,10 @@ def parse_action(response): return "" -class SciWorldWorkflow(MultiTurnWorkflow): +class SciWorldWorkflow(RepeatableMultiTurnWorkflow): """A workflow for sciworld task.""" is_async: bool = True - can_repeat: bool = True def __init__( self, diff --git a/trinity/common/workflows/on_policy_distill_workflow.py b/trinity/common/workflows/on_policy_distill_workflow.py index f84818d35c3..3fb67f5e628 100644 --- a/trinity/common/workflows/on_policy_distill_workflow.py +++ b/trinity/common/workflows/on_policy_distill_workflow.py @@ -17,10 +17,10 @@ from trinity.common.experience import Experience from trinity.common.models.model import ModelWrapper from trinity.common.rewards.qwen25_eval import verify_math_answer -from trinity.common.workflows.workflow import Task, Workflow +from trinity.common.workflows.workflow import RepeatableWorkflow, Task -class OnPolicyDistillWorkflow(Workflow): +class OnPolicyDistillWorkflow(RepeatableWorkflow): """On-policy distillation workflow. Computes and stores teacher_logprobs in experience.info. @@ -34,7 +34,6 @@ class OnPolicyDistillWorkflow(Workflow): is_async: bool = True can_reset: bool = True - can_repeat: bool = True def __init__( self, diff --git a/trinity/common/workflows/workflow.py b/trinity/common/workflows/workflow.py index 3947859acf3..b73b87b3d6b 100644 --- a/trinity/common/workflows/workflow.py +++ b/trinity/common/workflows/workflow.py @@ -33,6 +33,9 @@ def ok(self) -> bool: return self.completed_runs == self.total_runs +Metrics = Dict[str, float] + + @dataclass class Task(dict): """A Task class that defines a task and its associated reward function / workflow.""" @@ -94,10 +97,16 @@ def to_dict(self) -> dict: class Workflow: """The base workflow class. - A workflow is a runnable object that executes rollout logic and returns a - :class:`Status`. Training experiences are captured by the rollout model's - built-in recording path during generation; workflows should update rewards - on those recorded experiences before returning. + A workflow is a runnable object that implements rollout logic in ``run`` or + ``run_async``. User code normally returns metrics for the completed run + and lets :meth:`execute` assemble the internal :class:`Status` consumed by + the runner and scheduler. + + Training experiences are captured by the rollout model's built-in recording + path during generation; workflows should update rewards on those recorded + experiences before returning. Advanced workflows may override + :meth:`execute` directly when they need full control over partial success, + run ids, or error messages. Attributes: auxiliary_model_wrappers: List of ModelWrapper instances for auxiliary models. @@ -105,17 +114,8 @@ class Workflow: """ can_reset: bool = False # whether the workflow can be reset with a new task. If true, `reset()` must be implemented. - can_repeat: bool = False # whether the workflow can be repeated multiple times. If true, `set_repeat_times()` must be implemented. is_async: bool = False # whether the workflow runs in async mode. If true, `run_async()` must be implemented, else `run()` must be implemented. - # Mirrors ``config.explorer.rollout_model.enable_history``. When it is set, - # the rollout model records generated experiences in its model-side store - # under ``record_key`` groups. Workflows use that rollout-model history to - # update rewards before returning Status; the coordinator later drains a - # whole step by prefix and sends the recorded experiences to the - # ExperiencePipeline. - enable_history: bool = False - def __init__( self, *, @@ -134,7 +134,7 @@ def __init__( self.auxiliary_models = [m.get_openai_async_client() for m in auxiliary_models] else: self.auxiliary_models = [m.get_openai_client() for m in auxiliary_models] - self.run_id_base = 0 + self.run_id = 0 self.logger = get_logger(__name__) @property @@ -142,44 +142,43 @@ def resettable(self): """Deprecated, use cls.can_reset instead.""" return self.__class__.can_reset - @property - def repeatable(self): - """Deprecated, use cls.can_repeat instead. - A workflow is repeatable if it can be run multiple times within the run() or run_async() method. - """ - return self.__class__.can_repeat - @property def asynchronous(self): """Deprecated, use cls.is_async instead. Whether the workflow runs in async mode.""" return self.__class__.is_async + def set_execution_context(self, run_id: int) -> None: + """Set the execution context used by ``execute`` and recording helpers.""" + self.run_id = run_id + def reset(self, task: Task): """Reset the workflow.""" raise NotImplementedError - def set_repeat_times(self, repeat_times: int, run_id_base: int) -> None: - """ - Set the number of times to repeat the workflow. - Args: - repeat_times (int): number of times to repeat the workflow (if repeatable). - run_id_base (int): base run_id for recording keys and returned statuses. - """ - raise NotImplementedError( - "set_repeat_times() must be implemented for a repeatable workflow." - ) - - def run(self) -> Status: - """Run workflow and return its execution status.""" + def run(self) -> Metrics: + """Run workflow and return metrics for the completed run.""" raise NotImplementedError - async def run_async(self) -> Status: - """Run workflow in async and return its execution status.""" + async def run_async(self) -> Metrics: + """Run workflow asynchronously and return metrics for the completed run.""" raise NotImplementedError + async def execute(self) -> Status: + """Execute the workflow and normalize the user return value to Status.""" + result = await self.run_async() if self.asynchronous else self.run() + return self._to_status(result) + + def _to_status(self, result: Metrics) -> Status: + return Status( + completed_runs=1, + total_runs=1, + metrics=[result], + successful_run_ids=[self.run_id], + ) + def _build_record_key(self, run_id: Optional[int] = None) -> str: - run = self.run_id_base if run_id is None else run_id + run = self.run_id if run_id is None else run_id return f"{self.task.batch_id}/{self.task.task_id}/{run}" async def update_reward( @@ -190,8 +189,6 @@ async def update_reward( run_id: Optional[int] = None, ) -> None: """Update recorded experiences for one run with reward and optional info.""" - if not self.enable_history: - return await self.model.update_experience_reward_async( record_key=self._build_record_key(run_id), reward=reward, @@ -200,6 +197,57 @@ async def update_reward( ) +class RepeatableWorkflow(Workflow): + """Workflow base class for implementations that run repeats internally. + + ``run`` or ``run_async`` should return one metrics dict per successful + repeat. Workflows that need custom partial-success semantics can override + ``execute`` and return :class:`Status` directly. + """ + + def __init__( + self, + *, + task: Task, + model: ModelWrapper, + auxiliary_models: Optional[List[ModelWrapper]] = None, + ): + super().__init__(task=task, model=model, auxiliary_models=auxiliary_models) + self.repeat_times = 1 + self.run_id_base = 0 + + def set_repeat_times(self, repeat_times: int, run_id_base: int) -> None: + """Set repeat count and base run id for one repeatable execution.""" + self.repeat_times = repeat_times + self.run_id_base = run_id_base + self.run_id = run_id_base + + def set_execution_context(self, repeat_times: int, run_id_base: int) -> None: + """Set the execution context used by ``execute`` and recording helpers.""" + self.set_repeat_times(repeat_times, run_id_base) + + def run(self) -> List[Metrics]: + """Run workflow and return metrics for completed repeats.""" + raise NotImplementedError + + async def run_async(self) -> List[Metrics]: + """Run workflow asynchronously and return metrics for completed repeats.""" + raise NotImplementedError + + def _to_status(self, result: List[Metrics]) -> Status: + completed_runs = len(result) + return Status( + completed_runs=completed_runs, + total_runs=self.repeat_times, + metrics=result, + successful_run_ids=list(range(self.run_id_base, self.run_id_base + completed_runs)), + ) + + def _build_record_key(self, run_id: Optional[int] = None) -> str: + run = self.run_id_base if run_id is None else run_id + return f"{self.task.batch_id}/{self.task.task_id}/{run}" + + class MultiTurnWorkflow(Workflow): """ The base workflow class for concatenated multi-turn tasks. @@ -218,10 +266,6 @@ def __init__( auxiliary_models=auxiliary_models, ) - def set_repeat_times(self, repeat_times, run_id_base): - self.repeat_times = repeat_times - self.run_id_base = run_id_base - def _build_experience_from_converted( self, converted_experience, reward, info={}, truncate_status=None ) -> Experience: @@ -298,7 +342,11 @@ async def process_messages_to_experience_async( ) -class BaseSimpleWorkflow(Workflow): +class RepeatableMultiTurnWorkflow(MultiTurnWorkflow, RepeatableWorkflow): + """Multi-turn workflow variant that runs repeats inside one execute call.""" + + +class BaseSimpleWorkflow(RepeatableWorkflow): def __init__( self, *, @@ -384,17 +432,13 @@ class SimpleWorkflow(BaseSimpleWorkflow): """A workflow for simple single-round task.""" can_reset: bool = True - can_repeat: bool = True def run(self) -> List[Experience]: # TODO: Optimize the generate function messages = self.format_messages() self.logger.debug("start chat") - if self.enable_history: - return self._run_recorded(messages) - responses = self.model.chat(messages, **self.rollout_args) - return self._attach_rewards(responses, base=self.run_id_base) + return self._run_recorded(messages) def _run_recorded(self, messages) -> List[Experience]: # One chat call per sample (n=1) so each gets a distinct record_key @@ -420,10 +464,7 @@ async def run_async(self) -> List[Experience]: messages = self.format_messages() self.logger.info("start chat") - if self.enable_history: - return await self._run_recorded_async(messages) - responses = await self.model.chat_async(messages, **self.rollout_args) - return self._attach_rewards(responses, base=self.run_id_base) + return await self._run_recorded_async(messages) async def _run_recorded_async(self, messages) -> List[Experience]: rollout_args = dict(self.rollout_args) diff --git a/trinity/explorer/workflow_runner.py b/trinity/explorer/workflow_runner.py index db68046b77c..f9c0f764695 100644 --- a/trinity/explorer/workflow_runner.py +++ b/trinity/explorer/workflow_runner.py @@ -10,9 +10,10 @@ from trinity.buffer import get_buffer_reader, get_buffer_writer from trinity.common.config import Config, StorageConfig from trinity.common.constants import LOG_DIR_ENV_VAR, LOG_LEVEL_ENV_VAR +from trinity.common.experience import Experience from trinity.common.models.allocator import Allocator from trinity.common.models.model import ModelWrapper -from trinity.common.workflows import Status, Task, Workflow +from trinity.common.workflows import RepeatableWorkflow, Status, Task, Workflow from trinity.utils.log import get_logger @@ -86,7 +87,7 @@ def _build_record_key(self, task: Task, run_index: int) -> str: return f"{task.batch_id}/{task.task_id}/{run_index}" def _set_record_key(self, model_wrapper: ModelWrapper, record_key: Optional[str]) -> None: - if self._enable_history_recording() and record_key is not None: + if record_key is not None: model_wrapper.set_api_key(record_key) def _create_workflow_instance(self, task: Task, record_key: Optional[str] = None) -> Workflow: @@ -105,20 +106,13 @@ def _create_workflow_instance(self, task: Task, record_key: Optional[str] = None ) else: self.workflow_instance.reset(task) - self.workflow_instance.enable_history = self._enable_history_recording() return self.workflow_instance - def _enable_history_recording(self) -> bool: - return bool(self.config.explorer.rollout_model.enable_history) - async def _run_workflow(self, workflow_instance: Workflow) -> Status: - if workflow_instance.asynchronous: - status = await workflow_instance.run_async() - else: - status = workflow_instance.run() + status = await workflow_instance.execute() if not isinstance(status, Status): raise TypeError( - f"{workflow_instance.__class__.__name__}.run must return Status, " + f"{workflow_instance.__class__.__name__}.execute must return Status, " f"got {type(status).__name__}." ) return status @@ -136,7 +130,6 @@ def _create_isolated_workflow_instance( model_wrapper, self.auxiliary_model_wrappers, ) - wf.enable_history = self._enable_history_recording() return wf, model_wrapper def _build_status( @@ -268,6 +261,7 @@ async def _execute_single_run( self._set_record_key(model_wrapper, record_key) await model_wrapper.clean_workflow_state() run_id = run_id_base + run_index + workflow.set_execution_context(run_id=run_id) self.runner_state["workflow_id"] = self._build_record_key(task, run_id) self.runner_state["terminate_time"] = None self.runner_state["begin_time"] = st @@ -312,10 +306,10 @@ async def _run_task( collect_partial_runs: bool = True, ) -> Status: """Init workflow from the task and run it.""" - if task.workflow.can_repeat: + if issubclass(task.workflow, RepeatableWorkflow): record_key = self._build_record_key(task, run_id_base) workflow_instance = self._create_workflow_instance(task, record_key=record_key) - workflow_instance.set_repeat_times(repeat_times, run_id_base) + workflow_instance.set_execution_context(repeat_times, run_id_base) st = time.time() self._set_record_key(self.model_wrapper, record_key) await self.model_wrapper.clean_workflow_state() @@ -506,7 +500,8 @@ async def debug(self) -> None: experiences = [] if self.config.explorer.rollout_model.enable_history: try: - experiences = await self.model_wrapper.drain_experience_records_async("debug") + payload = await self.model_wrapper.drain_experience_records_bytes_async("debug") + experiences = Experience.deserialize_many(payload) if payload else [] except Exception: experiences = [] if not status.ok and not experiences: From 33422258b3ec48525bb02a7724dea8397f57f247 Mon Sep 17 00:00:00 2001 From: "panxuchen.pxc" Date: Mon, 29 Jun 2026 18:44:50 +0800 Subject: [PATCH 38/84] finish model wrapper interface --- trinity/common/models/model.py | 59 +++++++++++++++++++++++++--------- 1 file changed, 44 insertions(+), 15 deletions(-) diff --git a/trinity/common/models/model.py b/trinity/common/models/model.py index 3fd02341a5e..5e2f6e90317 100644 --- a/trinity/common/models/model.py +++ b/trinity/common/models/model.py @@ -151,15 +151,10 @@ async def extract_experience_from_history( enable recording must set ``self.recorder`` (a ``Recorder`` whose ``.store`` is a ``RecordStore``); this base implementation is shared. """ - recorder = getattr(self, "recorder", None) - if recorder is None: - raise ValueError("Recording is not enabled for this model.") - await recorder.flush() - exps = recorder.store.get(record_key) - if clear_history: - recorder.store.remove(record_key) - recorder.forget_record(record_key) - return exps + return await self._collect_experiences( + record_key, + remove=clear_history, + ) async def update_experience_reward( self, @@ -182,19 +177,38 @@ async def update_experience_reward( sample_ids=sample_ids, ) + async def overwrite_history_experiences(self, record_key: str, payload: bytes) -> None: + """Overwrite recorded experiences under one complete record key.""" + recorder = getattr(self, "recorder", None) + if recorder is None: + raise ValueError("Recording is not enabled for this model.") + await recorder.flush() + recorder.store.overwrite(record_key, Experience.deserialize_many(payload)) + recorder.forget_record(record_key) + async def _drain_experience_records(self, prefix: str) -> List[Experience]: """Remove and return recorded experiences matching a key or prefix.""" + return await self._collect_experiences( + prefix, + remove=True, + ) + + async def _collect_experiences( + self, + key: str, + *, + remove: bool, + ) -> List[Experience]: + """Collect recorded experiences by exact key or store-supported prefix.""" recorder = getattr(self, "recorder", None) if recorder is None: raise ValueError("Recording is not enabled for this model.") await recorder.flush() - matched_keys = [ - key for key in recorder.store.keys() if key == prefix or key.startswith(f"{prefix}/") - ] - exps = recorder.store.remove(prefix) - for key in matched_keys: + if remove: + exps = recorder.store.remove(key) recorder.forget_record(key) - return exps + return exps + return recorder.store.get(key) async def drain_experience_records_bytes(self, prefix: str) -> bytes: """Remove matching recorded experiences and return serialized bytes.""" @@ -845,6 +859,21 @@ async def update_experience_reward_async( sample_ids=sample_ids, ) + async def overwrite_history_experiences_async( + self, + record_key: str, + experiences: List[Experience], + ) -> None: + """Overwrite recorded experiences under one complete record key.""" + if not self.enable_history: + raise ValueError("History recording is not enabled.") + if self.model is None: + raise ValueError("Recording overwrite requires an inference model actor.") + await self.model.overwrite_history_experiences.remote( + record_key=record_key, + payload=Experience.serialize_many(experiences), + ) + async def drain_experience_records_bytes_async(self, prefix: str) -> bytes: """Remove matching recorded experiences and return serialized bytes.""" if not self.enable_history: From e255c18d49f533b5ca05f2e226ffe5123f2dbc61 Mon Sep 17 00:00:00 2001 From: "panxuchen.pxc" Date: Tue, 30 Jun 2026 12:36:17 +0800 Subject: [PATCH 39/84] refactor workflow interface --- tests/buffer/memory_store_test.py | 18 +- tests/common/recording_recorder_test.py | 6 +- tests/common/recording_store_test.py | 13 +- tests/explorer/workflow_test.py | 8 +- trinity/buffer/store/__init__.py | 3 +- trinity/buffer/store/base_store.py | 18 +- trinity/buffer/store/memory_store.py | 14 +- trinity/common/models/model.py | 147 ++++--- trinity/common/models/sglang_model.py | 18 +- trinity/common/models/vllm_model.py | 20 +- trinity/common/workflows/__init__.py | 11 +- .../workflows/on_policy_distill_workflow.py | 6 +- trinity/common/workflows/workflow.py | 365 ++++++++---------- trinity/explorer/scheduler.py | 64 +-- trinity/explorer/workflow_runner.py | 118 ++---- 15 files changed, 419 insertions(+), 410 deletions(-) diff --git a/tests/buffer/memory_store_test.py b/tests/buffer/memory_store_test.py index d86187b189a..e5fadef10e2 100644 --- a/tests/buffer/memory_store_test.py +++ b/tests/buffer/memory_store_test.py @@ -3,7 +3,7 @@ import torch -from trinity.buffer.store import MemoryStore +from trinity.buffer.store import ExperienceUpdate, MemoryStore from trinity.common.experience import EID, Experience @@ -32,7 +32,11 @@ def test_add_update_get_remove(self): store.add(key, experiences) self.assertEqual(len(store), 3) - store.update(key, reward=1.0, info={"source": "reward_model"}, sample_ids=None) + store.update( + key, + update=ExperienceUpdate(reward=1.0, info={"source": "reward_model"}), + sample_ids=None, + ) result = store.get(key) self.assertEqual(len(result), 3) for exp in result: @@ -53,7 +57,12 @@ def test_update_subset_by_sample_ids(self): experiences = get_dummy_experience(2, request_id="req_b") store.add(key, experiences) - store.update(key, reward=2.0, info=None, sample_ids=["req_b:1"]) + teacher_logprobs = torch.ones(3) + store.update( + key, + update=ExperienceUpdate(reward=2.0, teacher_logprobs=teacher_logprobs), + sample_ids=["req_b:1"], + ) result = store.get(key) self.assertIsNone(result[0].reward) @@ -61,6 +70,7 @@ def test_update_subset_by_sample_ids(self): self.assertEqual(result[1].eid.batch, "0") self.assertEqual(result[1].eid.task, "task_a") self.assertEqual(result[1].eid.run, 1) + torch.testing.assert_close(result[1].teacher_logprobs, teacher_logprobs) def test_overwrite_replaces_existing_records(self): store = MemoryStore() @@ -91,7 +101,7 @@ def test_complete_key_required_for_mutations(self): with self.assertRaises(ValueError): store.overwrite("0/task_a", get_dummy_experience(1)) with self.assertRaises(ValueError): - store.update("0/task_a", reward=1.0, info=None, sample_ids=None) + store.update("0/task_a", update=ExperienceUpdate(reward=1.0), sample_ids=None) with self.assertRaises(ValueError): store.add("0/task_a/not_int", get_dummy_experience(1)) diff --git a/tests/common/recording_recorder_test.py b/tests/common/recording_recorder_test.py index 0572f9bb0e7..bede2cb445d 100644 --- a/tests/common/recording_recorder_test.py +++ b/tests/common/recording_recorder_test.py @@ -2,7 +2,7 @@ import torch -from trinity.buffer.store import MemoryStore, parse_record_key +from trinity.buffer.store import ExperienceUpdate, MemoryStore, parse_record_key from trinity.common.experience import EID, Experience from trinity.common.models.recording.recorder import Recorder @@ -75,10 +75,10 @@ async def test_prefix_experiences_merge_and_keep_final_sample_id(self): self.assertEqual(merged.info["merged_eid_suffixes"], ["req-1", "req-2"]) self.assertEqual(merged.info["merged_sample_ids"], ["req-1", "req-2"]) - store.update(record_key, reward=1.0, info=None, sample_ids=["req-2"]) + store.update(record_key, update=ExperienceUpdate(reward=1.0), sample_ids=["req-2"]) self.assertEqual(store.get(record_key)[0].reward, 1.0) with self.assertRaises(KeyError): - store.update(record_key, reward=2.0, info=None, sample_ids=["req-1"]) + store.update(record_key, update=ExperienceUpdate(reward=2.0), sample_ids=["req-1"]) async def test_non_prefix_experiences_do_not_merge(self): store = MemoryStore() diff --git a/tests/common/recording_store_test.py b/tests/common/recording_store_test.py index facd41cf897..e27b2e9ab38 100644 --- a/tests/common/recording_store_test.py +++ b/tests/common/recording_store_test.py @@ -2,7 +2,12 @@ import torch -from trinity.buffer.store import MemoryStore, get_record_key, parse_record_key +from trinity.buffer.store import ( + ExperienceUpdate, + MemoryStore, + get_record_key, + parse_record_key, +) from trinity.common.experience import EID, Experience @@ -29,7 +34,11 @@ async def test_update_reward_sets_eid_from_record_key(self): exp = make_exp("req_a", record_key) store.add(get_record_key(exp), [exp]) - store.update(record_key, reward=1.5, info={"source": "reward_model"}, sample_ids=None) + store.update( + record_key, + update=ExperienceUpdate(reward=1.5, info={"source": "reward_model"}), + sample_ids=None, + ) updated = store.remove(record_key) self.assertEqual(len(updated), 1) diff --git a/tests/explorer/workflow_test.py b/tests/explorer/workflow_test.py index 520f5a3d620..33832accc6d 100644 --- a/tests/explorer/workflow_test.py +++ b/tests/explorer/workflow_test.py @@ -756,11 +756,11 @@ async def test_workflow_runner_partial_success_non_repeatable( workflow=PartialFailureWorkflow, repeat_times=3, raw_task={"fail_call_ids": [1]}, + batch_id="test", + task_id=0, ) - status, exps = await runner.run_task( - task, batch_id="test", repeat_times=3, run_id_base=0 - ) + status, exps = await runner.run_task(task, repeat_times=3, run_id_base=0) exps = deserialize_experiences(exps) self.assertFalse(status.ok) @@ -805,6 +805,7 @@ async def test_workflow_runner_fail_fast_without_partial_collection(self, concur workflow=PartialFailureWorkflow, repeat_times=3, raw_task={"fail_call_ids": []}, + batch_id="test", ) await runner.prepare() @@ -837,7 +838,6 @@ async def mock_execute_single_run( status, exps = await runner.run_task( task, - batch_id="test", repeat_times=3, run_id_base=0, collect_partial_runs=False, diff --git a/trinity/buffer/store/__init__.py b/trinity/buffer/store/__init__.py index 78d60297c1e..a50d652c1ff 100644 --- a/trinity/buffer/store/__init__.py +++ b/trinity/buffer/store/__init__.py @@ -1,4 +1,4 @@ -from trinity.buffer.store.base_store import RecordStore +from trinity.buffer.store.base_store import ExperienceUpdate, RecordStore from trinity.buffer.store.memory_store import ( MemoryStore, get_record_key, @@ -8,6 +8,7 @@ __all__ = [ "MemoryStore", + "ExperienceUpdate", "RecordStore", "get_record_key", "get_sample_id", diff --git a/trinity/buffer/store/base_store.py b/trinity/buffer/store/base_store.py index d8fa635ed60..ed116f903c6 100644 --- a/trinity/buffer/store/base_store.py +++ b/trinity/buffer/store/base_store.py @@ -1,9 +1,21 @@ from abc import ABC, abstractmethod +from dataclasses import dataclass from typing import List +from torch import Tensor + from trinity.common.experience import Experience +@dataclass +class ExperienceUpdate: + """Fields that may be patched onto recorded experiences after generation.""" + + reward: float | None = None + info: dict | None = None + teacher_logprobs: Tensor | None = None + + class RecordStore(ABC): """Abstract base class for an in-process experience store. @@ -24,10 +36,8 @@ def replace(self, key: str, old_sample_id: str, exp: Experience) -> None: """Replace one experience under the given complete key.""" @abstractmethod - def update( - self, key: str, reward: float, info: dict | None, sample_ids: List[str] | None - ) -> None: - """Update reward, EID fields from key, and optional info for selected experiences.""" + def update(self, key: str, update: ExperienceUpdate, sample_ids: List[str] | None) -> None: + """Patch selected experiences and stamp EID fields from the complete key.""" @abstractmethod def get(self, key: str) -> List[Experience]: diff --git a/trinity/buffer/store/memory_store.py b/trinity/buffer/store/memory_store.py index 2835c2ce1f9..dc390c898ad 100644 --- a/trinity/buffer/store/memory_store.py +++ b/trinity/buffer/store/memory_store.py @@ -3,7 +3,7 @@ from collections import OrderedDict from typing import Iterable, List -from trinity.buffer.store.base_store import RecordStore +from trinity.buffer.store.base_store import ExperienceUpdate, RecordStore from trinity.common.experience import Experience @@ -108,8 +108,7 @@ def replace(self, key: str, old_sample_id: str, exp: Experience) -> None: def update( self, key: str, - reward: float, - info: dict | None, + update: ExperienceUpdate, sample_ids: List[str] | None, ) -> None: batch, task, run = self._parse_complete_key(key) # validate key format @@ -124,11 +123,14 @@ def update( exp.eid.batch = batch exp.eid.task = task exp.eid.run = run - exp.reward = reward - if info: + if update.reward is not None: + exp.reward = update.reward + if update.info: if exp.info is None: exp.info = {} - exp.info.update(info) + exp.info.update(update.info) + if update.teacher_logprobs is not None: + exp.teacher_logprobs = update.teacher_logprobs def get(self, key: str) -> List[Experience]: result: List[Experience] = [] diff --git a/trinity/common/models/model.py b/trinity/common/models/model.py index 5e2f6e90317..20713dc7998 100644 --- a/trinity/common/models/model.py +++ b/trinity/common/models/model.py @@ -14,6 +14,7 @@ from torch import Tensor from transformers import AutoConfig +from trinity.buffer.store import ExperienceUpdate from trinity.common.config import InferenceModelConfig from trinity.common.constants import RunningStatus, SyncMethod from trinity.common.experience import Experience @@ -158,33 +159,45 @@ async def extract_experience_from_history( async def update_experience_reward( self, - record_key: str, + key: str, reward: float, info: Optional[dict] = None, sample_ids: Optional[List[str]] = None, ) -> None: """Update reward and optional info on recorded experiences.""" + await self.update_experience_records( + key=key, + update=ExperienceUpdate(reward=reward, info=info), + sample_ids=sample_ids, + ) + + async def update_experience_records( + self, + key: str, + update: ExperienceUpdate, + sample_ids: Optional[List[str]] = None, + ) -> None: + """Patch recorded experiences with generation-time training signals.""" recorder = getattr(self, "recorder", None) if recorder is None: raise ValueError("Recording is not enabled for this model.") await recorder.flush() - if not recorder.store.get(record_key): + if not recorder.store.get(key): return recorder.store.update( - key=record_key, - reward=reward, - info=info, + key=key, + update=update, sample_ids=sample_ids, ) - async def overwrite_history_experiences(self, record_key: str, payload: bytes) -> None: + async def overwrite_history_experiences(self, key: str, payload: bytes) -> None: """Overwrite recorded experiences under one complete record key.""" recorder = getattr(self, "recorder", None) if recorder is None: raise ValueError("Recording is not enabled for this model.") await recorder.flush() - recorder.store.overwrite(record_key, Experience.deserialize_many(payload)) - recorder.forget_record(record_key) + recorder.store.overwrite(key, Experience.deserialize_many(payload)) + recorder.forget_record(key) async def _drain_experience_records(self, prefix: str) -> List[Experience]: """Remove and return recorded experiences matching a key or prefix.""" @@ -490,7 +503,6 @@ def __init__( self.logger = get_logger(__name__) self.enable_lora = config.enable_lora self.enable_history = config.enable_history - self.recording_history_offsets: Dict[str, int] = {} self.status = RunningStatus.RUNNING self.workflow_state: Dict = {} self.request_count = 0 @@ -531,49 +543,47 @@ async def prepare(self) -> None: f"API server at {self.api_address} not ready after {max_retries} attempts." ) - def generate(self, prompts: List[str], **kwargs) -> List[Experience]: + def generate( + self, prompts: List[str], enable_recording: bool = False, **kwargs + ) -> List[Experience]: """Generate a list of experiences from a list of prompts.""" lora_request = self.get_lora_request() - if self.config.enable_history and kwargs.get("record_key") is None: - kwargs["record_key"] = self._api_key + if self.config.enable_history and enable_recording: + kwargs["key"] = self._api_key results = ray.get( [self.model.generate.remote(prompt, lora_request, **kwargs) for prompt in prompts] ) return [exp for exps in results for exp in exps] - async def generate_async(self, prompts: List[str], **kwargs) -> List[Experience]: + async def generate_async( + self, prompts: List[str], enable_recording: bool = False, **kwargs + ) -> List[Experience]: """Generate a list of experiences from a list of prompts in async.""" lora_request = await self.get_lora_request_async() - if self.config.enable_history and kwargs.get("record_key") is None: - kwargs["record_key"] = self._api_key + if self.config.enable_history and enable_recording: + kwargs["key"] = self._api_key results = await asyncio.gather( *[self.model.generate.remote(prompt, lora_request, **kwargs) for prompt in prompts] ) return [exp for exps in results for exp in exps] def chat( - self, messages: List[dict], record_key: Optional[str] = None, **kwargs + self, messages: List[dict], enable_recording: bool = False, **kwargs ) -> List[Experience]: """Generate a list of experiences from a list of messages.""" lora_request = self.get_lora_request() - if self.config.enable_history and record_key is None: - record_key = self._api_key - return ray.get( - self.model.chat.remote( - messages, lora_request=lora_request, record_key=record_key, **kwargs - ) - ) + if self.config.enable_history and enable_recording: + kwargs["key"] = self._api_key + return ray.get(self.model.chat.remote(messages, lora_request=lora_request, **kwargs)) async def chat_async( - self, messages: List[dict], record_key: Optional[str] = None, **kwargs + self, messages: List[dict], enable_recording: bool = False, **kwargs ) -> List[Experience]: """Generate a list of experiences from a list of messages in async.""" lora_request = await self.get_lora_request_async() - if self.config.enable_history and record_key is None: - record_key = self._api_key - return await self.model.chat.remote( - messages, lora_request=lora_request, record_key=record_key, **kwargs - ) + if self.config.enable_history and enable_recording: + kwargs["key"] = self._api_key + return await self.model.chat.remote(messages, lora_request=lora_request, **kwargs) def logprobs(self, tokens: List[int], temperature: Optional[float] = None) -> Tensor: """Calculate the logprobs of the given tokens.""" @@ -609,6 +619,13 @@ async def convert_messages_to_experience_async( messages, tools=tools, temperature=temperature ) + @property + def base_url(self) -> str: + """Get the base URL of the API server.""" + if not self.api_address: + raise ValueError("API address is not set. Cannot get base URL.") + return self.api_address + @property def api_key(self) -> str: """Get the API key.""" @@ -829,16 +846,10 @@ def extract_experience_from_history( exps = ray.get( self.model.extract_experience_from_history.remote( record_key=record_key, - clear_history=False, + clear_history=clear_history, ) ) - offset = self.recording_history_offsets.get(record_key, 0) - if offset > len(exps): - offset = 0 - new_exps = exps[offset:] - if clear_history: - self.recording_history_offsets[record_key] = len(exps) - return new_exps + return exps async def update_experience_reward_async( self, @@ -848,21 +859,64 @@ async def update_experience_reward_async( sample_ids: Optional[List[str]] = None, ) -> None: """Update reward and optional info on recorded experiences.""" + await self.update_experience_records_async( + record_key=record_key, + update=ExperienceUpdate(reward=reward, info=info), + sample_ids=sample_ids, + ) + + async def update_experience_records_async( + self, + record_key: str, + update: ExperienceUpdate, + sample_ids: Optional[List[str]] = None, + ) -> None: + """Patch recorded experiences with generation-time training signals.""" if not self.enable_history: raise ValueError("History recording is not enabled.") if self.model is None: raise ValueError("Recording update requires an inference model actor.") - await self.model.update_experience_reward.remote( + await self.model.update_experience_records.remote( record_key=record_key, - reward=reward, - info=info, + update=update, sample_ids=sample_ids, ) - async def overwrite_history_experiences_async( + def update_experience_reward( self, record_key: str, - experiences: List[Experience], + reward: float, + info: Optional[dict] = None, + sample_ids: Optional[List[str]] = None, + ) -> None: + """Update reward and optional info on recorded experiences.""" + self.update_experience_records( + record_key=record_key, + update=ExperienceUpdate(reward=reward, info=info), + sample_ids=sample_ids, + ) + + def update_experience_records( + self, + record_key: str, + update: ExperienceUpdate, + sample_ids: Optional[List[str]] = None, + ) -> None: + """Patch recorded experiences with generation-time training signals.""" + if not self.enable_history: + raise ValueError("History recording is not enabled.") + if self.model is None: + raise ValueError("Recording update requires an inference model actor.") + ray.get( + self.model.update_experience_records.remote( + record_key=record_key, + update=update, + sample_ids=sample_ids, + ) + ) + + async def overwrite_history_experiences_async( + self, experiences: List[Experience], key: str ) -> None: """Overwrite recorded experiences under one complete record key.""" if not self.enable_history: @@ -870,7 +924,7 @@ async def overwrite_history_experiences_async( if self.model is None: raise ValueError("Recording overwrite requires an inference model actor.") await self.model.overwrite_history_experiences.remote( - record_key=record_key, + key=key, payload=Experience.serialize_many(experiences), ) @@ -900,8 +954,6 @@ async def clean_workflow_state(self) -> None: """Clean the state of workflow using the model.""" async with self.state_lock: self.workflow_state = {} - # TODO: clear engine side experiences - self.recording_history_offsets.clear() async def shutdown(self) -> None: """Shutdown all underlying model actors cleanly.""" @@ -917,10 +969,9 @@ async def get_workflow_state(self) -> Dict: async with self.state_lock: return self.workflow_state.copy() - def clone_with_isolated_history(self) -> "ModelWrapper": - """Clone the current ModelWrapper with isolated history.""" + def clone_with_isolated_state(self) -> "ModelWrapper": + """Clone the current ModelWrapper with isolated state.""" new_wrapper = copy.copy(self) new_wrapper.openai_async_client = None new_wrapper.openai_client = None - new_wrapper.recording_history_offsets = {} return new_wrapper diff --git a/trinity/common/models/sglang_model.py b/trinity/common/models/sglang_model.py index 463a43e5b1f..4ed87166f31 100644 --- a/trinity/common/models/sglang_model.py +++ b/trinity/common/models/sglang_model.py @@ -222,7 +222,7 @@ async def update_weights_from_disk( return success async def generate( - self, input_ids: List[int], record_key: Optional[str] = None, **kwargs + self, input_ids: List[int], key: Optional[str] = None, **kwargs ) -> Sequence[dict[str, Any]]: sampling_params = { "n": kwargs.get("n", 1), @@ -251,7 +251,7 @@ async def generate( "/generate", payload, timeout=kwargs.get("timeout", 300), - api_key_override=record_key, + api_key_override=key, ) if isinstance(response, dict) and response.get("error"): raise RuntimeError(f"Failed to generate with SGLang: {response['error']}") @@ -404,12 +404,12 @@ async def generate( self, prompt: str, lora_request=None, - record_key: Optional[str] = None, + key: Optional[str] = None, **kwargs, ) -> Sequence[Experience]: """Generate a response from the provided prompt in async. - When ``record_key`` is set, it is sent as the Authorization bearer so the + When ``key`` is set, it is sent as the Authorization bearer so the server-side recorder groups this turn under that key (the api_key doubles as the record_key on the Trinity path). The returned experiences are the client-side copy; the recorded copy is written to the in-process store by @@ -429,7 +429,7 @@ async def generate( return_logprob = logprobs is not None and logprobs is not False responses = await self.api_client.generate( input_ids=prompt_token_ids, - record_key=record_key, + key=key, n=kwargs.get("n", 1), temperature=kwargs.get("temperature", self.config.temperature), top_p=kwargs.get("top_p", self.config.top_p), @@ -489,10 +489,10 @@ async def chat( self, messages: List[dict], lora_request=None, - record_key: Optional[str] = None, + key: Optional[str] = None, **kwargs, ) -> Sequence[Experience]: - # ``record_key`` is propagated to ``generate`` so the server-side recorder + # ``key`` is propagated to ``generate`` so the server-side recorder # groups this turn under the caller's key (sent as the Authorization # bearer, same as vLLM's RecordingIdentityMiddleware path). if self.tokenizer is None: @@ -500,9 +500,7 @@ async def chat( normalized_messages = self._normalize_chat_messages(messages) prompt = self.apply_chat_template(self.tokenizer, normalized_messages) - return await self.generate( - prompt=prompt, lora_request=lora_request, record_key=record_key, **kwargs - ) + return await self.generate(prompt=prompt, lora_request=lora_request, key=key, **kwargs) async def logprobs(self, token_ids: List[int], **kwargs) -> torch.Tensor: raise NotImplementedError("SGLangRolloutModel does not support logprobs.") diff --git a/trinity/common/models/vllm_model.py b/trinity/common/models/vllm_model.py index 3386778694c..896db6b80f6 100644 --- a/trinity/common/models/vllm_model.py +++ b/trinity/common/models/vllm_model.py @@ -220,14 +220,14 @@ async def chat( self, messages: List[Dict], lora_request=None, - record_key: Optional[str] = None, + key: Optional[str] = None, **kwargs, ) -> Sequence[Experience]: """Chat with the model with a list of messages in async. Args: messages (List[dict]): The input history messages. - record_key (Optional[str]): Recording identity for the in-vLLM + key (Optional[str]): Recording identity for the in-vLLM recorder (the MemoryStore group key). Propagated to ``generate`` via ``record_key_ctx`` so the recorder stamps it into ``Experience.eid`` without an HTTP hop. None skips @@ -257,22 +257,20 @@ async def chat( "prompt": prompt, "multi_modal_data": multi_modal_data or {}, } - return await self.generate( - prompt=prompt, lora_request=lora_request, record_key=record_key, **kwargs - ) + return await self.generate(prompt=prompt, lora_request=lora_request, key=key, **kwargs) async def generate( self, prompt: Union[str, Dict], lora_request=None, - record_key: Optional[str] = None, + key: Optional[str] = None, **kwargs, ) -> Sequence[Experience]: """Generate a response from the provided prompt in async. Args: prompt (str): The input prompt. - record_key (Optional[str]): Recording identity propagated to the + key (Optional[str]): Recording identity propagated to the in-vLLM recorder via ``record_key_ctx`` (see ``chat``). kwargs (dict): A dictionary of sampling parameters. @@ -292,15 +290,15 @@ async def generate( # persist these dummies directly under the record_key — masked # experiences must still be tracked for history extraction and # the buffer/trainer (they are popped by record_key on consume). - if self.recorder is not None and record_key is not None: - batch, task, run = parse_record_key(record_key) + if self.recorder is not None and key is not None: + batch, task, run = parse_record_key(key) for exp in returned_seq: exp.eid.batch = batch exp.eid.task = task exp.eid.run = run exp.info["rank"] = self.recorder.rank exp.info["model_version"] = self.model_version - self.recorder.store.add(record_key, [exp]) + self.recorder.store.add(key, [exp]) return returned_seq prompt = { "prompt_token_ids": returned_seq @@ -310,7 +308,7 @@ async def generate( # Propagate the recording identity to the engine-level recorder (same # async task, same process) so the recorded experience is grouped under # this record key in the MemoryStore. - record_key_token = record_key_ctx.set(record_key) + record_key_token = record_key_ctx.set(key) try: output = await self._generate_internal( prompt=prompt, lora_request=lora_request, **kwargs diff --git a/trinity/common/workflows/__init__.py b/trinity/common/workflows/__init__.py index 11f099d1e1e..34bd27be59c 100644 --- a/trinity/common/workflows/__init__.py +++ b/trinity/common/workflows/__init__.py @@ -1,13 +1,6 @@ # -*- coding: utf-8 -*- """Workflow module""" -from trinity.common.workflows.workflow import ( - Metrics, - RepeatableMultiTurnWorkflow, - RepeatableWorkflow, - Status, - Task, - Workflow, -) +from trinity.common.workflows.workflow import Metrics, Status, Task, Workflow from trinity.utils.registry import Registry WORKFLOWS: Registry = Registry( @@ -66,8 +59,6 @@ "Task", "Status", "Metrics", - "RepeatableMultiTurnWorkflow", - "RepeatableWorkflow", "Workflow", "WORKFLOWS", ] diff --git a/trinity/common/workflows/on_policy_distill_workflow.py b/trinity/common/workflows/on_policy_distill_workflow.py index 3fb67f5e628..e9251033741 100644 --- a/trinity/common/workflows/on_policy_distill_workflow.py +++ b/trinity/common/workflows/on_policy_distill_workflow.py @@ -17,10 +17,10 @@ from trinity.common.experience import Experience from trinity.common.models.model import ModelWrapper from trinity.common.rewards.qwen25_eval import verify_math_answer -from trinity.common.workflows.workflow import RepeatableWorkflow, Task +from trinity.common.workflows.workflow import Task, Workflow -class OnPolicyDistillWorkflow(RepeatableWorkflow): +class OnPolicyDistillWorkflow(Workflow): """On-policy distillation workflow. Computes and stores teacher_logprobs in experience.info. @@ -34,6 +34,7 @@ class OnPolicyDistillWorkflow(RepeatableWorkflow): is_async: bool = True can_reset: bool = True + can_repeat: bool = True def __init__( self, @@ -116,7 +117,6 @@ async def run_async(self) -> List[Experience]: resp_start = response.prompt_length - 1 teacher_resp_logprobs = teacher_logprobs[resp_start:] student_resp_logprobs = response.logprobs - assert student_resp_logprobs is not None, "Student logprobs should not be None." # Verify lengths match (they should be equal for the same token sequence) assert len(teacher_resp_logprobs) == len(student_resp_logprobs), ( diff --git a/trinity/common/workflows/workflow.py b/trinity/common/workflows/workflow.py index b73b87b3d6b..19727791252 100644 --- a/trinity/common/workflows/workflow.py +++ b/trinity/common/workflows/workflow.py @@ -3,10 +3,10 @@ from __future__ import annotations +from abc import abstractmethod from dataclasses import asdict, dataclass, field from typing import TYPE_CHECKING, Dict, List, Optional, Type, Union -from trinity.buffer.store import parse_record_key from trinity.common.config import FormatConfig, GenerationConfig from trinity.common.experience import Experience from trinity.common.rewards.reward_fn import RewardFn @@ -25,7 +25,7 @@ class Status: completed_runs: int total_runs: int metrics: List[Dict[str, float]] - successful_run_ids: List[int] = field(default_factory=list) + successful_run_ids: List[str] = field(default_factory=list) message: Optional[str] = None @property @@ -53,6 +53,7 @@ class Task(dict): # automatically assigned ids batch_id: Union[int, str] = "" task_id: Union[int, str] = "" + run_id: Union[int, str] = "" index: dict = field(default_factory=dict) @@ -90,23 +91,41 @@ def truth(self) -> Union[str, None]: response_key = self.format_args.response_key return self.raw_task[response_key] if response_key in self.raw_task else None # type: ignore + @property + def api_key(self) -> str: + if self.batch_id is None or self.task_id is None or self.run_id is None: + raise ValueError("batch_id, task_id, and run_id must be set before generating API_KEY.") + return f"{self.batch_id}/{self.task_id}/{self.run_id}" + def to_dict(self) -> dict: return self.raw_task # type: ignore -class Workflow: - """The base workflow class. +class BaseWorkflow: + """The base workflow interface.""" + + def __init__(self, task: Task, model: ModelWrapper) -> None: + self.task = task + self.model = model + self.model.set_api_key(task.api_key) # set the API key for the rollout model + self.logger = get_logger(__name__) + + @abstractmethod + async def execute(self) -> Status: + """Execute the workflow and return a Status object.""" + + def reset(self, task: Task): + """Reset the workflow with a new task.""" + self.task.batch_id = task.batch_id + self.task.task_id = task.task_id + self.task.run_id = task.run_id + self.model.set_api_key(task.api_key) # set the API key for the rollout model - A workflow is a runnable object that implements rollout logic in ``run`` or - ``run_async``. User code normally returns metrics for the completed run - and lets :meth:`execute` assemble the internal :class:`Status` consumed by - the runner and scheduler. - Training experiences are captured by the rollout model's built-in recording - path during generation; workflows should update rewards on those recorded - experiences before returning. Advanced workflows may override - :meth:`execute` directly when they need full control over partial success, - run ids, or error messages. +class Workflow(BaseWorkflow): + """The base workflow class. + + A workflow is a runnable object which generates a list of experiences. Attributes: auxiliary_model_wrappers: List of ModelWrapper instances for auxiliary models. @@ -114,6 +133,7 @@ class Workflow: """ can_reset: bool = False # whether the workflow can be reset with a new task. If true, `reset()` must be implemented. + can_repeat: bool = False # whether the workflow can be repeated multiple times. If true, `set_repeat_times()` must be implemented. is_async: bool = False # whether the workflow runs in async mode. If true, `run_async()` must be implemented, else `run()` must be implemented. def __init__( @@ -123,8 +143,7 @@ def __init__( model: ModelWrapper, auxiliary_models: Optional[List[ModelWrapper]] = None, ): - self.task = task - self.model = model + super().__init__(task=task, model=model) # Store ModelWrapper instances self.auxiliary_model_wrappers = auxiliary_models # Get OpenAI clients from ModelWrapper (async or sync based on workflow type) @@ -134,13 +153,8 @@ def __init__( self.auxiliary_models = [m.get_openai_async_client() for m in auxiliary_models] else: self.auxiliary_models = [m.get_openai_client() for m in auxiliary_models] - self.run_id = 0 - self.logger = get_logger(__name__) - - @property - def resettable(self): - """Deprecated, use cls.can_reset instead.""" - return self.__class__.can_reset + self.run_id_base = 0 + self.repeat_times = 1 @property def asynchronous(self): @@ -148,124 +162,45 @@ def asynchronous(self): Whether the workflow runs in async mode.""" return self.__class__.is_async - def set_execution_context(self, run_id: int) -> None: - """Set the execution context used by ``execute`` and recording helpers.""" - self.run_id = run_id - - def reset(self, task: Task): - """Reset the workflow.""" - raise NotImplementedError - - def run(self) -> Metrics: - """Run workflow and return metrics for the completed run.""" - raise NotImplementedError - - async def run_async(self) -> Metrics: - """Run workflow asynchronously and return metrics for the completed run.""" - raise NotImplementedError - - async def execute(self) -> Status: - """Execute the workflow and normalize the user return value to Status.""" - result = await self.run_async() if self.asynchronous else self.run() - return self._to_status(result) - - def _to_status(self, result: Metrics) -> Status: - return Status( - completed_runs=1, - total_runs=1, - metrics=[result], - successful_run_ids=[self.run_id], - ) - - def _build_record_key(self, run_id: Optional[int] = None) -> str: - run = self.run_id if run_id is None else run_id - return f"{self.task.batch_id}/{self.task.task_id}/{run}" - - async def update_reward( - self, - reward: float, - info: Optional[dict] = None, - sample_ids: Optional[List[str]] = None, - run_id: Optional[int] = None, - ) -> None: - """Update recorded experiences for one run with reward and optional info.""" - await self.model.update_experience_reward_async( - record_key=self._build_record_key(run_id), - reward=reward, - info=info, - sample_ids=sample_ids, - ) - - -class RepeatableWorkflow(Workflow): - """Workflow base class for implementations that run repeats internally. - - ``run`` or ``run_async`` should return one metrics dict per successful - repeat. Workflows that need custom partial-success semantics can override - ``execute`` and return :class:`Status` directly. - """ - - def __init__( - self, - *, - task: Task, - model: ModelWrapper, - auxiliary_models: Optional[List[ModelWrapper]] = None, - ): - super().__init__(task=task, model=model, auxiliary_models=auxiliary_models) - self.repeat_times = 1 - self.run_id_base = 0 - def set_repeat_times(self, repeat_times: int, run_id_base: int) -> None: - """Set repeat count and base run id for one repeatable execution.""" + """ + Set the number of times to repeat the workflow. + Args: + repeat_times (int): number of times to repeat the workflow (if repeatable). + run_id_base (int): base run_id for setting run_id in experiences. + """ self.repeat_times = repeat_times self.run_id_base = run_id_base - self.run_id = run_id_base - def set_execution_context(self, repeat_times: int, run_id_base: int) -> None: - """Set the execution context used by ``execute`` and recording helpers.""" - self.set_repeat_times(repeat_times, run_id_base) - - def run(self) -> List[Metrics]: - """Run workflow and return metrics for completed repeats.""" + def run(self) -> List[Experience]: + """Run workflow and return a list of experiences.""" raise NotImplementedError - async def run_async(self) -> List[Metrics]: - """Run workflow asynchronously and return metrics for completed repeats.""" + async def run_async(self) -> List[Experience]: + """Run workflow in async and return a list of experiences.""" raise NotImplementedError - def _to_status(self, result: List[Metrics]) -> Status: - completed_runs = len(result) + async def execute(self) -> Status: + if self.asynchronous: + exps = await self.run_async() + else: + exps = self.run() + await self.model.overwrite_history_experiences_async( + experiences=exps, key=self.task.api_key + ) return Status( - completed_runs=completed_runs, - total_runs=self.repeat_times, - metrics=result, - successful_run_ids=list(range(self.run_id_base, self.run_id_base + completed_runs)), + completed_runs=self.__class__.can_repeat and self.repeat_times or 1, + total_runs=self.__class__.can_repeat and self.repeat_times or 1, + metrics=[exp.metrics for exp in exps if exp.metrics is not None], + successful_run_ids=[self.task.api_key], ) - def _build_record_key(self, run_id: Optional[int] = None) -> str: - run = self.run_id_base if run_id is None else run_id - return f"{self.task.batch_id}/{self.task.task_id}/{run}" - class MultiTurnWorkflow(Workflow): """ The base workflow class for concatenated multi-turn tasks. """ - def __init__( - self, - *, - task: Task, - model: ModelWrapper, - auxiliary_models: Optional[List[ModelWrapper]] = None, - ): - super().__init__( - task=task, - model=model, - auxiliary_models=auxiliary_models, - ) - def _build_experience_from_converted( self, converted_experience, reward, info={}, truncate_status=None ) -> Experience: @@ -311,15 +246,6 @@ def _build_experience_from_converted( def process_messages_to_experience( self, messages, reward, info={}, truncate_status=None ) -> Experience: - # TODO(recording): when enable_history is on, this client-side - # conversion is redundant — the vLLM recorder's build_experience already - # captured the authoritative heavy data (real logprobs without an extra - # forward, real routed_experts) into the MemoryStore, keyed by the - # record_key the chat call carried. Replace this with an in-process - # lookup by record_key (store.get / update), then concatenate the - # session's turns (info["sample_index"] orders them) into one - # experience here. Requires threading the per-call record_key down to - # this call site. converted_experience = self.model.convert_messages_to_experience(messages) return self._build_experience_from_converted( converted_experience, @@ -331,8 +257,6 @@ def process_messages_to_experience( async def process_messages_to_experience_async( self, messages, reward, info={}, truncate_status=None ) -> Experience: - # TODO(recording): see process_messages_to_experience — replace with a - # MemoryStore lookup by record_key once it is threaded here. converted_experience = await self.model.convert_messages_to_experience_async(messages) return self._build_experience_from_converted( converted_experience, @@ -342,11 +266,10 @@ async def process_messages_to_experience_async( ) -class RepeatableMultiTurnWorkflow(MultiTurnWorkflow, RepeatableWorkflow): - """Multi-turn workflow variant that runs repeats inside one execute call.""" +class BaseSimpleWorkflow(Workflow): + """A simple workflow for single-round tasks, which use the batch generation + API to generate multiple responses in one call.""" - -class BaseSimpleWorkflow(RepeatableWorkflow): def __init__( self, *, @@ -379,8 +302,8 @@ def reset(self, task: Task): def set_repeat_times(self, repeat_times, run_id_base): self.repeat_times = repeat_times - self.task.rollout_args.n = repeat_times self.run_id_base = run_id_base + self.task.rollout_args.n = repeat_times @property def rollout_args(self): @@ -396,14 +319,21 @@ def format_messages(self): messages.append({"role": "assistant", "content": self.reply_prefix}) return messages - # -- recording-path helpers (shared by SimpleWorkflow / AsyncSimpleWorkflow) - - def _record_key(self, run_index: int) -> str: - """Per-sample recording identity (the MemoryStore group key).""" - return f"{self.task.batch_id}/{self.task.task_id}/{self.run_id_base + run_index}" - def _attach_rewards(self, responses, *, base: int) -> List[Experience]: +class SimpleWorkflow(BaseSimpleWorkflow): + """A workflow for simple single-round task.""" + + can_reset: bool = True + can_repeat: bool = True + is_async: bool = False + + def run(self) -> List[Experience]: + # TODO: Optimize the generate function + messages = self.format_messages() + + self.logger.debug("start chat") + responses = self.model.chat(messages, **self.rollout_args) for i, response in enumerate(responses): - run = base + i reward_dict = self.reward_fn( # type: ignore [misc] response=response.response_text, # type: ignore [arg-type] truth=self.truth, @@ -413,50 +343,17 @@ def _attach_rewards(self, responses, *, base: int) -> List[Experience]: response.metrics.update(reward_dict) reward = sum(reward_dict.values()) response.reward = reward - response.eid.run = run + response.eid.run = i + self.run_id_base + self.logger.debug( - f"self.task_desc: {self.task_desc}, response: {response.response_text}, reward: {reward}" + f"self.task_desc: {self.task_desc}, messages: {messages}, response: {response.response_text}, reward: {reward}" ) return responses - @staticmethod - def _stamp_record_key(exps: List[Experience], record_key: str) -> None: - batch, task, run = parse_record_key(record_key) - for exp in exps: - exp.eid.batch = batch - exp.eid.task = task - exp.eid.run = run - - -class SimpleWorkflow(BaseSimpleWorkflow): - """A workflow for simple single-round task.""" - - can_reset: bool = True - - def run(self) -> List[Experience]: - # TODO: Optimize the generate function - messages = self.format_messages() - - self.logger.debug("start chat") - return self._run_recorded(messages) - - def _run_recorded(self, messages) -> List[Experience]: - # One chat call per sample (n=1) so each gets a distinct record_key - # (the recording group key == reward unit). The runner later reports - # {record_key: reward} and the coordinator joins reward in-store. - rollout_args = dict(self.rollout_args) - rollout_args["n"] = 1 - exps: List[Experience] = [] - for i in range(self.repeat_times): - record_key = self._record_key(i) - responses = self.model.chat(messages, record_key=record_key, **rollout_args) - rewarded = self._attach_rewards(responses, base=self.run_id_base + i) - self._stamp_record_key(rewarded, record_key) - exps.extend(rewarded) - return exps - class AsyncSimpleWorkflow(BaseSimpleWorkflow): + can_reset: bool = True + can_repeat: bool = True is_async: bool = True async def run_async(self) -> List[Experience]: @@ -464,19 +361,24 @@ async def run_async(self) -> List[Experience]: messages = self.format_messages() self.logger.info("start chat") - return await self._run_recorded_async(messages) + responses = await self.model.chat_async(messages, **self.rollout_args) + for i, response in enumerate(responses): + reward_dict = self.reward_fn( # type: ignore [misc] + response=response.response_text, # type: ignore [arg-type] + truth=self.truth, + ) - async def _run_recorded_async(self, messages) -> List[Experience]: - rollout_args = dict(self.rollout_args) - rollout_args["n"] = 1 - exps: List[Experience] = [] - for i in range(self.repeat_times): - record_key = self._record_key(i) - responses = await self.model.chat_async(messages, record_key=record_key, **rollout_args) - rewarded = self._attach_rewards(responses, base=self.run_id_base + i) - self._stamp_record_key(rewarded, record_key) - exps.extend(rewarded) - return exps + if response.metrics is None: + response.metrics = {} + response.metrics.update(reward_dict) + reward = sum(reward_dict.values()) + response.reward = reward + response.eid.run = i + self.run_id_base + + self.logger.debug( + f"self.task_desc: {self.task_desc}, messages: {messages}, response: {response.response_text}, reward: {reward}" + ) + return responses class MathWorkflow(SimpleWorkflow): @@ -512,3 +414,76 @@ def reset(self, task: Task): class AsyncMathWorkflow(AsyncSimpleWorkflow, MathWorkflow): pass + + +class WorkflowWithRecording(BaseWorkflow): + """A workflow that using the rollout model's built-in recording path to capture + experience data. + + This interface is designed for complex agentic workflows (e.g., QwenPaw, Claude Code) + which are hard to extract experience data from the agent itself. + + It provides `base_url` and `api_key` to the OpenAI API of the rollout model, and the + workflow can use them to call the model and the model will record the experience data + automatically. + After the agentic workflow is completed, the workflow can call `update_reward` to update + the recorded experience data with the reward and optional info. + """ + + def __init__( + self, + *, + task: Task, + model: ModelWrapper, + auxiliary_models: Optional[List[ModelWrapper]] = None, + ): + super().__init__(task=task, model=model) + # Store ModelWrapper instances + self.auxiliary_model_wrappers = auxiliary_models + # Get OpenAI clients from ModelWrapper + self.auxiliary_models = [m.get_openai_async_client() for m in auxiliary_models] + + @property + def base_url(self) -> str: + """BASE_URL of the OpenAI API of the rollout model.""" + return self.model.base_url + + @property + def api_key(self) -> str: + """API_KEY of the OpenAI API of the rollout model.""" + return self.task.api_key + + def reset(self, task: Task): + """Reset the workflow.""" + raise NotImplementedError + + async def run(self) -> Metrics: + """Run workflow asynchronously and return metrics for the completed run.""" + raise NotImplementedError + + async def execute(self) -> Status: + """Execute the workflow and normalize the user return value to Status.""" + result = await self.run() + return self._to_status(result) + + def _to_status(self, result: Metrics) -> Status: + return Status( + completed_runs=1, + total_runs=1, + metrics=[result], + successful_run_ids=[self.task.api_key], + ) + + async def update_reward( + self, + reward: float, + info: Optional[dict] = None, + sample_ids: Optional[List[str]] = None, + ) -> None: + """Update recorded experiences for one run with reward and optional info.""" + await self.model.update_experience_reward_async( + record_key=self.api_key, + reward=reward, + info=info, + sample_ids=sample_ids, + ) diff --git a/trinity/explorer/scheduler.py b/trinity/explorer/scheduler.py index 23a74f47d81..e19c32d509f 100644 --- a/trinity/explorer/scheduler.py +++ b/trinity/explorer/scheduler.py @@ -12,7 +12,7 @@ from ray.actor import ActorHandle from trinity.common.config import Config -from trinity.common.workflows import Task +from trinity.common.workflows import RepeatableWorkflow, Task from trinity.explorer.workflow_runner import Status, WorkflowRunner from trinity.utils.log import get_logger from trinity.utils.metrics import calculate_task_level_metrics @@ -25,18 +25,25 @@ class TaskWrapper: """ task: Task - batch_id: Union[int, str] sub_task_num: int = 1 # number of sub tasks splitted from this task # if max_repeat_times_per_runner is set, one task may be splitted into multiple sub tasks finished_sub_task_num: int = 0 completed_runs: int = 0 total_runs: int = 0 # total planned runs for the whole task metrics: List[Dict[str, float]] = field(default_factory=list) - successful_run_ids: List[int] = field(default_factory=list) + successful_run_ids: List[str] = field(default_factory=list) experience_payloads: List[bytes] = field(default_factory=list) first_error: Optional[str] = None emitted: bool = False + @property + def batch_id(self) -> Union[int, str]: + return self.task.batch_id + + @property + def task_id(self) -> Union[int, str]: + return self.task.task_id + @dataclass(frozen=True) class CompletedTaskResult: @@ -107,11 +114,11 @@ async def update_state(self) -> None: async def prepare(self): await self.runner.prepare.remote() - def _enable_history_recording(self) -> bool: - return bool(self.config.explorer.rollout_model.enable_history) + def _task_level_record_key(self, task: TaskWrapper) -> str: + return f"{task.batch_id}/{task.task_id}" - def _record_key(self, task: TaskWrapper, run_id: int) -> str: - return f"{task.batch_id}/{task.task.task_id}/{run_id}" + def _run_level_record_key(self, task: TaskWrapper, run_id: int) -> str: + return f"{task.batch_id}/{task.task_id}/{run_id}" async def _drain_records(self, prefix: str) -> bytes: try: @@ -127,25 +134,26 @@ async def _delete_records(self, prefix: str) -> None: self.logger.error("records delete from rollout actor failed: %s", exc) async def _consume_finished_records(self, task: TaskWrapper, status: Status) -> List[bytes]: - if not self._enable_history_recording(): - return [] if not status.successful_run_ids: return [] + if issubclass(task.task.workflow, RepeatableWorkflow): + prefix = self._task_level_record_key(task) + if task.task.is_eval: + await self._delete_records(prefix) + return [] + + payload = await self._drain_records(prefix) + return [payload] if payload else [] + if task.task.is_eval: await asyncio.gather( - *[ - self._delete_records(self._record_key(task, run_id)) - for run_id in status.successful_run_ids - ] + *[self._delete_records(run_id) for run_id in status.successful_run_ids] ) return [] payloads = await asyncio.gather( - *[ - self._drain_records(self._record_key(task, run_id)) - for run_id in status.successful_run_ids - ] + *[self._drain_records(run_id) for run_id in status.successful_run_ids] ) return [payload for payload in payloads if payload] @@ -153,7 +161,7 @@ async def run_with_retry( self, task: TaskWrapper, repeat_times: int, - run_id_base: int, + run_base: int, timeout: float, collect_partial_runs: bool, ) -> Tuple[Status, List[bytes], int, float]: @@ -187,9 +195,8 @@ async def run_with_retry( try: run_task_ref = self.runner.run_task.remote( task=task2run, - batch_id=str(task.batch_id), repeat_times=repeat_times, - run_id_base=run_id_base, + run_base=run_base, collect_partial_runs=collect_partial_runs, ) status = await asyncio.wait_for( @@ -403,13 +410,13 @@ async def _schedule_pending_tasks(self) -> None: task_queue = self.pending_tasks[batch_id] while task_queue and self.idle_runners: - task, repeat_times, run_id_base = task_queue.pop() + task, repeat_times, run_base = task_queue.pop() runner_id = self.idle_runners.pop() future = asyncio.create_task( self.runners[runner_id].run_with_retry( task, repeat_times=repeat_times, - run_id_base=run_id_base, + run_base=run_base, timeout=self.dynamic_timeout(), collect_partial_runs=self.config.explorer.over_rollout.return_partial_tasks, ) @@ -432,7 +439,7 @@ def task_done_callback(self, async_task: asyncio.Task): self.busy_runners.pop(runner_id, None) self.idle_runners.add(runner_id) elif async_task.exception(): - self.logger.error(f"Task {task.task.task_id} failed: {async_task.exception()}") + self.logger.error(f"Task {task.task_id} failed: {async_task.exception()}") self._schedule_runner_restart(runner_id) else: status, experience_payloads, runner_id, run_time = async_task.result() @@ -487,7 +494,7 @@ def _emit_task_result(self, task: TaskWrapper) -> None: if task.emitted: return status = self._build_task_result(task) - task_id = task.task.task_id + task_id = task.task_id completed_result = CompletedTaskResult( batch_id=task.batch_id, task_id=task_id, @@ -520,7 +527,7 @@ def _emit_partial_tasks_for_batch(self, batch_id: Union[int, str]) -> None: continue self._emit_task_result(task) self.logger.debug( - f"Task partially completed and emitted (batch_id {task.batch_id}, task_id {task.task.task_id})." + f"Task partially completed and emitted (batch_id {task.batch_id}, task_id {task.task_id})." ) def _clear_timeout_tasks(self, batch_id: Union[int, str]) -> List[asyncio.Future]: @@ -624,7 +631,6 @@ def _split_and_submit_tasks(self, tasks: List[Task], batch_id: Union[int, str]) assert task.repeat_times is not None, "Task repeat_times should not be None" task_wrapper = TaskWrapper( task=replace(task, batch_id=batch_id, task_id=i), - batch_id=batch_id, total_runs=task.repeat_times, ) if self.max_repeat_times is None: @@ -632,9 +638,9 @@ def _split_and_submit_tasks(self, tasks: List[Task], batch_id: Union[int, str]) self.pending_tasks[batch_id].appendleft((task_wrapper, task.repeat_times, 0)) continue sub_tasks = [] - for run_id_base in range(0, task.repeat_times, self.max_repeat_times): - repeat_times = min(self.max_repeat_times, task.repeat_times - run_id_base) - sub_tasks.append((task_wrapper, repeat_times, run_id_base)) + for run_base in range(0, task.repeat_times, self.max_repeat_times): + repeat_times = min(self.max_repeat_times, task.repeat_times - run_base) + sub_tasks.append((task_wrapper, repeat_times, run_base)) task_wrapper.sub_task_num = len(sub_tasks) self.pending_tasks[batch_id].extendleft(sub_tasks) diff --git a/trinity/explorer/workflow_runner.py b/trinity/explorer/workflow_runner.py index f9c0f764695..b54395068ed 100644 --- a/trinity/explorer/workflow_runner.py +++ b/trinity/explorer/workflow_runner.py @@ -13,7 +13,7 @@ from trinity.common.experience import Experience from trinity.common.models.allocator import Allocator from trinity.common.models.model import ModelWrapper -from trinity.common.workflows import RepeatableWorkflow, Status, Task, Workflow +from trinity.common.workflows import Status, Task, Workflow from trinity.utils.log import get_logger @@ -83,17 +83,9 @@ async def prepare(self) -> None: def is_alive(self): return True - def _build_record_key(self, task: Task, run_index: int) -> str: - return f"{task.batch_id}/{task.task_id}/{run_index}" - - def _set_record_key(self, model_wrapper: ModelWrapper, record_key: Optional[str]) -> None: - if record_key is not None: - model_wrapper.set_api_key(record_key) - - def _create_workflow_instance(self, task: Task, record_key: Optional[str] = None) -> Workflow: + def _create_workflow_instance(self, task: Task) -> Workflow: if task.workflow is None: raise ValueError("Workflow is not set in the task.") - self._set_record_key(self.model_wrapper, record_key) if ( self.workflow_instance is None or not self.workflow_instance.__class__ == task.workflow @@ -118,14 +110,12 @@ async def _run_workflow(self, workflow_instance: Workflow) -> Status: return status def _create_isolated_workflow_instance( - self, task: Task, record_key: Optional[str] = None + self, task: Task, run_id: int ) -> Tuple[Workflow, ModelWrapper]: - model_wrapper = ( - self.model_wrapper.clone_with_isolated_history() - if self.config.explorer.rollout_model.enable_history - else self.model_wrapper - ) - self._set_record_key(model_wrapper, record_key) + model_wrapper = self.model_wrapper.clone_with_isolated_state() + # only a shallow copy is enough + task = task.copy() + task.run_id = run_id wf = task.to_workflow( model_wrapper, self.auxiliary_model_wrappers, @@ -136,7 +126,7 @@ def _build_status( self, total_runs: int, metrics: List[Dict[str, float]], - successful_run_ids: List[int], + successful_run_ids: List[str], first_error: Optional[str] = None, ) -> Status: completed_runs = len(successful_run_ids) @@ -191,16 +181,11 @@ async def _run_parallel_runs( use_threads: bool = False, ) -> Status: async def run_single(i: int) -> Status: - run_index = run_id_base + i - record_key = self._build_record_key(task, run_index) - workflow, model_wrapper = self._create_isolated_workflow_instance(task, record_key) + run_id = run_id_base + i + workflow, model_wrapper = self._create_isolated_workflow_instance(task, run_id) return await self._execute_single_run( - workflow, - task, - i, - run_id_base, + workflow=workflow, model_wrapper=model_wrapper, - record_key=record_key, ) if collect_partial_runs: @@ -250,40 +235,19 @@ async def run_single(i: int) -> Status: async def _execute_single_run( self, workflow: Workflow, - task: Task, - run_index: int, - run_id_base: int, - model_wrapper: Optional[ModelWrapper] = None, - record_key: Optional[str] = None, + model_wrapper: ModelWrapper, ) -> Status: st = time.time() - model_wrapper = model_wrapper or self.model_wrapper - self._set_record_key(model_wrapper, record_key) await model_wrapper.clean_workflow_state() - run_id = run_id_base + run_index - workflow.set_execution_context(run_id=run_id) - self.runner_state["workflow_id"] = self._build_record_key(task, run_id) self.runner_state["terminate_time"] = None self.runner_state["begin_time"] = st try: status = await self._run_workflow(workflow) et = time.time() self.runner_state["terminate_time"] = et - metrics = [dict(metric) for metric in status.metrics] - if not metrics: - metrics = [{}] - for metric in metrics: - metric["time/run_execution"] = et - st - successful_run_ids = status.successful_run_ids or ( - [run_id] if status.completed_runs > 0 else [] - ) - status = Status( - completed_runs=len(successful_run_ids), - total_runs=status.total_runs, - metrics=metrics, - successful_run_ids=successful_run_ids, - message=status.message, - ) + if status.metrics: + for metric in status.metrics: + metric["time/run_execution"] = et - st return status except Exception as exc: self.runner_state["terminate_time"] = time.time() @@ -306,14 +270,13 @@ async def _run_task( collect_partial_runs: bool = True, ) -> Status: """Init workflow from the task and run it.""" - if issubclass(task.workflow, RepeatableWorkflow): - record_key = self._build_record_key(task, run_id_base) - workflow_instance = self._create_workflow_instance(task, record_key=record_key) - workflow_instance.set_execution_context(repeat_times, run_id_base) + if getattr(task.workflow, "can_repeat", False): + task.run_id = run_id_base + workflow_instance = self._create_workflow_instance(task) + workflow_instance.set_repeat_times(repeat_times, run_id_base) st = time.time() - self._set_record_key(self.model_wrapper, record_key) await self.model_wrapper.clean_workflow_state() - self.runner_state["workflow_id"] = record_key + self.runner_state["workflow_id"] = task.api_key self.runner_state["terminate_time"] = None self.runner_state["begin_time"] = st status = await self._run_workflow(workflow_instance) @@ -322,9 +285,9 @@ async def _run_task( run_metrics = [dict(metric) for metric in status.metrics] for metric in run_metrics: metric["time/run_execution"] = et - st - successful_run_ids = status.successful_run_ids or list( - range(run_id_base, run_id_base + status.completed_runs) - ) + # repeatable workflow shares the same run_id, so we can only return + # the run_id of the first run + successful_run_ids = [task.api_key] return self._build_status( total_runs=repeat_times, metrics=run_metrics, @@ -348,15 +311,11 @@ async def _sequential_run( ) -> Status: results = [] for i in range(repeat_times): - run_index = run_id_base + i - record_key = self._build_record_key(task, run_index) - workflow = self._create_workflow_instance(task, record_key=record_key) + task.run_id = run_id_base + i + workflow = self._create_workflow_instance(task) result = await self._execute_single_run( - workflow, - task, - i, - run_id_base, - record_key=record_key, + workflow=workflow, + model_wrapper=self.model_wrapper, ) results.append(result) if collect_partial_runs: @@ -374,9 +333,9 @@ async def _asynchronous_run( collect_partial_runs: bool = True, ) -> Status: return await self._run_parallel_runs( - task, - repeat_times, - run_id_base, + task=task, + repeat_times=repeat_times, + run_id_base=run_id_base, collect_partial_runs=collect_partial_runs, ) @@ -388,9 +347,9 @@ async def _multi_threading_run( collect_partial_runs: bool = True, ) -> Status: return await self._run_parallel_runs( - task, - repeat_times, - run_id_base, + task=task, + repeat_times=repeat_times, + run_id_base=run_id_base, collect_partial_runs=collect_partial_runs, use_threads=True, ) @@ -404,7 +363,6 @@ async def get_runner_state(self) -> Dict: async def run_task( self, task: Task, - batch_id: str, repeat_times: int = 1, run_id_base: int = 0, collect_partial_runs: bool = True, @@ -415,7 +373,7 @@ async def run_task( model_version = await self.model_wrapper.model_version_async self.runner_state["model_version"] = model_version self.logger.info( - f"Starting task: step={batch_id}, model_version={model_version}, repeat_times={repeat_times}, run_id_base={run_id_base}" + f"Starting task: step={task.batch_id}, task={task.task_id}, model_version={model_version}, repeat_times={repeat_times}, run_id_base={run_id_base}" ) status = await self._run_task( task, @@ -487,16 +445,16 @@ async def debug(self) -> None: """Run the debug workflow.""" tasks = await self.taskset.read(batch_size=1) task = tasks[0] + task.batch_id = "debug" + task.task_id = 0 self.logger.info(f"Start debugging task:\n{task.raw_task}") if not self.enable_profiling: - status = await self.run_task(task=task, batch_id="debug", repeat_times=1, run_id_base=0) + status = await self.run_task(task=task, repeat_times=1, run_id_base=0) else: from viztracer import VizTracer with VizTracer(output_file=self.output_profiling_file): - status = await self.run_task( - task=task, batch_id="debug", repeat_times=1, run_id_base=0 - ) + status = await self.run_task(task=task, repeat_times=1, run_id_base=0) experiences = [] if self.config.explorer.rollout_model.enable_history: try: From d422bbd9ef965778c6a5a28e4b719fb2085eb04f Mon Sep 17 00:00:00 2001 From: "panxuchen.pxc" Date: Tue, 30 Jun 2026 14:15:00 +0800 Subject: [PATCH 40/84] fix instance --- tests/explorer/scheduler_test.py | 8 +- tests/explorer/workflow_test.py | 147 +++++++----------- tests/utils/plugins/main.py | 9 +- tests/utils/plugins/my_workflow.py | 11 +- trinity/common/models/model.py | 30 ++-- trinity/common/workflows/__init__.py | 11 +- .../envs/alfworld/RAFT_alfworld_workflow.py | 5 +- .../envs/sciworld/sciworld_workflow.py | 5 +- .../workflows/on_policy_distill_workflow.py | 1 + trinity/common/workflows/workflow.py | 22 +-- trinity/explorer/scheduler.py | 16 +- trinity/explorer/workflow_runner.py | 19 ++- 12 files changed, 140 insertions(+), 144 deletions(-) diff --git a/tests/explorer/scheduler_test.py b/tests/explorer/scheduler_test.py index 0d4f1ace6b5..76bcf5c59da 100644 --- a/tests/explorer/scheduler_test.py +++ b/tests/explorer/scheduler_test.py @@ -16,12 +16,12 @@ from trinity.common.experience import EID, Experience from trinity.common.models.allocator import Allocator from trinity.common.models.model import InferenceModel, ModelWrapper -from trinity.common.workflows import WORKFLOWS, RepeatableWorkflow, Task, Workflow +from trinity.common.workflows import WORKFLOWS, Task, Workflow from trinity.explorer.scheduler import Scheduler @WORKFLOWS.register_module("dummy_workflow") -class DummyWorkflow(RepeatableWorkflow): +class DummyWorkflow(Workflow): def __init__(self, *, task, model, auxiliary_models): super().__init__(task=task, model=model, auxiliary_models=auxiliary_models) self.step_num = task.workflow_args.get("step_num", 1) @@ -185,7 +185,7 @@ def run(self): @WORKFLOWS.register_module("dummy_async_workflow") -class DummyAsyncWorkflow(RepeatableWorkflow): +class DummyAsyncWorkflow(Workflow): is_async: bool = True def __init__(self, *, task, model, auxiliary_models): @@ -219,7 +219,7 @@ def run(self): @WORKFLOWS.register_module("dummy_workflow_with_state") -class DummyWorkflowWithState(RepeatableWorkflow): +class DummyWorkflowWithState(Workflow): is_async: bool = True def __init__(self, *, task, model: ModelWrapper, auxiliary_models): diff --git a/tests/explorer/workflow_test.py b/tests/explorer/workflow_test.py index 33832accc6d..a657001c546 100644 --- a/tests/explorer/workflow_test.py +++ b/tests/explorer/workflow_test.py @@ -29,21 +29,11 @@ from trinity.common.experience import EID, Experience from trinity.common.models.allocator import Allocator from trinity.common.models.model import ModelWrapper -from trinity.common.workflows import WORKFLOWS, RepeatableWorkflow, Workflow +from trinity.common.workflows import WORKFLOWS, Workflow from trinity.common.workflows.customized_math_workflows import MathBoxedWorkflow from trinity.common.workflows.eval_workflow import MathEvalWorkflow -from trinity.common.workflows.workflow import ( - MathWorkflow, - RepeatableMultiTurnWorkflow, - Task, -) -from trinity.explorer.workflow_runner import WorkflowRunner - - -def deserialize_experiences(exp_payload: bytes) -> list[Experience]: - if not exp_payload: - return [] - return Experience.deserialize_many(exp_payload) +from trinity.common.workflows.workflow import MathWorkflow, MultiTurnWorkflow, Task +from trinity.explorer.workflow_runner import Status, WorkflowRunner def patch_runner_models(*wrappers): @@ -67,7 +57,7 @@ class MockResponse: action_mask: Optional[Tensor] = None -class DummyWorkflow(RepeatableWorkflow): +class DummyWorkflow(Workflow): can_reset: bool = True def __init__(self, model, task: Task, auxiliary_models=None): @@ -110,7 +100,7 @@ def run(self): raise ValueError("Invalid output format") -class DummyAsyncWorkflow(RepeatableWorkflow): +class DummyAsyncWorkflow(Workflow): can_reset: bool = True is_async: bool = True @@ -155,7 +145,7 @@ async def run_async(self): raise ValueError("Invalid output format") -class DummyMultiTurnWorkflow(RepeatableMultiTurnWorkflow): +class DummyMultiTurnWorkflow(MultiTurnWorkflow): def __init__(self, model, task: Task, auxiliary_models=None): super().__init__(task=task, model=model, auxiliary_models=auxiliary_models) self.contents = task.raw_task["contents"] # type: ignore @@ -171,7 +161,7 @@ def run(self): return experience_list -class DummyAsyncMultiTurnWorkflow(RepeatableMultiTurnWorkflow): +class DummyAsyncMultiTurnWorkflow(MultiTurnWorkflow): is_async: bool = True def __init__(self, model, task: Task, auxiliary_models=None): @@ -707,14 +697,12 @@ async def test_workflow_runner(self): workflow_args={"output_format": "json"}, ) - status, exps = await runner.run_task( - task, batch_id="test", repeat_times=3, run_id_base=0 - ) - exps = deserialize_experiences(exps) + status = await runner.run_task(task, repeat_times=3, run_id_base=0) self.assertTrue(status.ok) - self.assertIsInstance(exps, list) - self.assertEqual(len(exps), 3) + self.assertEqual(status.completed_runs, 3) + self.assertEqual(status.total_runs, 3) + self.assertEqual(len(status.metrics), 3) task = Task( workflow=DummyAsyncWorkflow, @@ -723,13 +711,11 @@ async def test_workflow_runner(self): workflow_args={"output_format": "yaml"}, ) - status, exps = await runner.run_task( - task, batch_id="test", repeat_times=2, run_id_base=0 - ) - exps = deserialize_experiences(exps) + status = await runner.run_task(task, repeat_times=2, run_id_base=0) self.assertTrue(status.ok) - self.assertIsInstance(exps, list) - self.assertEqual(len(exps), 2) + self.assertEqual(status.completed_runs, 2) + self.assertEqual(status.total_runs, 2) + self.assertEqual(len(status.metrics), 2) @parameterized.expand( [ @@ -760,13 +746,11 @@ async def test_workflow_runner_partial_success_non_repeatable( task_id=0, ) - status, exps = await runner.run_task(task, repeat_times=3, run_id_base=0) - exps = deserialize_experiences(exps) + status = await runner.run_task(task, repeat_times=3, run_id_base=0) self.assertFalse(status.ok) self.assertEqual(status.completed_runs, expected_success_runs) self.assertEqual(status.total_runs, 3) - self.assertEqual(len(exps), expected_success_runs) # One internal run fails with call_id=1, so runner-level metrics should # retain only the successful runs from this single subtask: call_id=0 and 2. @@ -775,14 +759,11 @@ async def test_workflow_runner_partial_success_non_repeatable( sorted(metric["run_metrics"] for metric in status.metrics), [0.0, 2.0], ) - - # Experiences returned from the runner should match the same successful - # run set, proving failed runs do not leak into partial-return outputs. - self.assertEqual( - sorted(exp.metrics["run_metrics"] for exp in exps if exp.metrics), - [0.0, 2.0], + assert status.message is not None + self.assertIn( + f"{expected_success_runs}/3 runs completed successfully", + status.message, ) - self.assertIn(f"{expected_success_runs}/3 runs completed successfully", status.message) # type: ignore[arg-type] @parameterized.expand( [ @@ -811,44 +792,50 @@ async def test_workflow_runner_fail_fast_without_partial_collection(self, concur async def mock_execute_single_run( workflow: Workflow, - task: Task, - run_index: int, - run_id_base: int, + model_wrapper: ModelWrapper, ): + run_index = int(workflow.task.run_id) if run_index == 0: await asyncio.sleep(0.01) - exp = Experience( - tokens=Tensor([0, 1, 2]), - prompt_length=1, - metrics={"run_metrics": 0.0}, + return Status( + completed_runs=1, + total_runs=1, + metrics=[{"run_metrics": 0.0}], + successful_ids=[workflow.task.api_key], ) - return True, [exp], {"run_metrics": 0.0}, None if run_index == 1: await asyncio.sleep(0.02) - return False, [], None, "planned failure" + return Status( + completed_runs=0, + total_runs=1, + metrics=[], + message="planned failure", + ) await asyncio.sleep(0.5) - exp = Experience( - tokens=Tensor([0, 1, 2]), - prompt_length=1, - metrics={"run_metrics": 2.0}, + return Status( + completed_runs=1, + total_runs=1, + metrics=[{"run_metrics": 2.0}], + successful_ids=[workflow.task.api_key], ) - return True, [exp], {"run_metrics": 2.0}, None runner._execute_single_run = AsyncMock(side_effect=mock_execute_single_run) - status, exps = await runner.run_task( + status = await runner.run_task( task, repeat_times=3, run_id_base=0, collect_partial_runs=False, ) - exps = deserialize_experiences(exps) self.assertFalse(status.ok) self.assertEqual(status.completed_runs, 1) self.assertEqual(status.total_runs, 3) - self.assertEqual(len(exps), 1) - self.assertIn("1/3 runs completed successfully", status.message) # type: ignore[arg-type] + assert status.message is not None + self.assertIn( + "1/3 runs completed successfully", + status.message, + ) async def test_workflow_runner_get_state(self): config = get_template_config() @@ -905,7 +892,7 @@ async def monitor_routine(): await asyncio.gather( *[ monitor_routine(), - runner.run_task(task, batch_id="test", repeat_times=3, run_id_base=0), + runner.run_task(task, repeat_times=3, run_id_base=0), ] ) @@ -938,20 +925,13 @@ async def test_workflow_with_openai(self): ), ] - status, exps = await runner.run_task( - tasks[0], batch_id="test", repeat_times=2, run_id_base=0 - ) # test exception handling - exps = deserialize_experiences(exps) + status = await runner.run_task(tasks[0], repeat_times=2, run_id_base=0) self.assertEqual(status.ok, False) - self.assertEqual(len(exps), 0) exps = runner.model_wrapper.extract_experience_from_history(clear_history=False) self.assertEqual(len(exps), 1) - status, exps = await runner.run_task( - tasks[1], batch_id="test", repeat_times=2, run_id_base=0 - ) # normal run - exps = deserialize_experiences(exps) + status = await runner.run_task(tasks[1], repeat_times=2, run_id_base=0) self.assertEqual(status.ok, True) - self.assertEqual(len(exps), 2) + self.assertEqual(status.completed_runs, 2) exps = runner.model_wrapper.extract_experience_from_history(clear_history=False) self.assertEqual(len(exps), 0) self.assertEqual(len(rollout_model), 1) @@ -1068,43 +1048,30 @@ async def test_concurrent_workflow_runner(self): ) # warmup - async_status, async_exps = await async_runner.run_task.remote( - task, batch_id="test", repeat_times=2, run_id_base=0 - ) + async_status = await async_runner.run_task.remote(task, repeat_times=2, run_id_base=0) st = time.time() - async_status, async_exps = await async_runner.run_task.remote( - task, batch_id="test", repeat_times=4, run_id_base=0 - ) + async_status = await async_runner.run_task.remote(task, repeat_times=4, run_id_base=0) async_runtime = time.time() - st # warmup - thread_status, thread_exps = await thread_runner.run_task.remote( - task, batch_id="test", repeat_times=1, run_id_base=0 - ) + thread_status = await thread_runner.run_task.remote(task, repeat_times=1, run_id_base=0) st = time.time() - thread_status, thread_exps = await thread_runner.run_task.remote( - task, batch_id="test", repeat_times=4, run_id_base=0 - ) + thread_status = await thread_runner.run_task.remote(task, repeat_times=4, run_id_base=0) thread_runtime = time.time() - st st = time.time() - sequential_status, sequential_exps = await sequential_runner.run_task.remote( - task, batch_id="test", repeat_times=4, run_id_base=0 + sequential_status = await sequential_runner.run_task.remote( + task, repeat_times=4, run_id_base=0 ) sequential_runtime = time.time() - st self.assertTrue(async_status.ok) self.assertTrue(thread_status.ok) self.assertTrue(sequential_status.ok) - - async_exps = deserialize_experiences(async_exps) - thread_exps = deserialize_experiences(thread_exps) - sequential_exps = deserialize_experiences(sequential_exps) - - self.assertEqual(len(async_exps), 8) - self.assertEqual(len(thread_exps), 8) - self.assertEqual(len(sequential_exps), 8) + self.assertEqual(async_status.completed_runs, 4) + self.assertEqual(thread_status.completed_runs, 4) + self.assertEqual(sequential_status.completed_runs, 4) self.assertLessEqual(async_runtime * 2, sequential_runtime) self.assertLessEqual(thread_runtime * 2, sequential_runtime) diff --git a/tests/utils/plugins/main.py b/tests/utils/plugins/main.py index 92129097295..fd217b08fe8 100644 --- a/tests/utils/plugins/main.py +++ b/tests/utils/plugins/main.py @@ -1,10 +1,15 @@ from tests.utils.plugins.dependencies import DEPENDENCY_VALUE, dependency_func -from trinity.common.workflows.workflow import RepeatableWorkflow +from trinity.common.workflows.workflow import Workflow -class MainDummyWorkflow(RepeatableWorkflow): +class MainDummyWorkflow(Workflow): + can_repeat: bool = True + def __init__(self, *, task, model, auxiliary_models=None): super().__init__(task=task, model=model, auxiliary_models=auxiliary_models) + def set_repeat_times(self, repeat_times, run_id_base): + pass + def run(self) -> list: return [DEPENDENCY_VALUE, dependency_func()] diff --git a/tests/utils/plugins/my_workflow.py b/tests/utils/plugins/my_workflow.py index 505f624271e..471b2371cff 100644 --- a/tests/utils/plugins/my_workflow.py +++ b/tests/utils/plugins/my_workflow.py @@ -1,12 +1,19 @@ from typing import List -from trinity.common.workflows import WORKFLOWS, RepeatableWorkflow +from trinity.common.workflows import WORKFLOWS, Workflow @WORKFLOWS.register_module("my_workflow") -class MyWorkflow(RepeatableWorkflow): +class MyWorkflow(Workflow): def __init__(self, *, task, model, auxiliary_models=None): super().__init__(task=task, model=model, auxiliary_models=auxiliary_models) + @property + def repeatable(self): + return True + + def set_repeat_times(self, repeat_times, run_id_base): + pass + def run(self) -> List: return ["Hello world", "Hi"] diff --git a/trinity/common/models/model.py b/trinity/common/models/model.py index 20713dc7998..1afc444e35e 100644 --- a/trinity/common/models/model.py +++ b/trinity/common/models/model.py @@ -142,7 +142,7 @@ def get_api_key(self) -> str: return "EMPTY" async def extract_experience_from_history( - self, record_key: str, clear_history: bool = True + self, key: str, clear_history: bool = True ) -> List[Experience]: """Extract recorded experiences by record key from the in-process store. @@ -153,7 +153,7 @@ async def extract_experience_from_history( ``.store`` is a ``RecordStore``); this base implementation is shared. """ return await self._collect_experiences( - record_key, + key, remove=clear_history, ) @@ -833,19 +833,19 @@ async def sync_model_weights( self._model_path = await self.model.get_model_path.remote() def extract_experience_from_history( - self, clear_history: bool = True, record_key: Optional[str] = None + self, clear_history: bool = True, key: Optional[str] = None ) -> List[Experience]: """Extract experiences from the history.""" if not self.enable_history: raise ValueError("History recording is not enabled.") if self.model is None: raise ValueError("Recording extraction requires an inference model actor.") - record_key = record_key or self._api_key - if record_key is None: - raise ValueError("record_key is required when recording is enabled.") + key = key or self._api_key + if key is None: + raise ValueError("key is required when recording is enabled.") exps = ray.get( self.model.extract_experience_from_history.remote( - record_key=record_key, + key=key, clear_history=clear_history, ) ) @@ -853,21 +853,21 @@ def extract_experience_from_history( async def update_experience_reward_async( self, - record_key: str, + key: str, reward: float, info: Optional[dict] = None, sample_ids: Optional[List[str]] = None, ) -> None: """Update reward and optional info on recorded experiences.""" await self.update_experience_records_async( - record_key=record_key, + key=key, update=ExperienceUpdate(reward=reward, info=info), sample_ids=sample_ids, ) async def update_experience_records_async( self, - record_key: str, + key: str, update: ExperienceUpdate, sample_ids: Optional[List[str]] = None, ) -> None: @@ -877,28 +877,28 @@ async def update_experience_records_async( if self.model is None: raise ValueError("Recording update requires an inference model actor.") await self.model.update_experience_records.remote( - record_key=record_key, + key=key, update=update, sample_ids=sample_ids, ) def update_experience_reward( self, - record_key: str, + key: str, reward: float, info: Optional[dict] = None, sample_ids: Optional[List[str]] = None, ) -> None: """Update reward and optional info on recorded experiences.""" self.update_experience_records( - record_key=record_key, + key=key, update=ExperienceUpdate(reward=reward, info=info), sample_ids=sample_ids, ) def update_experience_records( self, - record_key: str, + key: str, update: ExperienceUpdate, sample_ids: Optional[List[str]] = None, ) -> None: @@ -909,7 +909,7 @@ def update_experience_records( raise ValueError("Recording update requires an inference model actor.") ray.get( self.model.update_experience_records.remote( - record_key=record_key, + key=key, update=update, sample_ids=sample_ids, ) diff --git a/trinity/common/workflows/__init__.py b/trinity/common/workflows/__init__.py index 34bd27be59c..140fbcd2217 100644 --- a/trinity/common/workflows/__init__.py +++ b/trinity/common/workflows/__init__.py @@ -1,6 +1,13 @@ # -*- coding: utf-8 -*- """Workflow module""" -from trinity.common.workflows.workflow import Metrics, Status, Task, Workflow +from trinity.common.workflows.workflow import ( + Metrics, + Status, + Task, + Workflow, + WorkflowBase, + WorkflowWithRecording, +) from trinity.utils.registry import Registry WORKFLOWS: Registry = Registry( @@ -60,5 +67,7 @@ "Status", "Metrics", "Workflow", + "WorkflowBase", + "WorkflowWithRecording", "WORKFLOWS", ] diff --git a/trinity/common/workflows/envs/alfworld/RAFT_alfworld_workflow.py b/trinity/common/workflows/envs/alfworld/RAFT_alfworld_workflow.py index deb13de4990..4bc34833cb8 100644 --- a/trinity/common/workflows/envs/alfworld/RAFT_alfworld_workflow.py +++ b/trinity/common/workflows/envs/alfworld/RAFT_alfworld_workflow.py @@ -13,10 +13,10 @@ process_messages_to_experience_async, validate_trajectory_format, ) -from trinity.common.workflows.workflow import RepeatableWorkflow, Task +from trinity.common.workflows.workflow import Task, Workflow -class RAFTAlfworldWorkflow(RepeatableWorkflow): +class RAFTAlfworldWorkflow(Workflow): """ RAFT workflow for alfworld using trajectory context. @@ -26,6 +26,7 @@ class RAFTAlfworldWorkflow(RepeatableWorkflow): """ can_reset: bool = True + can_repeat: bool = True is_async: bool = True def __init__( diff --git a/trinity/common/workflows/envs/sciworld/sciworld_workflow.py b/trinity/common/workflows/envs/sciworld/sciworld_workflow.py index fb2bd6ecfa7..d5aeb1f2178 100644 --- a/trinity/common/workflows/envs/sciworld/sciworld_workflow.py +++ b/trinity/common/workflows/envs/sciworld/sciworld_workflow.py @@ -4,7 +4,7 @@ from trinity.common.experience import Experience from trinity.common.models.model import ModelWrapper -from trinity.common.workflows.workflow import RepeatableMultiTurnWorkflow, Task +from trinity.common.workflows.workflow import MultiTurnWorkflow, Task SCIWORLD_SYSTEM_PROMPT = """ You are an agent, your job is to do some scientific experiment in a virtual text-based environment. @@ -55,10 +55,11 @@ def parse_action(response): return "" -class SciWorldWorkflow(RepeatableMultiTurnWorkflow): +class SciWorldWorkflow(MultiTurnWorkflow): """A workflow for sciworld task.""" is_async: bool = True + can_repeat: bool = True def __init__( self, diff --git a/trinity/common/workflows/on_policy_distill_workflow.py b/trinity/common/workflows/on_policy_distill_workflow.py index e9251033741..f84818d35c3 100644 --- a/trinity/common/workflows/on_policy_distill_workflow.py +++ b/trinity/common/workflows/on_policy_distill_workflow.py @@ -117,6 +117,7 @@ async def run_async(self) -> List[Experience]: resp_start = response.prompt_length - 1 teacher_resp_logprobs = teacher_logprobs[resp_start:] student_resp_logprobs = response.logprobs + assert student_resp_logprobs is not None, "Student logprobs should not be None." # Verify lengths match (they should be equal for the same token sequence) assert len(teacher_resp_logprobs) == len(student_resp_logprobs), ( diff --git a/trinity/common/workflows/workflow.py b/trinity/common/workflows/workflow.py index 19727791252..f4d1179912b 100644 --- a/trinity/common/workflows/workflow.py +++ b/trinity/common/workflows/workflow.py @@ -25,7 +25,7 @@ class Status: completed_runs: int total_runs: int metrics: List[Dict[str, float]] - successful_run_ids: List[str] = field(default_factory=list) + successful_ids: List[str] = field(default_factory=list) message: Optional[str] = None @property @@ -101,13 +101,14 @@ def to_dict(self) -> dict: return self.raw_task # type: ignore -class BaseWorkflow: +class WorkflowBase: """The base workflow interface.""" def __init__(self, task: Task, model: ModelWrapper) -> None: self.task = task self.model = model - self.model.set_api_key(task.api_key) # set the API key for the rollout model + if self.model is not None: + self.model.set_api_key(task.api_key) # set the API key for the rollout model self.logger = get_logger(__name__) @abstractmethod @@ -119,10 +120,11 @@ def reset(self, task: Task): self.task.batch_id = task.batch_id self.task.task_id = task.task_id self.task.run_id = task.run_id - self.model.set_api_key(task.api_key) # set the API key for the rollout model + if self.model is not None: + self.model.set_api_key(task.api_key) # set the API key for the rollout model -class Workflow(BaseWorkflow): +class Workflow(WorkflowBase): """The base workflow class. A workflow is a runnable object which generates a list of experiences. @@ -192,7 +194,7 @@ async def execute(self) -> Status: completed_runs=self.__class__.can_repeat and self.repeat_times or 1, total_runs=self.__class__.can_repeat and self.repeat_times or 1, metrics=[exp.metrics for exp in exps if exp.metrics is not None], - successful_run_ids=[self.task.api_key], + successful_ids=[self.task.api_key], ) @@ -416,7 +418,7 @@ class AsyncMathWorkflow(AsyncSimpleWorkflow, MathWorkflow): pass -class WorkflowWithRecording(BaseWorkflow): +class WorkflowWithRecording(WorkflowBase): """A workflow that using the rollout model's built-in recording path to capture experience data. @@ -441,7 +443,7 @@ def __init__( # Store ModelWrapper instances self.auxiliary_model_wrappers = auxiliary_models # Get OpenAI clients from ModelWrapper - self.auxiliary_models = [m.get_openai_async_client() for m in auxiliary_models] + self.auxiliary_models = [m.get_openai_async_client() for m in (auxiliary_models or [])] @property def base_url(self) -> str: @@ -471,7 +473,7 @@ def _to_status(self, result: Metrics) -> Status: completed_runs=1, total_runs=1, metrics=[result], - successful_run_ids=[self.task.api_key], + successful_ids=[self.task.api_key], ) async def update_reward( @@ -482,7 +484,7 @@ async def update_reward( ) -> None: """Update recorded experiences for one run with reward and optional info.""" await self.model.update_experience_reward_async( - record_key=self.api_key, + key=self.api_key, reward=reward, info=info, sample_ids=sample_ids, diff --git a/trinity/explorer/scheduler.py b/trinity/explorer/scheduler.py index e19c32d509f..c0f2237c34d 100644 --- a/trinity/explorer/scheduler.py +++ b/trinity/explorer/scheduler.py @@ -12,7 +12,7 @@ from ray.actor import ActorHandle from trinity.common.config import Config -from trinity.common.workflows import RepeatableWorkflow, Task +from trinity.common.workflows import Task from trinity.explorer.workflow_runner import Status, WorkflowRunner from trinity.utils.log import get_logger from trinity.utils.metrics import calculate_task_level_metrics @@ -31,7 +31,7 @@ class TaskWrapper: completed_runs: int = 0 total_runs: int = 0 # total planned runs for the whole task metrics: List[Dict[str, float]] = field(default_factory=list) - successful_run_ids: List[str] = field(default_factory=list) + successful_ids: List[str] = field(default_factory=list) experience_payloads: List[bytes] = field(default_factory=list) first_error: Optional[str] = None emitted: bool = False @@ -134,10 +134,10 @@ async def _delete_records(self, prefix: str) -> None: self.logger.error("records delete from rollout actor failed: %s", exc) async def _consume_finished_records(self, task: TaskWrapper, status: Status) -> List[bytes]: - if not status.successful_run_ids: + if not status.successful_ids: return [] - if issubclass(task.task.workflow, RepeatableWorkflow): + if getattr(task.task.workflow, "can_repeat", False): prefix = self._task_level_record_key(task) if task.task.is_eval: await self._delete_records(prefix) @@ -148,12 +148,12 @@ async def _consume_finished_records(self, task: TaskWrapper, status: Status) -> if task.task.is_eval: await asyncio.gather( - *[self._delete_records(run_id) for run_id in status.successful_run_ids] + *[self._delete_records(run_id) for run_id in status.successful_ids] ) return [] payloads = await asyncio.gather( - *[self._drain_records(run_id) for run_id in status.successful_run_ids] + *[self._drain_records(run_id) for run_id in status.successful_ids] ) return [payload for payload in payloads if payload] @@ -467,7 +467,7 @@ def _accumulate_task_result( task.finished_sub_task_num += 1 task.completed_runs += status.completed_runs task.metrics.extend(status.metrics) - task.successful_run_ids.extend(status.successful_run_ids) + task.successful_ids.extend(status.successful_ids) task.experience_payloads.extend(experience_payloads) if not status.ok and task.first_error is None: task.first_error = status.message @@ -485,7 +485,7 @@ def _build_task_result(self, task: TaskWrapper) -> Status: completed_runs=task.completed_runs, total_runs=task.total_runs, metrics=[calculate_task_level_metrics(task.metrics, task.task.is_eval)], - successful_run_ids=sorted(task.successful_run_ids), + successful_ids=sorted(task.successful_ids), message=message, ) return status diff --git a/trinity/explorer/workflow_runner.py b/trinity/explorer/workflow_runner.py index b54395068ed..a15f207cd16 100644 --- a/trinity/explorer/workflow_runner.py +++ b/trinity/explorer/workflow_runner.py @@ -125,11 +125,11 @@ def _create_isolated_workflow_instance( def _build_status( self, total_runs: int, + completed_runs: int, metrics: List[Dict[str, float]], - successful_run_ids: List[str], + successful_ids: List[str], first_error: Optional[str] = None, ) -> Status: - completed_runs = len(successful_run_ids) if first_error is None: message = None elif completed_runs > 0: @@ -144,7 +144,7 @@ def _build_status( completed_runs=completed_runs, total_runs=total_runs, metrics=list(metrics), - successful_run_ids=list(successful_run_ids), + successful_ids=list(successful_ids), message=message, ) @@ -154,21 +154,24 @@ def _aggregate_run_results( results: List[Status], ) -> Status: run_metrics = [] - successful_run_ids = [] + successful_ids = [] + completed_runs = 0 first_error = None for status in results: + completed_runs += status.completed_runs if status.ok: run_metrics.extend(status.metrics) - successful_run_ids.extend(status.successful_run_ids) + successful_ids.extend(status.successful_ids) continue if first_error is None: first_error = status.message return self._build_status( total_runs=total_runs, + completed_runs=completed_runs, metrics=run_metrics, - successful_run_ids=successful_run_ids, + successful_ids=successful_ids, first_error=first_error, ) @@ -287,11 +290,11 @@ async def _run_task( metric["time/run_execution"] = et - st # repeatable workflow shares the same run_id, so we can only return # the run_id of the first run - successful_run_ids = [task.api_key] return self._build_status( total_runs=repeat_times, + completed_runs=status.completed_runs, metrics=run_metrics, - successful_run_ids=successful_run_ids, + successful_ids=status.successful_ids or [task.api_key], first_error=status.message, ) else: From 81f308c39e2873907d6e7946eebb93f5cd0cb400 Mon Sep 17 00:00:00 2001 From: "panxuchen.pxc" Date: Tue, 30 Jun 2026 16:28:42 +0800 Subject: [PATCH 41/84] optimizer prefix merger --- tests/common/recording_recorder_test.py | 105 +++++++++++++++ trinity/common/models/recording/merger.py | 157 ++++++++++++++++++---- 2 files changed, 239 insertions(+), 23 deletions(-) diff --git a/tests/common/recording_recorder_test.py b/tests/common/recording_recorder_test.py index bede2cb445d..fa945673d38 100644 --- a/tests/common/recording_recorder_test.py +++ b/tests/common/recording_recorder_test.py @@ -158,6 +158,111 @@ async def test_merge_head_replaces_only_matching_sample_stream(self): ) ) + async def test_interleaved_branches_with_shared_sample_index_merge_independently(self): + store = MemoryStore() + recorder = Recorder( + store=store, + build_experiences=lambda *_args, **_kwargs: [], + enabled=True, + ) + record_key = "0/task_a/1" + branch_a_first = make_turn( + request_id="req-a1", + record_key=record_key, + tokens=[10, 11, 20], + prompt_length=2, + logprobs=[-0.2], + sample_index=0, + ) + branch_b_first = make_turn( + request_id="req-b1", + record_key=record_key, + tokens=[10, 12, 21], + prompt_length=2, + logprobs=[-0.3], + sample_index=0, + ) + branch_a_final = make_turn( + request_id="req-a2", + record_key=record_key, + tokens=[10, 11, 20, 13, 30], + prompt_length=4, + logprobs=[-0.4], + sample_index=0, + ) + branch_b_final = make_turn( + request_id="req-b2", + record_key=record_key, + tokens=[10, 12, 21, 14, 31], + prompt_length=4, + logprobs=[-0.5], + sample_index=0, + ) + + await recorder._safe_append(branch_a_first) + await recorder._safe_append(branch_b_first) + await recorder._safe_append(branch_a_final) + await recorder._safe_append(branch_b_final) + + recorded = store.get(record_key) + self.assertEqual(len(recorded), 2) + self.assertEqual({exp.eid.suffix for exp in recorded}, {"req-a2", "req-b2"}) + merged_by_suffix = {exp.eid.suffix: exp for exp in recorded} + self.assertEqual( + merged_by_suffix["req-a2"].info["merged_eid_suffixes"], ["req-a1", "req-a2"] + ) + self.assertEqual( + merged_by_suffix["req-b2"].info["merged_eid_suffixes"], ["req-b1", "req-b2"] + ) + + async def test_multi_head_merge_uses_longest_matching_prefix(self): + store = MemoryStore() + recorder = Recorder( + store=store, + build_experiences=lambda *_args, **_kwargs: [], + enabled=True, + ) + record_key = "0/task_a/1" + short_prefix = make_turn( + request_id="req-short", + record_key=record_key, + tokens=[10, 11, 20], + prompt_length=2, + logprobs=[-0.2], + ) + long_prefix = make_turn( + request_id="req-long", + record_key=record_key, + tokens=[10, 11, 20, 12, 30], + prompt_length=4, + logprobs=[-0.3], + ) + unrelated = make_turn( + request_id="req-other", + record_key=record_key, + tokens=[10, 13, 21], + prompt_length=2, + logprobs=[-0.4], + ) + final = make_turn( + request_id="req-final", + record_key=record_key, + tokens=[10, 11, 20, 12, 30, 14, 40], + prompt_length=6, + logprobs=[-0.5], + ) + + await recorder._safe_append(short_prefix) + await recorder._safe_append(long_prefix) + await recorder._safe_append(unrelated) + await recorder._safe_append(final) + + recorded = store.get(record_key) + self.assertEqual(len(recorded), 2) + merged = next(exp for exp in recorded if exp.eid.suffix == "req-final") + self.assertEqual(merged.info["merged_eid_suffixes"], ["req-short", "req-long", "req-final"]) + self.assertEqual(merged.info["merged_turn_count"], 3) + async def test_stale_merge_head_falls_back_to_append(self): store = MemoryStore() recorder = Recorder( diff --git a/trinity/common/models/recording/merger.py b/trinity/common/models/recording/merger.py index cf3bc4295d8..4f328dadee2 100644 --- a/trinity/common/models/recording/merger.py +++ b/trinity/common/models/recording/merger.py @@ -1,6 +1,7 @@ """Prefix-based merging for recorded multi-turn experiences.""" from collections.abc import Sequence +from dataclasses import dataclass from typing import Any, Optional import torch @@ -8,6 +9,8 @@ from trinity.buffer.store import RecordStore, get_sample_id from trinity.common.experience import Experience +MAX_HEADS_PER_STREAM = 128 + class PrefixExperienceMerger: """Merge same-record experiences whose tokens form a strict prefix chain. @@ -15,35 +18,27 @@ class PrefixExperienceMerger: Strategy: * Experiences are grouped by record key and a best-effort sample stream key (sample_index, then default). - * Each stream tracks one latest/longest head experience. A new experience - merges only when the head tokens are a strict prefix of the new tokens. - * If no head exists yet, the store is scanned once to find the longest - prefix-matching experience for that stream. - - Limitation: - This assumes each record/sample stream is a single linear conversation - branch. If one task has concurrent writers sharing the same record/sample - stream (for example, multi-agent rollouts under one record key), the latest - head may belong to a different branch, so the prefix hit can be missed or - become ambiguous. - - TODO(recording): support branching/concurrent histories by tracking multiple - heads per record/sample stream, keyed by a stable conversation/thread id or - by token-prefix fingerprints. + * Each stream tracks multiple latest/longest heads so interleaved + branches sharing one record/sample stream do not evict each other. + * A length index tries longer heads first; exact token prefix comparison + remains the source of truth. + * If no cached head exists yet, the store is scanned once to seed the + stream cache from previously appended experiences. """ def __init__(self, store: RecordStore) -> None: self.store = store - self._heads: dict[str, dict[tuple[str, Any], Experience]] = {} + self._heads: dict[str, dict[tuple[str, Any], _StreamHeads]] = {} def try_merge(self, record_key: str, exp: Experience) -> bool: stream_key = _sample_stream_key(exp) heads = self._heads.setdefault(record_key, {}) - candidate = heads.get(stream_key) - if candidate is None: + stream_heads = heads.setdefault(stream_key, _StreamHeads()) + candidate = stream_heads.find_longest_prefix(exp) + if candidate is None and stream_heads.is_empty(): candidate = _find_longest_prefix_experience(self.store.get(record_key), exp) - elif not _is_strict_token_prefix(candidate.tokens, exp.tokens): - return False + if candidate is not None: + stream_heads.remember(candidate) if candidate is None: return False @@ -52,18 +47,114 @@ def try_merge(self, record_key: str, exp: Experience) -> bool: try: self.store.replace(record_key, old_sample_id, merged) except KeyError: - heads.pop(stream_key, None) + stream_heads.discard_sample_id(old_sample_id) return False - heads[stream_key] = merged + stream_heads.discard_sample_id(old_sample_id) + stream_heads.remember(merged) return True def remember(self, record_key: str, exp: Experience) -> None: - self._heads.setdefault(record_key, {})[_sample_stream_key(exp)] = exp + heads = self._heads.setdefault(record_key, {}) + heads.setdefault(_sample_stream_key(exp), _StreamHeads()).remember(exp) def forget_record(self, record_key: str) -> None: self._heads.pop(record_key, None) +@dataclass +class _HeadEntry: + exp: Experience + sequence: int + signature: tuple[int, ...] + + +class _StreamHeads: + """Small in-memory index of possible heads for one record/sample stream.""" + + def __init__(self, max_heads: int = MAX_HEADS_PER_STREAM) -> None: + self.max_heads = max_heads + self._heads_by_sample_id: dict[str, _HeadEntry] = {} + self._sample_ids_by_length: dict[int, set[str]] = {} + self._sample_ids_by_fingerprint: dict[tuple[int, tuple[int, ...]], set[str]] = {} + self._lengths_desc: list[int] = [] + self._sequence = 0 + + def is_empty(self) -> bool: + return not self._heads_by_sample_id + + def remember(self, exp: Experience) -> None: + sample_id = get_sample_id(exp) + self.discard_sample_id(sample_id) + self._sequence += 1 + length = len(exp.tokens) + signature = _prefix_signature(exp.tokens, length) + self._heads_by_sample_id[sample_id] = _HeadEntry( + exp=exp, + sequence=self._sequence, + signature=signature, + ) + sample_ids = self._sample_ids_by_length.setdefault(length, set()) + if not sample_ids: + self._insert_length(length) + sample_ids.add(sample_id) + self._sample_ids_by_fingerprint.setdefault((length, signature), set()).add(sample_id) + self._evict_excess_heads() + + def discard_sample_id(self, sample_id: str) -> None: + entry = self._heads_by_sample_id.pop(sample_id, None) + if entry is None: + return + length = len(entry.exp.tokens) + fingerprint_key = (length, entry.signature) + fingerprint_sample_ids = self._sample_ids_by_fingerprint.get(fingerprint_key) + if fingerprint_sample_ids is not None: + fingerprint_sample_ids.discard(sample_id) + if not fingerprint_sample_ids: + self._sample_ids_by_fingerprint.pop(fingerprint_key, None) + sample_ids = self._sample_ids_by_length.get(length) + if sample_ids is None: + return + sample_ids.discard(sample_id) + if not sample_ids: + self._sample_ids_by_length.pop(length, None) + self._lengths_desc.remove(length) + + def find_longest_prefix(self, exp: Experience) -> Optional[Experience]: + exp_length = len(exp.tokens) + for length in self._lengths_desc: + if length >= exp_length: + continue + signature = _prefix_signature(exp.tokens, length) + best_entry = None + sample_ids = self._sample_ids_by_fingerprint.get((length, signature), ()) + for sample_id in sample_ids: + entry = self._heads_by_sample_id.get(sample_id) + if entry is None: + continue + if _is_strict_token_prefix(entry.exp.tokens, exp.tokens): + if best_entry is None or entry.sequence < best_entry.sequence: + best_entry = entry + if best_entry is not None: + return best_entry.exp + return None + + def _insert_length(self, length: int) -> None: + index = 0 + while index < len(self._lengths_desc) and self._lengths_desc[index] > length: + index += 1 + self._lengths_desc.insert(index, length) + + def _evict_excess_heads(self) -> None: + while len(self._heads_by_sample_id) > self.max_heads: + shortest_length = self._lengths_desc[-1] + sample_ids = self._sample_ids_by_length[shortest_length] + oldest_sample_id = min( + sample_ids, + key=lambda sample_id: self._heads_by_sample_id[sample_id].sequence, + ) + self.discard_sample_id(oldest_sample_id) + + def _find_longest_prefix_experience( existing: Sequence[Experience], exp: Experience, @@ -86,6 +177,26 @@ def _same_sample_stream(left: Experience, right: Experience) -> bool: return _sample_stream_key(left) == _sample_stream_key(right) +def _prefix_signature(tokens: torch.Tensor, length: int) -> tuple[int, ...]: + """Return a cheap, collision-tolerant signature for ``tokens[:length]``. + + This only narrows candidates. ``_is_strict_token_prefix`` still performs the + exact comparison before any merge. + """ + if length <= 0: + return () + positions = { + 0, + length // 3, + (2 * length) // 3, + max(0, length - 4), + max(0, length - 3), + max(0, length - 2), + length - 1, + } + return tuple(int(tokens[position].item()) for position in sorted(positions)) + + def _sample_stream_key(exp: Experience) -> tuple[str, Any]: info = exp.info or {} sample_index = info.get("sample_index") From 6b7f4dd13b790c8964b7515240f5d27467e26730 Mon Sep 17 00:00:00 2001 From: "panxuchen.pxc" Date: Tue, 30 Jun 2026 16:43:19 +0800 Subject: [PATCH 42/84] fix sub prefix --- tests/common/recording_recorder_test.py | 31 +++++++++++++++++++++++ trinity/common/models/recording/merger.py | 14 +++++++--- 2 files changed, 41 insertions(+), 4 deletions(-) diff --git a/tests/common/recording_recorder_test.py b/tests/common/recording_recorder_test.py index fa945673d38..05820aee53e 100644 --- a/tests/common/recording_recorder_test.py +++ b/tests/common/recording_recorder_test.py @@ -263,6 +263,37 @@ async def test_multi_head_merge_uses_longest_matching_prefix(self): self.assertEqual(merged.info["merged_eid_suffixes"], ["req-short", "req-long", "req-final"]) self.assertEqual(merged.info["merged_turn_count"], 3) + async def test_same_prompt_independent_completions_do_not_merge(self): + store = MemoryStore() + recorder = Recorder( + store=store, + build_experiences=lambda *_args, **_kwargs: [], + enabled=True, + ) + record_key = "0/task_a/1" + short_completion = make_turn( + request_id="req-short", + record_key=record_key, + tokens=[10, 11, 20, 21], + prompt_length=2, + logprobs=[-0.2, -0.3], + ) + long_completion = make_turn( + request_id="req-long", + record_key=record_key, + tokens=[10, 11, 20, 21, 22, 23], + prompt_length=2, + logprobs=[-0.4, -0.5, -0.6, -0.7], + ) + + await recorder._safe_append(short_completion) + await recorder._safe_append(long_completion) + + recorded = store.get(record_key) + self.assertEqual(len(recorded), 2) + self.assertEqual({exp.eid.suffix for exp in recorded}, {"req-short", "req-long"}) + self.assertTrue(all("merged_turn_count" not in exp.info for exp in recorded)) + async def test_stale_merge_head_falls_back_to_append(self): store = MemoryStore() recorder = Recorder( diff --git a/trinity/common/models/recording/merger.py b/trinity/common/models/recording/merger.py index 4f328dadee2..3642cb5e7c1 100644 --- a/trinity/common/models/recording/merger.py +++ b/trinity/common/models/recording/merger.py @@ -131,7 +131,7 @@ def find_longest_prefix(self, exp: Experience) -> Optional[Experience]: entry = self._heads_by_sample_id.get(sample_id) if entry is None: continue - if _is_strict_token_prefix(entry.exp.tokens, exp.tokens): + if _is_mergeable_turn_prefix(entry.exp, exp): if best_entry is None or entry.sequence < best_entry.sequence: best_entry = entry if best_entry is not None: @@ -167,7 +167,7 @@ def _find_longest_prefix_experience( continue if not _same_sample_stream(candidate, exp): continue - if _is_strict_token_prefix(candidate.tokens, exp.tokens): + if _is_mergeable_turn_prefix(candidate, exp): best_candidate = candidate best_length = candidate_length return best_candidate @@ -215,11 +215,17 @@ def _is_strict_token_prefix(prefix: torch.Tensor, tokens: torch.Tensor) -> bool: return bool(torch.equal(prefix.detach().cpu(), tokens[:prefix_len].detach().cpu())) +def _is_mergeable_turn_prefix(prefix_exp: Experience, final_exp: Experience) -> bool: + prefix_len = len(prefix_exp.tokens) + if prefix_len > final_exp.prompt_length: + return False + return _is_strict_token_prefix(prefix_exp.tokens, final_exp.tokens) + + def _merge_prefix_experiences(prefix_exp: Experience, final_exp: Experience) -> Experience: prefix_len = len(prefix_exp.tokens) final_prompt_length = final_exp.prompt_length - if final_prompt_length < prefix_len: - final_prompt_length = prefix_len + assert final_prompt_length >= prefix_len gap_len = final_prompt_length - prefix_len final_response_len = len(final_exp.tokens) - final_prompt_length From f3f13e0838ed72a6d75a3bf959cbf607df7f9ad9 Mon Sep 17 00:00:00 2001 From: pxc Date: Tue, 30 Jun 2026 15:16:10 +0800 Subject: [PATCH 43/84] fix recording tests --- tests/common/sglang_test.py | 23 +++++++++++++++-------- tests/common/vllm_test.py | 6 +++--- trinity/common/models/sglang_model.py | 4 ++-- trinity/explorer/explorer.py | 2 +- 4 files changed, 21 insertions(+), 14 deletions(-) diff --git a/tests/common/sglang_test.py b/tests/common/sglang_test.py index a7aa337225e..32886d5d918 100644 --- a/tests/common/sglang_test.py +++ b/tests/common/sglang_test.py @@ -386,16 +386,18 @@ class TestRecording(RayUnittestBaseAsync): per-task ``record_key`` (captured by ``RecordingIdentityMiddleware``), matching vLLM (which sets no api_key auth in recording mode). - ``enable_history`` forces ``enable_return_routed_experts`` in the - Allocator, so this test uses a MoE checkpoint (``get_moe_model_path``). + ``enable_router_replay`` (mirrored to ``enable_return_routed_experts`` by + ``check_and_update``) is on, so this test uses a MoE checkpoint + (``get_moe_model_path``) and asserts routed_experts shapes. """ async def asyncSetUp(self): self.config = get_template_config() self.config.mode = "explore" - # enable_history forces enable_return_routed_experts -> needs a MoE - # model (otherwise routed_experts is absent and the shape asserts below - # would fail). Use a Qwen3-MoE checkpoint. + # enable_router_replay drives enable_return_routed_experts (see + # ``config_validator``) -> needs a MoE model (otherwise routed_experts + # is absent and the shape asserts below would fail). Use a Qwen3-MoE + # checkpoint. self.config.model.model_path = get_moe_model_path() self.tokenizer = AutoTokenizer.from_pretrained( self.config.model.model_path, @@ -413,6 +415,11 @@ async def asyncSetUp(self): self.config.explorer.rollout_model.enable_openai_api = True self.config.explorer.rollout_model.enable_history = True self.config.explorer.rollout_model.enable_expert_parallel = True + # enable_router_replay is mirrored to enable_return_routed_experts by + # ``check_and_update`` (config_validator); it is NOT implied by + # enable_history. The routed-experts asserts below require it on, so + # the in-SGLang recorder captures routed_experts on every path. + self.config.algorithm.enable_router_replay = True # Tool-call parsing coverage (qwen3_coder matches the Qwen3.5 chat # template). SGLang enables tool calling via tool_call_parser (no # separate enable_auto_tool_choice flag); enable_auto_tool_choice is @@ -472,7 +479,7 @@ def _assert_recorded_experience(self, exp: Experience, record_key: str): self.assertGreater(len(exp.response_text), 0) def _assert_recorded_routed_experts(self, exp: Experience): - # enable_return_routed_experts is forced on by enable_history. + # enable_router_replay -> enable_return_routed_experts is on for this test. self.assertIsNotNone(exp.routed_experts) re = exp.routed_experts self.assertEqual(re.dtype, torch.uint8) @@ -489,7 +496,7 @@ async def test_record(self): # noqa: C901 # ===== 1. Ray-direct generate (record_key via Authorization bearer) ===== rk_gen = "0/t_gen/1" await self.model_wrapper.generate_async( - ["Hello, world!"], n=1, temperature=1.0, max_tokens=16, record_key=rk_gen + ["Hello, world!"], n=1, temperature=1.0, max_tokens=16, key=rk_gen ) consumed = await self._consume(rk_gen, reward=0.5) self.assertEqual(len(consumed), 1) @@ -502,7 +509,7 @@ async def test_record(self): # noqa: C901 # ===== 2. Ray-direct chat, n=2 (one record-key group, two samples) ===== rk_chat = "0/t_chat/2" chat_exps = await self.model_wrapper.chat_async( - messages, n=2, temperature=1.0, max_tokens=16, record_key=rk_chat + messages, n=2, temperature=1.0, max_tokens=16, key=rk_chat ) self.assertEqual(len(chat_exps), 2) consumed = await self._consume(rk_chat, reward=0.8) diff --git a/tests/common/vllm_test.py b/tests/common/vllm_test.py index 0b3e1522ac7..34f076a6bbd 100644 --- a/tests/common/vllm_test.py +++ b/tests/common/vllm_test.py @@ -1937,7 +1937,7 @@ async def test_record(self): # noqa: C901 # ===== 1. Ray-direct generate (record_key via record_key_ctx) ===== rk_gen = "0/t_gen/1" await self.model_wrapper.generate_async( - ["Hello, world!"], n=1, temperature=1.0, max_tokens=16, record_key=rk_gen + ["Hello, world!"], n=1, temperature=1.0, max_tokens=16, key=rk_gen ) consumed = await self._consume(rk_gen, reward=0.5) self.assertEqual(len(consumed), 1) @@ -1950,7 +1950,7 @@ async def test_record(self): # noqa: C901 # ===== 2. Ray-direct chat, n=2 (one record-key group, two samples) ===== rk_chat = "0/t_chat/2" chat_exps = await self.model_wrapper.chat_async( - messages, n=2, temperature=1.0, max_tokens=16, record_key=rk_chat + messages, n=2, temperature=1.0, max_tokens=16, key=rk_chat ) self.assertEqual(len(chat_exps), 2) consumed = await self._consume(rk_chat, reward=0.8) @@ -1966,7 +1966,7 @@ async def test_record(self): # noqa: C901 self._assert_recorded_experience(exp, rk_chat) self._assert_recorded_routed_experts(exp) - # ===== 3. OpenAI regular (HTTP; record_key = Bearer api_key) ===== + # ===== 3. OpenAI regular (HTTP; key = Bearer api_key) ===== rk_oai = "0/t_oai/3" client = await self._openai_client(rk_oai) model_id = await self._get_model_id(client) diff --git a/trinity/common/models/sglang_model.py b/trinity/common/models/sglang_model.py index 4ed87166f31..7c40f6e67de 100644 --- a/trinity/common/models/sglang_model.py +++ b/trinity/common/models/sglang_model.py @@ -506,8 +506,8 @@ async def logprobs(self, token_ids: List[int], **kwargs) -> torch.Tensor: raise NotImplementedError("SGLangRolloutModel does not support logprobs.") # NOTE: if implemented later, the auxiliary forward must avoid being # recorded. Unlike vLLM, ``skip_recording_ctx`` does NOT cross the HTTP - # hop to the server; instead omit ``record_key`` for that call so the - # server-side recorder skips it (record_key is None -> no record). + # hop to the server; instead omit ``key`` for that call so the + # server-side recorder skips it (key is None -> no record). async def convert_messages_to_experience( self, diff --git a/trinity/explorer/explorer.py b/trinity/explorer/explorer.py index 0188b14ee93..d718bae60af 100644 --- a/trinity/explorer/explorer.py +++ b/trinity/explorer/explorer.py @@ -295,7 +295,6 @@ async def explore(self) -> str: """ while True: try: - self.logger.info(f"Explore step {self.explore_step_num + 1} started.") explore_contionue = await self.explore_step() if not explore_contionue: # TODO: support eval on last checkpoint @@ -343,6 +342,7 @@ async def explore_step(self) -> bool: ) await self._finish_explore_step(step=oldest_step) self.last_monitored_step = oldest_step + self.logger.info(f"Explore step {self.explore_step_num} started.") await self.rollout_coordinator.submit_batch.remote( batch_id=self.explore_step_num, tasks=tasks, From 1939b315c31f47a02475f91d4fcf8d9ffc0f5366 Mon Sep 17 00:00:00 2001 From: pxc Date: Tue, 30 Jun 2026 15:51:24 +0800 Subject: [PATCH 44/84] fix workflow test --- tests/common/vllm_test.py | 1 - tests/explorer/scheduler_test.py | 1 - tests/explorer/workflow_test.py | 170 ++++++++---------------- trinity/common/config_validator.py | 21 ++- trinity/common/models/model.py | 18 --- trinity/common/workflows/workflow.py | 5 +- trinity/explorer/rollout_coordinator.py | 6 - trinity/explorer/scheduler.py | 120 +---------------- trinity/explorer/workflow_runner.py | 52 +++----- 9 files changed, 98 insertions(+), 296 deletions(-) diff --git a/tests/common/vllm_test.py b/tests/common/vllm_test.py index 34f076a6bbd..23c769ca7b9 100644 --- a/tests/common/vllm_test.py +++ b/tests/common/vllm_test.py @@ -1035,7 +1035,6 @@ async def test_logprobs_api(self): ) # test openai api and vllm engine logprobs consistency - await self.model_wrapper.clean_workflow_state() _ = await self.model_client.chat.completions.create( model=self.model_client.model_path, messages=messages, diff --git a/tests/explorer/scheduler_test.py b/tests/explorer/scheduler_test.py index 76bcf5c59da..3b056840bee 100644 --- a/tests/explorer/scheduler_test.py +++ b/tests/explorer/scheduler_test.py @@ -246,7 +246,6 @@ async def run_async(self) -> List[Experience]: ) run_level_exps[-1].metrics = run_level_metrics self.logger.info(f"Setting workflow state to repeat_cnt={i}") - await self.model.set_workflow_state({"repeat_cnt": i}) await asyncio.sleep(1) exps.extend(run_level_exps) return exps diff --git a/tests/explorer/workflow_test.py b/tests/explorer/workflow_test.py index a657001c546..26170ea4af9 100644 --- a/tests/explorer/workflow_test.py +++ b/tests/explorer/workflow_test.py @@ -1,12 +1,12 @@ # -*- coding: utf-8 -*- """Test for the workflow module""" import asyncio +import copy import os import shutil import threading import time import unittest -from collections import defaultdict from dataclasses import dataclass, field from typing import Dict, Optional from unittest import mock @@ -24,7 +24,6 @@ get_template_config, get_unittest_dataset_config, ) -from trinity.common.config import InferenceModelConfig from trinity.common.constants import LOG_DIR_ENV_VAR, LOG_LEVEL_ENV_VAR from trinity.common.experience import EID, Experience from trinity.common.models.allocator import Allocator @@ -494,48 +493,6 @@ def tearDown(self): ray.shutdown(_exiting_interpreter=True) -class StateRecordingWorkflow(Workflow): - is_async: bool = True - - def __init__(self, *, task, model: ModelWrapper, auxiliary_models): - super().__init__(task=task, model=model, auxiliary_models=auxiliary_models) - self.wait_time = task.workflow_args.get("wait_time", 1) - - async def run_async(self): - for i in range(self.wait_time): - await self.model.set_workflow_state({"step": i}) - await asyncio.sleep(1) - return [Experience(tokens=Tensor([0, 1, 2]), prompt_length=1, reward=1.0)] - - -class TestWorkflowStateRecording(unittest.IsolatedAsyncioTestCase): - async def test_workflow_state_recording(self): - model = MagicMock() - model_wrapper = ModelWrapper(model, config=InferenceModelConfig(model_path="dummy_model")) - - task = Task( - workflow=StateRecordingWorkflow, - repeat_times=3, - raw_task={}, - workflow_args={"wait_time": 3}, - ) - workflow = task.to_workflow(model_wrapper) - - async def monitor_routine(): - old_state = {} - count = 0 - for i in range(20): - await asyncio.sleep(0.2) - new_state = await model_wrapper.get_workflow_state() - if new_state.get("step") != old_state.get("step"): - old_state = new_state - count += 1 - self.assertEqual(count, 3) - return count - - await asyncio.gather(*[monitor_routine(), workflow.run_async()]) - - class TestAgentScopeWorkflowAdapter(unittest.IsolatedAsyncioTestCase): async def test_adapter_v1(self): try: @@ -592,20 +549,31 @@ async def as_judge_func(task, response) -> JudgeOutput: class DummyModelWrapper: def __init__(self, model, **kwargs): - pass + self._api_key = "EMPTY" async def prepare(self): return + def set_api_key(self, api_key: str) -> None: + """Mirror ModelWrapper.set_api_key for the refactored WorkflowBase.""" + self._api_key = api_key + + def clone_with_isolated_state(self) -> "DummyModelWrapper": + """Mirror ModelWrapper.clone_with_isolated_state for the runner's + isolated workflow instances used in async/multi-threading modes.""" + return copy.copy(self) + + async def overwrite_history_experiences_async(self, experiences, key: str) -> None: + """Mirror ModelWrapper.overwrite_history_experiences_async; a no-op for + tests since DummyWorkflow does not record history.""" + return + def get_openai_client(self): return openai.OpenAI(api_key="EMPTY") def get_openai_async_client(self): return openai.AsyncOpenAI(api_key="EMPTY") - async def clean_workflow_state(self): - return - @property async def model_version_async(self): return 0 @@ -792,7 +760,6 @@ async def test_workflow_runner_fail_fast_without_partial_collection(self, concur async def mock_execute_single_run( workflow: Workflow, - model_wrapper: ModelWrapper, ): run_index = int(workflow.task.run_id) if run_index == 0: @@ -837,65 +804,6 @@ async def mock_execute_single_run( status.message, ) - async def test_workflow_runner_get_state(self): - config = get_template_config() - - async def mock_get_api_server_url_remote(): - return None - - async def mock_get_model_version_remote(): - return 1 - - async def mock_get_api_key_remote(): - return "dummy_api_key" - - async def mock_get_model_config_remote(): - return InferenceModelConfig(model_path="dummy_model") - - model = MagicMock() - model.get_api_server_url.remote = MagicMock(side_effect=mock_get_api_server_url_remote) - model.get_model_version.remote = MagicMock(side_effect=mock_get_model_version_remote) - model.get_api_key.remote = MagicMock(side_effect=mock_get_api_key_remote) - model.get_model_config.remote = MagicMock(side_effect=mock_get_model_config_remote) - - with patch_runner_models( - ModelWrapper(model, config=InferenceModelConfig(model_path="dummy_model")) - ): - runner = WorkflowRunner( - config, - rollout_model_id=0, - runner_id=1, - ) - await runner.prepare() - task = Task( - workflow=StateRecordingWorkflow, - raw_task={}, - workflow_args={"wait_time": 2}, - batch_id=1, - task_id=2, - ) - - async def monitor_routine(): - state_history = defaultdict(set) - count = 0 - for i in range(20): - await asyncio.sleep(0.4) - new_state = await runner.get_runner_state() - for k, v in new_state.items(): - state_history[k].add(v) - self.assertEqual(len(state_history["model_version"]), 1) - self.assertEqual(len(state_history["workflow_id"]), 3) - self.assertEqual(len(state_history["begin_time"]), 3) - self.assertEqual(len(state_history["step"]), 2) - return count - - await asyncio.gather( - *[ - monitor_routine(), - runner.run_task(task, repeat_times=3, run_id_base=0), - ] - ) - async def test_workflow_with_openai(self): config = get_template_config() config.mode = "explore" @@ -917,23 +825,33 @@ async def test_workflow_with_openai(self): workflow=APIWorkflow, raw_task={"raise_except": True}, repeat_times=2, + batch_id="openai_test", + task_id=0, ), Task( workflow=APIWorkflow, raw_task={}, repeat_times=2, + batch_id="openai_test", + task_id=1, ), ] status = await runner.run_task(tasks[0], repeat_times=2, run_id_base=0) self.assertEqual(status.ok, False) + # The run raised after the chat call, so the partial experience recorded + # under the last run's key persists (execute/overwrite is never reached). exps = runner.model_wrapper.extract_experience_from_history(clear_history=False) self.assertEqual(len(exps), 1) status = await runner.run_task(tasks[1], repeat_times=2, run_id_base=0) self.assertEqual(status.ok, True) self.assertEqual(status.completed_runs, 2) + # A successful run extracts the recorded history (clearing it) and then + # `Workflow.execute` overwrites the final experiences back under the key, + # so the last run's key still holds one experience (drained later by the + # coordinator, not by run_task). exps = runner.model_wrapper.extract_experience_from_history(clear_history=False) - self.assertEqual(len(exps), 0) + self.assertEqual(len(exps), 1) self.assertEqual(len(rollout_model), 1) await rollout_model[0].shutdown() @@ -950,20 +868,26 @@ def __init__(self, model: ModelWrapper, task: Task, auxiliary_models=None): async def run_async(self): assert self.task.raw_task is not None - _ = await self.model.chat_async([{"role": "user", "content": self.task.raw_task["text"]}]) + text = self.task.raw_task["text"] + # Both calls opt into recording under the run's record key + # (enable_recording=True is required for chat_async to stamp the key; + # otherwise the engine recorder skips the turn entirely). Distinct prompts + # guarantee the two recorded experiences never form a token-prefix chain, + # so the prefix merger leaves them as two separate experiences. + _ = await self.model.chat_async([{"role": "user", "content": text}], enable_recording=True) await asyncio.sleep(1.0) _ = await self.client.chat.completions.create( model=self.client.model_path, - messages=[{"role": "user", "content": self.task.raw_task["text"]}], + messages=[{"role": "user", "content": "What is the result of one plus one?"}], ) history_exps = self.model.extract_experience_from_history() - assert len(history_exps) == 2 - assert history_exps[0].prompt_length == history_exps[1].prompt_length - prompt_length = history_exps[0].prompt_length - assert ( - history_exps[0].tokens[:prompt_length].shape - == history_exps[1].tokens[:prompt_length].shape + assert len(history_exps) == 2, "Expected 2 experiences from history, got {}".format( + len(history_exps) ) + for exp in history_exps: + assert exp.prompt_length > 0, "Expected a positive prompt length, got {}".format( + exp.prompt_length + ) self.logger.debug("[DEBUG MESSAGE]") self.logger.info("[INFO MESSAGE]") self.logger.warning("[WARNING MESSAGE]") @@ -1045,22 +969,34 @@ async def test_concurrent_workflow_runner(self): workflow=ConcurrentTestWorkflow, repeat_times=4, raw_task={"text": "Hello, world!"}, + batch_id="concurrent", + task_id=0, ) + # Each run_task call uses a distinct batch_id so the record keys + # (//) never collide across calls on the shared + # rollout-model store. `Workflow.execute` overwrites the final experiences + # back under each key, so reusing a key would let a later call observe the + # previous call's leftovers and break the per-run `assert len==2`. # warmup + task.batch_id = "concurrent_async_warmup" async_status = await async_runner.run_task.remote(task, repeat_times=2, run_id_base=0) st = time.time() + task.batch_id = "concurrent_async" async_status = await async_runner.run_task.remote(task, repeat_times=4, run_id_base=0) async_runtime = time.time() - st # warmup + task.batch_id = "concurrent_thread_warmup" thread_status = await thread_runner.run_task.remote(task, repeat_times=1, run_id_base=0) st = time.time() + task.batch_id = "concurrent_thread" thread_status = await thread_runner.run_task.remote(task, repeat_times=4, run_id_base=0) thread_runtime = time.time() - st st = time.time() + task.batch_id = "concurrent_sequential" sequential_status = await sequential_runner.run_task.remote( task, repeat_times=4, run_id_base=0 ) diff --git a/trinity/common/config_validator.py b/trinity/common/config_validator.py index 92be9f6d34b..9a5558b45bc 100644 --- a/trinity/common/config_validator.py +++ b/trinity/common/config_validator.py @@ -708,15 +708,30 @@ def validate(self, config: Config) -> None: # noqa: C901 if config.explorer.rollout_model.engine_type.startswith("vllm") or ( config.explorer.rollout_model.engine_type == "sglang" ): - # enable_history is the single switch for engine-side recording. - if config.explorer.rollout_model.enable_history: - config.explorer.rollout_model.enable_openai_api = True + # enable_history is the single switch for engine-side recording and is + # mandatory for the rollout model: the Workflow.execute overwrite path + # and the Scheduler drain both rely on experiences being captured into + # the in-process store keyed by the recording identity. + if not config.explorer.rollout_model.enable_history: + config.explorer.rollout_model.enable_history = True + self.logger.warning( + "`explorer.rollout_model.enable_history` is required for the rollout " + "model's recording flow; force-set to True." + ) + # enable_history implies the OpenAI API server (the recording runner). + config.explorer.rollout_model.enable_openai_api = True self._validate_inference_parallel_config(config.explorer.rollout_model, "rollout_model") # auxiliary models for aux_model in config.explorer.auxiliary_models: if not aux_model.model_path: raise ValueError("auxiliary model's model_path is required.") aux_model.ray_namespace = config.ray_namespace + # auxiliary models must not record history; only the rollout model does. + if aux_model.enable_history: + self.logger.warning( + "`enable_history` is not supported on auxiliary models and is " + "force-set to False." + ) aux_model.enable_history = False aux_model.enable_openai_api = True for args in model_args: diff --git a/trinity/common/models/model.py b/trinity/common/models/model.py index 1afc444e35e..7eade887ab1 100644 --- a/trinity/common/models/model.py +++ b/trinity/common/models/model.py @@ -504,9 +504,7 @@ def __init__( self.enable_lora = config.enable_lora self.enable_history = config.enable_history self.status = RunningStatus.RUNNING - self.workflow_state: Dict = {} self.request_count = 0 - self.state_lock = asyncio.Lock() async def prepare(self) -> None: """Prepare some necessary information for the model before inference.""" @@ -944,17 +942,6 @@ async def delete_experience_records_async(self, prefix: str) -> None: raise ValueError("Recording delete requires an inference model actor.") await self.model.delete_experience_records.remote(prefix=prefix) - # Workflow state management methods - async def set_workflow_state(self, state: Dict) -> None: - """Set the state of workflow using the model.""" - async with self.state_lock: - self.workflow_state.update(state) - - async def clean_workflow_state(self) -> None: - """Clean the state of workflow using the model.""" - async with self.state_lock: - self.workflow_state = {} - async def shutdown(self) -> None: """Shutdown all underlying model actors cleanly.""" try: @@ -964,11 +951,6 @@ async def shutdown(self) -> None: f"Error during model {self.config.model_path}[{self.config.engine_id}:{self.config.node_rank}] shutdown: {e}" ) - async def get_workflow_state(self) -> Dict: - """Get the state of workflow using the model.""" - async with self.state_lock: - return self.workflow_state.copy() - def clone_with_isolated_state(self) -> "ModelWrapper": """Clone the current ModelWrapper with isolated state.""" new_wrapper = copy.copy(self) diff --git a/trinity/common/workflows/workflow.py b/trinity/common/workflows/workflow.py index f4d1179912b..1bb0bed1a65 100644 --- a/trinity/common/workflows/workflow.py +++ b/trinity/common/workflows/workflow.py @@ -432,6 +432,9 @@ class WorkflowWithRecording(WorkflowBase): the recorded experience data with the reward and optional info. """ + can_reset: bool = True + is_async: bool = True + def __init__( self, *, @@ -459,7 +462,7 @@ def reset(self, task: Task): """Reset the workflow.""" raise NotImplementedError - async def run(self) -> Metrics: + async def run_async(self) -> Metrics: """Run workflow asynchronously and return metrics for the completed run.""" raise NotImplementedError diff --git a/trinity/explorer/rollout_coordinator.py b/trinity/explorer/rollout_coordinator.py index ad439b215d8..8891632caa6 100644 --- a/trinity/explorer/rollout_coordinator.py +++ b/trinity/explorer/rollout_coordinator.py @@ -68,10 +68,6 @@ def __init__( # recording residual cleanup. self._rollout_actors: Dict[int, ActorHandle] = {} - def _enable_history_recording(self) -> bool: - """Whether the recording-consume path is active for train batches.""" - return bool(self.config.explorer.rollout_model.enable_history) - def _resolve_rollout_actors(self) -> Dict[int, ActorHandle]: """Resolve each rollout engine's actor handle via named Ray actors. @@ -323,8 +319,6 @@ async def _finalize_train_batch( async def _discard_recorded_experiences(self, prefix: str) -> None: """Delete recorded experiences matching a prefix from all rollout ranks.""" - if not self._enable_history_recording(): - return actors = self._resolve_rollout_actors() results = await asyncio.gather( *[actor.delete_experience_records.remote(prefix=prefix) for actor in actors.values()], diff --git a/trinity/explorer/scheduler.py b/trinity/explorer/scheduler.py index c0f2237c34d..72e68423b9f 100644 --- a/trinity/explorer/scheduler.py +++ b/trinity/explorer/scheduler.py @@ -84,7 +84,6 @@ def __init__( self.timeout = config.explorer.max_timeout self.namespace = config.ray_namespace self.runner = self._create_runner() - self.state = {} self.rollout_actor = rollout_actor def _create_runner(self): @@ -106,11 +105,6 @@ def _create_runner(self): ) ) - async def update_state(self) -> None: - """Get the runner state.""" - self.state = await self.runner.get_runner_state.remote() - self.state["running_time"] = time.time() - self.state.get("begin_time", time.time()) - async def prepare(self): await self.runner.prepare.remote() @@ -161,7 +155,7 @@ async def run_with_retry( self, task: TaskWrapper, repeat_times: int, - run_base: int, + run_id_base: int, timeout: float, collect_partial_runs: bool, ) -> Tuple[Status, List[bytes], int, float]: @@ -196,7 +190,7 @@ async def run_with_retry( run_task_ref = self.runner.run_task.remote( task=task2run, repeat_times=repeat_times, - run_base=run_base, + run_id_base=run_id_base, collect_partial_runs=collect_partial_runs, ) status = await asyncio.wait_for( @@ -326,7 +320,6 @@ def __init__( self.background_tasks: set[asyncio.Task] = set() self.scheduler_task: Optional[asyncio.Task] = None - self.monitor_task: Optional[asyncio.Task] = None self.total_running_time = 0.0 self.total_completed_steps = 0 @@ -383,24 +376,6 @@ async def _scheduler_loop(self) -> None: await asyncio.sleep(0.1) self.logger.info("Scheduler loop stopped.") - async def _monitor_runner_state_loop(self) -> None: - interval = self.config.explorer.runner_state_report_interval - if interval <= 0: - self.logger.info("Runner state monitoring loop disabled.") - return - - self.logger.info("Runner state monitoring loop started.") - while self.running: - try: - await asyncio.gather(*[runner.update_state() for runner in self.runners.values()]) - self.print_all_state() - except Exception: - self.logger.error( - f"Error in runner state monitoring loop:\n{traceback.format_exc()}" - ) - await asyncio.sleep(interval) - self.logger.info("Runner state monitoring loop stopped.") - async def _schedule_pending_tasks(self) -> None: if not self.idle_runners: return @@ -410,13 +385,13 @@ async def _schedule_pending_tasks(self) -> None: task_queue = self.pending_tasks[batch_id] while task_queue and self.idle_runners: - task, repeat_times, run_base = task_queue.pop() + task, repeat_times, run_id_base = task_queue.pop() runner_id = self.idle_runners.pop() future = asyncio.create_task( self.runners[runner_id].run_with_retry( task, repeat_times=repeat_times, - run_base=run_base, + run_id_base=run_id_base, timeout=self.dynamic_timeout(), collect_partial_runs=self.config.explorer.over_rollout.return_partial_tasks, ) @@ -578,7 +553,6 @@ async def _create_limited(i: int) -> None: self.scheduler_task = asyncio.create_task(self._scheduler_loop()) ready_refs = [runner.runner.__ray_ready__.remote() for runner in self.runners.values()] await asyncio.gather(*ready_refs) - self.monitor_task = asyncio.create_task(self._monitor_runner_state_loop()) self.logger.info(f"Starting Scheduler with {self.runner_num} runners") async def stop(self) -> None: @@ -603,12 +577,6 @@ async def stop(self) -> None: await self.scheduler_task except asyncio.CancelledError: pass - if self.monitor_task: - self.monitor_task.cancel() - try: - await self.monitor_task - except asyncio.CancelledError: - pass self.logger.info("Scheduler stopped") def schedule(self, tasks: List[Task], batch_id: Union[int, str]) -> None: @@ -638,9 +606,9 @@ def _split_and_submit_tasks(self, tasks: List[Task], batch_id: Union[int, str]) self.pending_tasks[batch_id].appendleft((task_wrapper, task.repeat_times, 0)) continue sub_tasks = [] - for run_base in range(0, task.repeat_times, self.max_repeat_times): - repeat_times = min(self.max_repeat_times, task.repeat_times - run_base) - sub_tasks.append((task_wrapper, repeat_times, run_base)) + for run_id_base in range(0, task.repeat_times, self.max_repeat_times): + repeat_times = min(self.max_repeat_times, task.repeat_times - run_id_base) + sub_tasks.append((task_wrapper, repeat_times, run_id_base)) task_wrapper.sub_task_num = len(sub_tasks) self.pending_tasks[batch_id].extendleft(sub_tasks) @@ -860,77 +828,3 @@ def has_step(self, batch_id: Union[int, str]) -> bool: or batch_id in self.pending_tasks or batch_id in self.running_tasks ) - - def get_key_state(self, key: str) -> Dict: - """Get the scheduler state. - - Args: - key (`str`): The key of the state to get. - - Returns: - `Dict`: A dictionary of runner ids to their state for the given key. - """ - result = {} - for runner in self.runners.values(): - runner_state = runner.state - if runner_state and key in runner_state: - result[runner.runner_id] = runner_state[key] - return result - - def get_runner_state(self, runner_id: int) -> Dict: - """Get the scheduler state. - - Args: - runner_id (`int`): The id of the runner. - - Returns: - `Dict`: The state of the runner. - """ - runner = self.runners.get(runner_id, None) - if runner: - return runner.state - else: - return {} - - def get_all_state(self) -> Dict: - """Get all runners' state. - - Returns: - `Dict`: The state of all runners. - """ - result = {} - for runner in self.runners.values(): - runner_state = runner.state - if runner_state: - result[runner.runner_id] = runner_state - return result - - def print_all_state(self) -> None: - """Print all runners' state in a clear, aligned table format.""" - all_keys = set() - for runner in self.runners.values(): - runner_state = runner.state - if runner_state: - all_keys.update(runner_state.keys()) - all_keys = sorted(all_keys) - # Prepare header - header = ["runner_id"] + all_keys # type: ignore [operator] - # Prepare rows - rows = [] - for runner in self.runners.values(): - runner_state = runner.state or {} - row = [str(runner.runner_id)] - for key in all_keys: - value = runner_state.get(key, "-") - row.append(str(value)) - rows.append(row) - # Calculate column widths - col_widths = [max(len(str(x)) for x in col) for col in zip(header, *rows)] - # Print header - header_line = " | ".join(str(h).ljust(w) for h, w in zip(header, col_widths)) - self.logger.info(header_line) - self.logger.info("-+-".join("-" * w for w in col_widths)) - # Print each row - for row in rows: - line = " | ".join(str(cell).ljust(w) for cell, w in zip(row, col_widths)) - self.logger.info(line) diff --git a/trinity/explorer/workflow_runner.py b/trinity/explorer/workflow_runner.py index a15f207cd16..23b6b212b29 100644 --- a/trinity/explorer/workflow_runner.py +++ b/trinity/explorer/workflow_runner.py @@ -2,10 +2,11 @@ """The Workflow Runner Module.""" import asyncio +import copy import os import time import traceback -from typing import Dict, List, Optional, Tuple +from typing import Dict, List, Optional from trinity.buffer import get_buffer_reader, get_buffer_writer from trinity.common.config import Config, StorageConfig @@ -89,7 +90,7 @@ def _create_workflow_instance(self, task: Task) -> Workflow: if ( self.workflow_instance is None or not self.workflow_instance.__class__ == task.workflow - or not self.workflow_instance.resettable + or not getattr(self.workflow_instance.__class__, "can_reset", True) ): # Pass ModelWrapper directly; Workflow.__init__ will get OpenAI clients automatically self.workflow_instance = task.to_workflow( @@ -109,18 +110,17 @@ async def _run_workflow(self, workflow_instance: Workflow) -> Status: ) return status - def _create_isolated_workflow_instance( - self, task: Task, run_id: int - ) -> Tuple[Workflow, ModelWrapper]: + def _create_isolated_workflow_instance(self, task: Task, run_id: int) -> Workflow: model_wrapper = self.model_wrapper.clone_with_isolated_state() - # only a shallow copy is enough - task = task.copy() + # only a shallow copy is enough; use copy.copy so the result stays a Task + # (Task inherits dict, so task.copy() would return a plain dict) + task = copy.copy(task) task.run_id = run_id wf = task.to_workflow( model_wrapper, self.auxiliary_model_wrappers, ) - return wf, model_wrapper + return wf def _build_status( self, @@ -185,11 +185,8 @@ async def _run_parallel_runs( ) -> Status: async def run_single(i: int) -> Status: run_id = run_id_base + i - workflow, model_wrapper = self._create_isolated_workflow_instance(task, run_id) - return await self._execute_single_run( - workflow=workflow, - model_wrapper=model_wrapper, - ) + workflow = self._create_isolated_workflow_instance(task, run_id) + return await self._execute_single_run(workflow=workflow) if collect_partial_runs: if use_threads: @@ -238,10 +235,8 @@ async def run_single(i: int) -> Status: async def _execute_single_run( self, workflow: Workflow, - model_wrapper: ModelWrapper, ) -> Status: st = time.time() - await model_wrapper.clean_workflow_state() self.runner_state["terminate_time"] = None self.runner_state["begin_time"] = st try: @@ -278,13 +273,8 @@ async def _run_task( workflow_instance = self._create_workflow_instance(task) workflow_instance.set_repeat_times(repeat_times, run_id_base) st = time.time() - await self.model_wrapper.clean_workflow_state() - self.runner_state["workflow_id"] = task.api_key - self.runner_state["terminate_time"] = None - self.runner_state["begin_time"] = st status = await self._run_workflow(workflow_instance) et = time.time() - self.runner_state["terminate_time"] = et run_metrics = [dict(metric) for metric in status.metrics] for metric in run_metrics: metric["time/run_execution"] = et - st @@ -316,10 +306,7 @@ async def _sequential_run( for i in range(repeat_times): task.run_id = run_id_base + i workflow = self._create_workflow_instance(task) - result = await self._execute_single_run( - workflow=workflow, - model_wrapper=self.model_wrapper, - ) + result = await self._execute_single_run(workflow=workflow) results.append(result) if collect_partial_runs: continue @@ -357,12 +344,6 @@ async def _multi_threading_run( use_threads=True, ) - async def get_runner_state(self) -> Dict: - """Get the runner state.""" - runner_state = self.runner_state.copy() - runner_state.update(await self.model_wrapper.get_workflow_state()) - return runner_state - async def run_task( self, task: Task, @@ -459,12 +440,11 @@ async def debug(self) -> None: with VizTracer(output_file=self.output_profiling_file): status = await self.run_task(task=task, repeat_times=1, run_id_base=0) experiences = [] - if self.config.explorer.rollout_model.enable_history: - try: - payload = await self.model_wrapper.drain_experience_records_bytes_async("debug") - experiences = Experience.deserialize_many(payload) if payload else [] - except Exception: - experiences = [] + try: + payload = await self.model_wrapper.drain_experience_records_bytes_async("debug") + experiences = Experience.deserialize_many(payload) if payload else [] + except Exception: + experiences = [] if not status.ok and not experiences: try: experiences = self.model_wrapper.extract_experience_from_history() From 570b9273515e4e867418fbd68176729fe69bbcaa Mon Sep 17 00:00:00 2001 From: "panxuchen.pxc" Date: Tue, 30 Jun 2026 17:20:35 +0800 Subject: [PATCH 45/84] fix eval batch_id --- .codex/AGENTS.md | 5 +++++ tests/common/recording_store_test.py | 18 ++++++++++++++++++ trinity/buffer/store/memory_store.py | 14 ++++++++++++-- 3 files changed, 35 insertions(+), 2 deletions(-) diff --git a/.codex/AGENTS.md b/.codex/AGENTS.md index 3233c538ab7..6d62f396713 100644 --- a/.codex/AGENTS.md +++ b/.codex/AGENTS.md @@ -12,3 +12,8 @@ ## Repository Convention Treat `docs/agents/` as the single source of truth for agent-facing process and navigation documents. + +## Local Python Environment + +- Always use the repository virtual environment for Python commands: `.venv/bin/python`. +- Run Python tools through that interpreter, for example `.venv/bin/python -m pytest ...`, instead of relying on globally installed commands. diff --git a/tests/common/recording_store_test.py b/tests/common/recording_store_test.py index e27b2e9ab38..126dae8fb86 100644 --- a/tests/common/recording_store_test.py +++ b/tests/common/recording_store_test.py @@ -80,6 +80,24 @@ async def test_delete_request_experience_keeps_other_experiences(self): remaining = store.get(record_key) self.assertEqual(remaining, [exp_b]) + async def test_eval_batch_record_key_allows_slash_in_batch_id(self): + store = MemoryStore() + record_key = "0/eval_short/1/0" + exp = make_exp("req_eval", record_key) + + batch, task, run = parse_record_key(record_key) + self.assertEqual(batch, "0/eval_short") + self.assertEqual(task, "1") + self.assertEqual(run, 0) + + store.add(get_record_key(exp), [exp]) + + self.assertEqual(store.get(record_key), [exp]) + self.assertEqual(store.get("0/eval_short"), [exp]) + self.assertEqual(store.get("0/eval_short/1"), [exp]) + self.assertEqual(store.remove("0/eval_short/1"), [exp]) + self.assertEqual(store.keys(), []) + def _find_request(store: MemoryStore, record_key: str, request_id: str) -> Experience | None: for exp in store.get(record_key): diff --git a/trinity/buffer/store/memory_store.py b/trinity/buffer/store/memory_store.py index dc390c898ad..241bfaad2cc 100644 --- a/trinity/buffer/store/memory_store.py +++ b/trinity/buffer/store/memory_store.py @@ -8,8 +8,12 @@ def parse_record_key(key: str) -> tuple[str, str, int]: - """Parse a complete ``//`` store key.""" - parts = key.split("/") + """Parse a complete ``//`` store key. + + ``batch_id`` may itself contain ``/`` for eval batches, for example + ``0/eval_short/1/0`` means batch ``0/eval_short``, task ``1`` and run ``0``. + """ + parts = key.rsplit("/", 2) if len(parts) != 3 or any(part == "" for part in parts): raise ValueError( f"Store key must be complete '//', got '{key}'." @@ -157,12 +161,18 @@ def _matching_keys(self, key: str) -> list[str]: return list(self._records.keys()) if key in self._records: return [key] + if key in self._batch_keys: + return list(self._batch_keys[key]) parts = key.split("/") if len(parts) == 1 and parts[0] != "": return list(self._batch_keys.get(parts[0], ())) if len(parts) == 2 and parts[0] != "" and parts[1] != "": return list(self._task_keys.get((parts[0], parts[1]), ())) + + batch, sep, task = key.rpartition("/") + if sep and batch and task: + return list(self._task_keys.get((batch, task), ())) return [] def _drop_key(self, key: str) -> list[Experience]: From d23db4dc4446c213b29e0eb525230e6169e467f6 Mon Sep 17 00:00:00 2001 From: "panxuchen.pxc" Date: Tue, 30 Jun 2026 17:55:00 +0800 Subject: [PATCH 46/84] add new workflow tests --- tests/explorer/workflow_test.py | 76 +++++++++++++++++++++++++++- trinity/common/workflows/workflow.py | 2 +- 2 files changed, 76 insertions(+), 2 deletions(-) diff --git a/tests/explorer/workflow_test.py b/tests/explorer/workflow_test.py index 26170ea4af9..135cd749964 100644 --- a/tests/explorer/workflow_test.py +++ b/tests/explorer/workflow_test.py @@ -24,6 +24,7 @@ get_template_config, get_unittest_dataset_config, ) +from trinity.buffer.store import get_record_key from trinity.common.constants import LOG_DIR_ENV_VAR, LOG_LEVEL_ENV_VAR from trinity.common.experience import EID, Experience from trinity.common.models.allocator import Allocator @@ -31,7 +32,13 @@ from trinity.common.workflows import WORKFLOWS, Workflow from trinity.common.workflows.customized_math_workflows import MathBoxedWorkflow from trinity.common.workflows.eval_workflow import MathEvalWorkflow -from trinity.common.workflows.workflow import MathWorkflow, MultiTurnWorkflow, Task +from trinity.common.workflows.workflow import ( + MathWorkflow, + Metrics, + MultiTurnWorkflow, + Task, + WorkflowWithRecording, +) from trinity.explorer.workflow_runner import Status, WorkflowRunner @@ -894,6 +901,42 @@ async def run_async(self): return history_exps +class ConcurrentRecordingWorkflow(WorkflowWithRecording): + def __init__(self, model: ModelWrapper, task: Task, auxiliary_models=None): + super().__init__(task=task, model=model, auxiliary_models=auxiliary_models) + self.client = openai.AsyncOpenAI(base_url=f"{self.base_url}/v1", api_key=self.api_key) + self.model_name = self.model.model_name + + def reset(self, task: Task): + self.task = task + self.model.set_api_key(task.api_key) + self.client.api_key = task.api_key + + async def _chat(self, messages): + return await self.client.chat.completions.create( + model=self.model_name, + messages=messages, + temperature=0.0, + max_tokens=16, + ) + + async def run_async(self) -> Metrics: + prefix_messages = [{"role": "user", "content": "Reply with the word alpha only."}] + first = await self._chat(prefix_messages) + first_text = first.choices[0].message.content or "" + + merged_messages = [ + *prefix_messages, + {"role": "assistant", "content": first_text}, + {"role": "user", "content": "Now reply with the word beta only."}, + ] + await self._chat(merged_messages) + + await self._chat([{"role": "user", "content": "This is an unrelated single-turn branch."}]) + await self.update_reward(0.75, info={"source": "workflow_with_recording"}) + return {"recording_workflow/updated_reward": 1.0} + + class TestConcurrentWorkflowRunner(RayUnittestBaseAsync): def setUp(self) -> None: config = get_template_config() @@ -1012,6 +1055,37 @@ async def test_concurrent_workflow_runner(self): self.assertLessEqual(async_runtime * 2, sequential_runtime) self.assertLessEqual(thread_runtime * 2, sequential_runtime) + recording_task = Task( + workflow=ConcurrentRecordingWorkflow, + repeat_times=1, + raw_task={}, + batch_id="concurrent_recording", + task_id=0, + ) + recording_status = await sequential_runner.run_task.remote( + recording_task, repeat_times=1, run_id_base=0 + ) + self.assertTrue(recording_status.ok) + self.assertEqual(recording_status.completed_runs, 1) + self.assertEqual(recording_status.successful_ids, ["concurrent_recording/0/0"]) + self.assertEqual(recording_status.metrics[0]["recording_workflow/updated_reward"], 1.0) + + recording_exps = rollout_model[0].extract_experience_from_history( + key="concurrent_recording/0/0" + ) + self.assertEqual(len(recording_exps), 2) + for exp in recording_exps: + self.assertEqual(get_record_key(exp), "concurrent_recording/0/0") + self.assertEqual(exp.reward, 0.75) + self.assertEqual(exp.info["source"], "workflow_with_recording") + + merged_exps = [exp for exp in recording_exps if "merged_turn_count" in (exp.info or {})] + branch_exps = [exp for exp in recording_exps if "merged_turn_count" not in (exp.info or {})] + self.assertEqual(len(merged_exps), 1) + self.assertEqual(len(branch_exps), 1) + self.assertEqual(merged_exps[0].info["merged_turn_count"], 2) + self.assertEqual(len(merged_exps[0].info["merged_eid_suffixes"]), 2) + # check log files sequential_log_path = os.path.join(self.config.log.save_dir, "explorer_runner_0.log") async_log_path = os.path.join(self.config.log.save_dir, "explorer_runner_1.log") diff --git a/trinity/common/workflows/workflow.py b/trinity/common/workflows/workflow.py index 1bb0bed1a65..1b9882551da 100644 --- a/trinity/common/workflows/workflow.py +++ b/trinity/common/workflows/workflow.py @@ -468,7 +468,7 @@ async def run_async(self) -> Metrics: async def execute(self) -> Status: """Execute the workflow and normalize the user return value to Status.""" - result = await self.run() + result = await self.run_async() return self._to_status(result) def _to_status(self, result: Metrics) -> Status: From ca402f7d79453a29a81e9fc3e5cd49ae2de84882 Mon Sep 17 00:00:00 2001 From: "panxuchen.pxc" Date: Tue, 30 Jun 2026 18:11:54 +0800 Subject: [PATCH 47/84] fix workflow tests --- tests/explorer/scheduler_test.py | 43 +++++++++++++++++++++++++++++++- 1 file changed, 42 insertions(+), 1 deletion(-) diff --git a/tests/explorer/scheduler_test.py b/tests/explorer/scheduler_test.py index 3b056840bee..8e8665af9ff 100644 --- a/tests/explorer/scheduler_test.py +++ b/tests/explorer/scheduler_test.py @@ -204,7 +204,12 @@ async def run_async(self): for step in range(self.step_num): run_level_exps.append( Experience( - eid=EID(run=i + self.run_id_base, step=step), + eid=EID( + batch=self.task.batch_id, + task=self.task.task_id, + run=i + self.run_id_base, + step=step, + ), tokens=torch.zeros(5), prompt_length=2, prompt_text="success", @@ -277,6 +282,7 @@ def __init__(self): from trinity.common.config import InferenceModelConfig super().__init__(InferenceModelConfig(model_path="dummy_model")) + self._history_payloads: Dict[str, bytes] = {} def sync_model_weights(self, model_version, sync_method, timeout): return True @@ -303,6 +309,38 @@ async def init_process_group( def get_api_server_url(self) -> Optional[str]: return None + async def overwrite_history_experiences(self, key: str, payload: bytes) -> None: + self._history_payloads[key] = payload + + async def drain_experience_records_bytes(self, prefix: str) -> bytes: + keys = self._matching_history_keys(prefix) + exps = [] + for key in keys: + exps.extend(Experience.deserialize_many(self._history_payloads.pop(key))) + return Experience.serialize_many(exps) + + async def delete_experience_records(self, prefix: str) -> None: + for key in self._matching_history_keys(prefix): + self._history_payloads.pop(key, None) + + async def extract_experience_from_history( + self, key: str, clear_history: bool = True + ) -> List[Experience]: + payload = self._history_payloads.get(key) + if payload is None: + return [] + if clear_history: + self._history_payloads.pop(key, None) + return Experience.deserialize_many(payload) + + def _matching_history_keys(self, prefix: str) -> List[str]: + if prefix == "": + return list(self._history_payloads) + if prefix in self._history_payloads: + return [prefix] + prefix_with_sep = f"{prefix}/" + return [key for key in self._history_payloads if key.startswith(prefix_with_sep)] + async def chat(self, messages: List[Dict], lora_request=None, **kwargs) -> Sequence[Experience]: prompt_length = sum(len(msg["content"]) for msg in messages) return [ @@ -409,6 +447,9 @@ def _assign_test_namespace(config) -> None: def _configure_dummy_models(config) -> None: + config.explorer.rollout_model.engine_type = "tinker" + config.explorer.rollout_model.enable_openai_api = False + config.explorer.rollout_model.enable_history = True for auxiliary_config in config.explorer.auxiliary_models: auxiliary_config.enable_openai_api = True From d6c00c8afd0829fa9e0757df1588de0a7b187fe2 Mon Sep 17 00:00:00 2001 From: "panxuchen.pxc" Date: Tue, 30 Jun 2026 18:38:45 +0800 Subject: [PATCH 48/84] fix workflow tests --- trinity/common/workflows/workflow.py | 10 ++++++++++ trinity/explorer/workflow_runner.py | 5 +++++ 2 files changed, 15 insertions(+) diff --git a/trinity/common/workflows/workflow.py b/trinity/common/workflows/workflow.py index 1b9882551da..5849d2e814f 100644 --- a/trinity/common/workflows/workflow.py +++ b/trinity/common/workflows/workflow.py @@ -174,6 +174,16 @@ def set_repeat_times(self, repeat_times: int, run_id_base: int) -> None: self.repeat_times = repeat_times self.run_id_base = run_id_base + def set_single_run_context(self, run_id_base: int) -> None: + """ + Set the workflow context for a single non-repeat run. + + This only updates runner bookkeeping fields and intentionally avoids + repeat-workflow side effects such as changing rollout_args.n. + """ + self.repeat_times = 1 + self.run_id_base = run_id_base + def run(self) -> List[Experience]: """Run workflow and return a list of experiences.""" raise NotImplementedError diff --git a/trinity/explorer/workflow_runner.py b/trinity/explorer/workflow_runner.py index 23b6b212b29..de6850acb34 100644 --- a/trinity/explorer/workflow_runner.py +++ b/trinity/explorer/workflow_runner.py @@ -99,8 +99,12 @@ def _create_workflow_instance(self, task: Task) -> Workflow: ) else: self.workflow_instance.reset(task) + self._set_single_run_context(self.workflow_instance, task.run_id) return self.workflow_instance + def _set_single_run_context(self, workflow: Workflow, run_id: int) -> None: + workflow.set_single_run_context(run_id) + async def _run_workflow(self, workflow_instance: Workflow) -> Status: status = await workflow_instance.execute() if not isinstance(status, Status): @@ -120,6 +124,7 @@ def _create_isolated_workflow_instance(self, task: Task, run_id: int) -> Workflo model_wrapper, self.auxiliary_model_wrappers, ) + self._set_single_run_context(wf, run_id) return wf def _build_status( From 74bc51e7103338d20cb79394fd4d1d410e8d8e2b Mon Sep 17 00:00:00 2001 From: "panxuchen.pxc" Date: Tue, 30 Jun 2026 18:40:45 +0800 Subject: [PATCH 49/84] clean code --- trinity/common/workflows/workflow.py | 2 +- trinity/explorer/workflow_runner.py | 7 ++----- 2 files changed, 3 insertions(+), 6 deletions(-) diff --git a/trinity/common/workflows/workflow.py b/trinity/common/workflows/workflow.py index 5849d2e814f..feef9982fe3 100644 --- a/trinity/common/workflows/workflow.py +++ b/trinity/common/workflows/workflow.py @@ -53,7 +53,7 @@ class Task(dict): # automatically assigned ids batch_id: Union[int, str] = "" task_id: Union[int, str] = "" - run_id: Union[int, str] = "" + run_id: int = 0 index: dict = field(default_factory=dict) diff --git a/trinity/explorer/workflow_runner.py b/trinity/explorer/workflow_runner.py index de6850acb34..a013c8ffef9 100644 --- a/trinity/explorer/workflow_runner.py +++ b/trinity/explorer/workflow_runner.py @@ -99,12 +99,9 @@ def _create_workflow_instance(self, task: Task) -> Workflow: ) else: self.workflow_instance.reset(task) - self._set_single_run_context(self.workflow_instance, task.run_id) + self.workflow_instance.set_single_run_context(task.run_id) return self.workflow_instance - def _set_single_run_context(self, workflow: Workflow, run_id: int) -> None: - workflow.set_single_run_context(run_id) - async def _run_workflow(self, workflow_instance: Workflow) -> Status: status = await workflow_instance.execute() if not isinstance(status, Status): @@ -124,7 +121,7 @@ def _create_isolated_workflow_instance(self, task: Task, run_id: int) -> Workflo model_wrapper, self.auxiliary_model_wrappers, ) - self._set_single_run_context(wf, run_id) + wf.set_single_run_context(run_id) return wf def _build_status( From 1223285537535e1f1a684fe45523a74e6cf4bab5 Mon Sep 17 00:00:00 2001 From: "panxuchen.pxc" Date: Tue, 30 Jun 2026 19:12:03 +0800 Subject: [PATCH 50/84] fix scheduler tests --- tests/explorer/scheduler_test.py | 42 +++++++++++++++++++++++++++----- 1 file changed, 36 insertions(+), 6 deletions(-) diff --git a/tests/explorer/scheduler_test.py b/tests/explorer/scheduler_test.py index 8e8665af9ff..a6de5741240 100644 --- a/tests/explorer/scheduler_test.py +++ b/tests/explorer/scheduler_test.py @@ -58,7 +58,12 @@ def run(self) -> List[Experience]: tokens=torch.zeros(5), prompt_length=2, prompt_text=self.error_type or "success", - eid=EID(run=i + self.run_id_base, step=step), + eid=EID( + batch=self.task.batch_id, + task=self.task.task_id, + run=i + self.run_id_base, + step=step, + ), info={"repeat_times": self.repeat_times}, ) ) @@ -86,7 +91,12 @@ def reset(self, task: Task): def run(self) -> List[Experience]: exps = [ Experience( - eid=EID(run=self.run_id_base, step=step), + eid=EID( + batch=self.task.batch_id, + task=self.task.task_id, + run=self.run_id_base, + step=step, + ), tokens=torch.zeros(5), prompt_length=2, prompt_text="success", @@ -122,7 +132,12 @@ def run(self) -> List[Experience]: return [ Experience( - eid=EID(step=0), + eid=EID( + batch=self.task.batch_id, + task=self.task.task_id, + run=self.run_id_base, + step=0, + ), tokens=torch.zeros(5), prompt_length=2, prompt_text=action, @@ -172,7 +187,12 @@ async def run_async(self) -> List[Experience]: return [ Experience( - eid=EID(step=0), + eid=EID( + batch=self.task.batch_id, + task=self.task.task_id, + run=self.run_id_base, + step=0, + ), tokens=torch.zeros(5), prompt_length=2, prompt_text=action, @@ -243,7 +263,12 @@ async def run_async(self) -> List[Experience]: for step in range(self.step_num): run_level_exps.append( Experience( - eid=EID(run=i + self.run_id_base, step=step), + eid=EID( + batch=self.task.batch_id, + task=self.task.task_id, + run=i + self.run_id_base, + step=step, + ), tokens=torch.zeros(5), prompt_length=2, prompt_text="success", @@ -268,7 +293,12 @@ async def run_async(self) -> List[Experience]: return [ Experience( - eid=EID(run=self.run_id_base, step=0), + eid=EID( + batch=self.task.batch_id, + task=self.task.task_id, + run=self.run_id_base, + step=0, + ), tokens=torch.zeros(5), prompt_length=2, prompt_text="success", From b2708d0b68496cdc719875f6f1cc1dc80c9e048f Mon Sep 17 00:00:00 2001 From: "panxuchen.pxc" Date: Tue, 30 Jun 2026 19:58:03 +0800 Subject: [PATCH 51/84] fix scheduler tests --- tests/explorer/scheduler_test.py | 76 +------------------------------- 1 file changed, 2 insertions(+), 74 deletions(-) diff --git a/tests/explorer/scheduler_test.py b/tests/explorer/scheduler_test.py index a6de5741240..e34dd6dfb47 100644 --- a/tests/explorer/scheduler_test.py +++ b/tests/explorer/scheduler_test.py @@ -22,6 +22,8 @@ @WORKFLOWS.register_module("dummy_workflow") class DummyWorkflow(Workflow): + can_repeat: bool = True + def __init__(self, *, task, model, auxiliary_models): super().__init__(task=task, model=model, auxiliary_models=auxiliary_models) self.step_num = task.workflow_args.get("step_num", 1) @@ -1410,77 +1412,3 @@ def tearDown(self): ray.shutdown() except Exception: pass - - -class TestRunnerStateCollection(unittest.IsolatedAsyncioTestCase): - def setUp(self): - ray.init(ignore_reinit_error=True) - self.config = get_template_config() - _assign_test_namespace(self.config) - _configure_dummy_models(self.config) - self.config.explorer.runner_per_model = 2 - self.config.explorer.runner_state_report_interval = 1 - self.config.explorer.max_repeat_times_per_runner = 2 - self.config.check_and_update() - self.model_actors = _create_named_model_actors(self.config) - - def _create_scheduler(self) -> Scheduler: - return Scheduler(self.config, rollout_actors=_resolve_rollout_actors(self.config)) - - def tearDown(self): - try: - _cleanup_named_model_actors(getattr(self, "model_actors", None)) - except Exception: - pass - try: - ray.shutdown() - except Exception: - pass - - async def test_runner_state_collection(self): - scheduler = self._create_scheduler() - # 4 runner in side the scheduler - await scheduler.start() - - tasks = [ - Task( - workflow=DummyWorkflowWithState, # type: ignore[type-abstract] - workflow_args={"step_num": 2}, - repeat_times=4, - raw_task={}, - ) - for _ in range(4) - ] - scheduler.schedule(tasks, batch_id=0) - - async def monitor_routine(): - runner_0_state_history = defaultdict(set) - await asyncio.sleep(self.config.explorer.runner_state_report_interval + 0.2) - for _ in range(16): - await asyncio.sleep(0.3) - states = scheduler.get_all_state() - self.assertEqual(len(states), 4) - for state in states.values(): - self.assertIn("workflow_id", state) - self.assertIn("model_version", state) - self.assertIn("begin_time", state) - self.assertIn("terminate_time", state) - self.assertIn("repeat_cnt", state) - ids = scheduler.get_key_state("workflow_id") - self.assertEqual(len(ids), 4) - self.assertEqual(len(set(ids.values())), 4) - runner_0_state = scheduler.get_runner_state(0) - for key, value in runner_0_state.items(): - runner_0_state_history[key].add(value) - self.assertEqual(len(runner_0_state_history["repeat_cnt"]), 2) # max_repeat_times is 2 - self.assertEqual(len(runner_0_state_history["model_version"]), 1) - self.assertEqual( - len(runner_0_state_history["workflow_id"]), 2 - ) # split into 2 sub tasks - self.assertEqual(len(runner_0_state_history["begin_time"]), 2) - - await asyncio.gather( - monitor_routine(), - collect_results(scheduler, batch_id=0), - ) - await scheduler.stop() From a43d6e1e4808f7dd462ab3a371e6fcd905f49847 Mon Sep 17 00:00:00 2001 From: "panxuchen.pxc" Date: Tue, 30 Jun 2026 20:09:45 +0800 Subject: [PATCH 52/84] fix coordinator tests --- docker/sync.sh | 9 +++++++-- tests/explorer/rollout_coordinator_test.py | 11 +++++++++++ 2 files changed, 18 insertions(+), 2 deletions(-) diff --git a/docker/sync.sh b/docker/sync.sh index 7464a904ce3..ecaba0dae78 100755 --- a/docker/sync.sh +++ b/docker/sync.sh @@ -65,7 +65,6 @@ ssh_port="${port_override:-$TRINITY_REMOTE_SSH_PORT}" rsync_args=( -az --itemize-changes - --files-from=- --from0 -e "ssh -p ${ssh_port} -o StrictHostKeyChecking=accept-new" ) @@ -83,6 +82,12 @@ if [[ -n "$untracked" ]]; then echo "" >&2 fi +# Write file list to a temp file to avoid the "Bad file descriptor" race +# condition that occurs when rsync reads --files-from stdin via a pipe. +tmpfile="$(mktemp -t trinity-sync-XXXXXX)" +trap 'rm -f "$tmpfile"' EXIT +git -C "$PROJECT_DIR" ls-files -z > "$tmpfile" + dest="${TRINITY_REMOTE_HOST}:${TRINITY_REMOTE_WORKSPACE}/" echo "Syncing git-tracked files: ${PROJECT_DIR}/ -> ${dest}" -git -C "$PROJECT_DIR" ls-files -z | rsync "${rsync_args[@]}" "${PROJECT_DIR}/" "$dest" +rsync "${rsync_args[@]}" --files-from="$tmpfile" "${PROJECT_DIR}/" "$dest" diff --git a/tests/explorer/rollout_coordinator_test.py b/tests/explorer/rollout_coordinator_test.py index aeac50743b4..d90d5c70b21 100644 --- a/tests/explorer/rollout_coordinator_test.py +++ b/tests/explorer/rollout_coordinator_test.py @@ -145,6 +145,7 @@ def __init__(self, config, *, pipeline, scheduler): self._test_pipeline = pipeline self._test_scheduler = scheduler + self.discard_recorded_prefixes = [] super().__init__(config) async def _init_experience_pipeline(self): @@ -158,6 +159,11 @@ async def _init_scheduler(self): self.scheduler = self._test_scheduler + async def _discard_recorded_experiences(self, prefix: str) -> None: + """Record cleanup requests without resolving real rollout actors.""" + + self.discard_recorded_prefixes.append(prefix) + class TestRolloutCoordinator(unittest.IsolatedAsyncioTestCase): """Focused behavioral tests for the first coordinator implementation.""" @@ -199,6 +205,7 @@ async def test_finalize_train_batch_processes_scheduler_payloads(self): self.assertEqual(result["metrics"]["experience_pipeline/experience_count"], 2.0) self.assertTrue(self.pipeline.prepare_called) self.assertEqual(self.pipeline.process_chunk_calls, [[b"payload-0", b"payload-1"]]) + self.assertEqual(self.coordinator.discard_recorded_prefixes[-1], "1") self.assertNotIn(1, self.coordinator.pending_batches) with self.assertRaisesRegex(KeyError, "not registered"): @@ -224,6 +231,7 @@ async def test_finalize_train_batch_supports_partial_finalize(self): self.assertEqual(result["finished_task_count"], 1) self.assertEqual(self.pipeline.process_chunk_calls[-1], [b"payload-0"]) self.assertEqual(self.scheduler.abort_calls[-1]["batch_id"], 2) + self.assertIn("2", self.coordinator.discard_recorded_prefixes) self.assertNotIn(2, self.coordinator.pending_batches) async def test_finalize_train_batch_times_out_without_any_results(self): @@ -258,6 +266,7 @@ async def test_finalize_eval_batch_aggregates_eval_metrics(self): self.assertEqual(result["metrics"]["eval/eval_set/run_metrics"], 4.0) self.assertEqual(self.pipeline.process_chunk_calls, []) self.assertEqual(self.scheduler.get_statuses_calls[0]["batch_id"], batch_id) + self.assertEqual(self.coordinator.discard_recorded_prefixes[-1], batch_id) self.assertNotIn(batch_id, self.coordinator.pending_batches) async def test_finalize_train_batch_rejects_eval_batches_before_waiting(self): @@ -284,6 +293,7 @@ async def test_terminal_batches_are_not_reusable_after_finalize(self): ) self.scheduler.batch_results[eval_batch_id] = ([_build_status(3.0)], []) await self.coordinator.finalize_eval_batch(eval_batch_id, timeout=1.0) + self.assertEqual(self.coordinator.discard_recorded_prefixes[-1], eval_batch_id) with self.assertRaisesRegex(KeyError, "not registered"): await self.coordinator.finalize_train_batch(eval_batch_id, timeout=0.1) @@ -317,6 +327,7 @@ async def test_abort_batch_marks_batch_aborted_and_evicts_it(self): await self.coordinator.abort_batch(4, reason="shutdown") self.assertEqual(self.scheduler.abort_calls[0]["batch_id"], 4) + self.assertEqual(self.coordinator.discard_recorded_prefixes[-1], "4") self.assertNotIn(4, self.coordinator.pending_batches) with self.assertRaisesRegex(KeyError, "not registered"): From 881f9af5c9eeb9028e5af3258416720e2a08afaa Mon Sep 17 00:00:00 2001 From: "panxuchen.pxc" Date: Tue, 30 Jun 2026 20:23:05 +0800 Subject: [PATCH 53/84] fix tests --- trinity/common/workflows/workflow.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/trinity/common/workflows/workflow.py b/trinity/common/workflows/workflow.py index feef9982fe3..26ca8ba8470 100644 --- a/trinity/common/workflows/workflow.py +++ b/trinity/common/workflows/workflow.py @@ -502,3 +502,8 @@ async def update_reward( info=info, sample_ids=sample_ids, ) + + def set_single_run_context(self, run_id_base: int) -> None: + """Only a placeholder to align with the Workflow interface. + This workflow does not support repeat runs.""" + pass From 7c14499677e82ba7bbe76046517bc3d93e7597ff Mon Sep 17 00:00:00 2001 From: "panxuchen.pxc" Date: Tue, 30 Jun 2026 20:35:08 +0800 Subject: [PATCH 54/84] fix sglang default key --- tests/common/sglang_test.py | 5 +++++ trinity/common/models/recording/context.py | 2 ++ trinity/common/models/sglang_model.py | 2 ++ 3 files changed, 9 insertions(+) diff --git a/tests/common/sglang_test.py b/tests/common/sglang_test.py index 32886d5d918..050c74ca104 100644 --- a/tests/common/sglang_test.py +++ b/tests/common/sglang_test.py @@ -108,6 +108,9 @@ async def asyncSetUp(self): allocator = Allocator(self.config.explorer) rollout_models, _ = await allocator.create_all_models() self.model_wrapper = rollout_models[0] + self.record_key = "0/sglang_openai_api/0" + if self.enable_history: + self.model_wrapper.set_api_key(self.record_key) self.tokenizer = AutoTokenizer.from_pretrained( self.config.model.model_path, trust_remote_code=self.config.explorer.rollout_model.trust_remote_code, @@ -288,6 +291,7 @@ async def test_chat_completions(self): chat_exps = await self.model_wrapper.chat_async( messages, + enable_recording=self.enable_history, n=2, temperature=0.7, max_tokens=32, @@ -328,6 +332,7 @@ async def test_chat_completions(self): generate_prompt = "Write one short sentence about Boston." generate_exps = await self.model_wrapper.generate_async( [generate_prompt], + enable_recording=self.enable_history, n=2, temperature=0.7, max_tokens=32, diff --git a/trinity/common/models/recording/context.py b/trinity/common/models/recording/context.py index c31af625f58..db55a1af054 100644 --- a/trinity/common/models/recording/context.py +++ b/trinity/common/models/recording/context.py @@ -27,6 +27,8 @@ def extract_bearer_token(authorization: Optional[str]) -> Optional[str]: if scheme.lower() != "bearer": return None token = token.strip() + if token == "EMPTY": + return None return token or None diff --git a/trinity/common/models/sglang_model.py b/trinity/common/models/sglang_model.py index 7c40f6e67de..15c8c414009 100644 --- a/trinity/common/models/sglang_model.py +++ b/trinity/common/models/sglang_model.py @@ -60,6 +60,8 @@ def _auth_header(self, api_key_override: Optional[str]) -> str: # back to the client's api_key (which, on the Trinity Ray-direct path, # equals the record_key injected by ModelWrapper). token = api_key_override if api_key_override is not None else self.api_key + if token == "EMPTY": + token = None return f"Bearer {token}" if token else "" async def _server_call( From 8f69ffdc35d18b33fd5260f9c33e85965c81a1b2 Mon Sep 17 00:00:00 2001 From: "panxuchen.pxc" Date: Tue, 30 Jun 2026 20:44:19 +0800 Subject: [PATCH 55/84] remove legacy history --- tests/common/sglang_test.py | 12 +++++++++--- 1 file changed, 9 insertions(+), 3 deletions(-) diff --git a/tests/common/sglang_test.py b/tests/common/sglang_test.py index 050c74ca104..dc5a285269d 100644 --- a/tests/common/sglang_test.py +++ b/tests/common/sglang_test.py @@ -125,7 +125,6 @@ def _assert_experience_matches_text(self, exp, prompt_contents, response_text): def _assert_history_matches_responses(self, expected_count, prompt_contents, response_texts): if not self.enable_history: - self.assertEqual(len(self.model_wrapper.history), 0) return [] exps = self.model_wrapper.extract_experience_from_history() @@ -229,12 +228,16 @@ async def test_chat_completions(self): prompt_contents = [message["content"] for message in messages] openai_client = self.model_wrapper.get_openai_async_client() + routed_experts_body = ( + {"return_routed_experts": True} if self.enable_return_routed_experts else None + ) response = await openai_client.chat.completions.create( model=openai_client.model_path, messages=messages, n=1, temperature=0.7, max_tokens=32, + **({"extra_body": routed_experts_body} if routed_experts_body else {}), ) self.assertEqual(len(response.choices), 1) @@ -250,6 +253,7 @@ async def test_chat_completions(self): tool_choice="none", temperature=0.7, max_tokens=32, + **({"extra_body": routed_experts_body} if routed_experts_body else {}), ) self.assertEqual(len(tool_response.choices), 1) @@ -327,7 +331,8 @@ async def test_chat_completions(self): self.expected_routed_experts_topk, ) else: - self.assertEqual(len(self.model_wrapper.history), 0) + with self.assertRaises(ValueError): + self.model_wrapper.extract_experience_from_history() generate_prompt = "Write one short sentence about Boston." generate_exps = await self.model_wrapper.generate_async( @@ -369,7 +374,8 @@ async def test_chat_completions(self): self.expected_routed_experts_topk, ) else: - self.assertEqual(len(self.model_wrapper.history), 0) + with self.assertRaises(ValueError): + self.model_wrapper.extract_experience_from_history() class TestRecording(RayUnittestBaseAsync): From 14722e338dc0ae2a9fc65c0f0084d8571572d1e9 Mon Sep 17 00:00:00 2001 From: "panxuchen.pxc" Date: Tue, 30 Jun 2026 20:49:34 +0800 Subject: [PATCH 56/84] fix prompt path --- tests/common/sglang_test.py | 1 - tests/common/vllm_test.py | 2 -- 2 files changed, 3 deletions(-) diff --git a/tests/common/sglang_test.py b/tests/common/sglang_test.py index dc5a285269d..c397c7feb3a 100644 --- a/tests/common/sglang_test.py +++ b/tests/common/sglang_test.py @@ -362,7 +362,6 @@ async def test_chat_completions(self): self.assertEqual(len(generate_history), 2) for exp, recorded_exp in zip(generate_exps, generate_history): self.assertEqual(recorded_exp.response_text, exp.response_text) - self.assertEqual(recorded_exp.prompt_text, exp.prompt_text) self._assert_experience_matches_text( recorded_exp, [generate_prompt], exp.response_text ) diff --git a/tests/common/vllm_test.py b/tests/common/vllm_test.py index 23c769ca7b9..23fa6cfa4ec 100644 --- a/tests/common/vllm_test.py +++ b/tests/common/vllm_test.py @@ -724,7 +724,6 @@ async def test_api(self): self.assertTrue(response.choices[0].token_ids is None) with self.assertRaises(ValueError): self.model_wrapper_no_history.extract_experience_from_history() - self.assertEqual(len(self.model_wrapper_no_history.history), 0) class TestQwen35APIServerReasoning(VLLMTestBase): @@ -1191,7 +1190,6 @@ async def test_api_async(self): self.assertTrue(response.choices[0].token_ids is None) with self.assertRaises(ValueError): self.model_wrapper_no_history.extract_experience_from_history() - self.assertEqual(len(self.model_wrapper_no_history.history), 0) @unittest.skipIf("TINKER_API_KEY" not in os.environ, "TINKER_API_KEY is not set") From 85d753c29406b235ee798eb927c7dfa216ad14e2 Mon Sep 17 00:00:00 2001 From: "panxuchen.pxc" Date: Tue, 30 Jun 2026 21:01:59 +0800 Subject: [PATCH 57/84] fix enable_history --- tests/common/sglang_test.py | 3 +++ tests/common/vllm_test.py | 4 ++++ 2 files changed, 7 insertions(+) diff --git a/tests/common/sglang_test.py b/tests/common/sglang_test.py index c397c7feb3a..f8b3973129f 100644 --- a/tests/common/sglang_test.py +++ b/tests/common/sglang_test.py @@ -105,6 +105,7 @@ async def asyncSetUp(self): self.config.explorer.rollout_model.base_port = 13000 self.config.algorithm.enable_router_replay = self.enable_return_routed_experts self.config.check_and_update() + self.config.explorer.rollout_model.enable_history = self.enable_history allocator = Allocator(self.config.explorer) rollout_models, _ = await allocator.create_all_models() self.model_wrapper = rollout_models[0] @@ -125,6 +126,8 @@ def _assert_experience_matches_text(self, exp, prompt_contents, response_text): def _assert_history_matches_responses(self, expected_count, prompt_contents, response_texts): if not self.enable_history: + with self.assertRaises(ValueError): + self.model_wrapper.extract_experience_from_history() return [] exps = self.model_wrapper.extract_experience_from_history() diff --git a/tests/common/vllm_test.py b/tests/common/vllm_test.py index 23fa6cfa4ec..e7af9a721e4 100644 --- a/tests/common/vllm_test.py +++ b/tests/common/vllm_test.py @@ -164,6 +164,8 @@ async def asyncSetUp(self): self.config.algorithm.repeat_times = self.repeat_times self.config.explorer.rollout_model.enable_history = self.enable_history self.config.explorer.rollout_model.enable_openai_api = self.enable_return_routed_experts + requested_enable_history = self.config.explorer.rollout_model.enable_history + requested_enable_openai_api = self.config.explorer.rollout_model.enable_openai_api self.config.explorer.rollout_model.chat_template = CHAT_TEMPLATE self.config.explorer.rollout_model.extra_engine_args = {"max_num_seqs": 24} if self.enable_return_routed_experts: @@ -171,6 +173,8 @@ async def asyncSetUp(self): self.config.explorer.rollout_model.extra_engine_args["gdn_prefill_backend"] = "triton" self.config.algorithm.enable_router_replay = self.enable_return_routed_experts self.config.check_and_update() + self.config.explorer.rollout_model.enable_history = requested_enable_history + self.config.explorer.rollout_model.enable_openai_api = requested_enable_openai_api self.engines, self.auxiliary_engines = await create_test_models(self.config) self.model_wrapper = self.engines[0] From aee8aa77a2def773ec4bd6951cbfb79521bc2828 Mon Sep 17 00:00:00 2001 From: "panxuchen.pxc" Date: Tue, 30 Jun 2026 21:06:28 +0800 Subject: [PATCH 58/84] fix sglang auth --- trinity/common/models/sglang_model.py | 11 +++++------ 1 file changed, 5 insertions(+), 6 deletions(-) diff --git a/trinity/common/models/sglang_model.py b/trinity/common/models/sglang_model.py index 15c8c414009..97da3c5d5be 100644 --- a/trinity/common/models/sglang_model.py +++ b/trinity/common/models/sglang_model.py @@ -55,13 +55,12 @@ def __init__(self, server_url: str, api_key: Optional[str], logger: Logger): self.logger = logger def _auth_header(self, api_key_override: Optional[str]) -> str: - # The record_key travels as the bearer so the server-side - # RecordingIdentityMiddleware stamps it into ``record_key_ctx``. Falls - # back to the client's api_key (which, on the Trinity Ray-direct path, - # equals the record_key injected by ModelWrapper). + # ``api_key_override`` is a per-request record_key when recording is on. + # Otherwise fall back to the configured API key for SGLang auth. The + # default ``EMPTY`` token is still a valid auth token for no-history + # servers; RecordingIdentityMiddleware separately ignores it as a + # record_key. token = api_key_override if api_key_override is not None else self.api_key - if token == "EMPTY": - token = None return f"Bearer {token}" if token else "" async def _server_call( From 47c755df3d501b3a81ac1f159bca5c249779a110 Mon Sep 17 00:00:00 2001 From: "panxuchen.pxc" Date: Wed, 1 Jul 2026 10:20:32 +0800 Subject: [PATCH 59/84] simplify client side args --- tests/common/sglang_test.py | 16 ---------------- tests/common/vllm_test.py | 29 +++++++++++------------------ 2 files changed, 11 insertions(+), 34 deletions(-) diff --git a/tests/common/sglang_test.py b/tests/common/sglang_test.py index f8b3973129f..e1410d39711 100644 --- a/tests/common/sglang_test.py +++ b/tests/common/sglang_test.py @@ -144,15 +144,6 @@ def _assert_history_matches_responses(self, expected_count, prompt_contents, res ) return exps - def _assert_openai_response_routed_experts(self, response): - if not self.enable_return_routed_experts: - return - self.assertTrue(hasattr(response, "sglext")) - self.assertIsNotNone(response.sglext) - self.assertTrue("routed_experts" in response.sglext) - self.assertIsInstance(response.sglext["routed_experts"], str) - self.assertGreater(len(response.sglext["routed_experts"]), 0) - def _get_tool_call_case(self): tool_messages = [ {"role": "system", "content": "You are a helpful assistant."}, @@ -231,20 +222,15 @@ async def test_chat_completions(self): prompt_contents = [message["content"] for message in messages] openai_client = self.model_wrapper.get_openai_async_client() - routed_experts_body = ( - {"return_routed_experts": True} if self.enable_return_routed_experts else None - ) response = await openai_client.chat.completions.create( model=openai_client.model_path, messages=messages, n=1, temperature=0.7, max_tokens=32, - **({"extra_body": routed_experts_body} if routed_experts_body else {}), ) self.assertEqual(len(response.choices), 1) - self._assert_openai_response_routed_experts(response) response_texts = await self._collect_response_texts(response) self._assert_history_matches_responses(1, prompt_contents, response_texts) @@ -256,11 +242,9 @@ async def test_chat_completions(self): tool_choice="none", temperature=0.7, max_tokens=32, - **({"extra_body": routed_experts_body} if routed_experts_body else {}), ) self.assertEqual(len(tool_response.choices), 1) - self._assert_openai_response_routed_experts(tool_response) tool_response_texts = await self._collect_response_texts(tool_response) self._assert_history_matches_responses(1, tool_prompt_contents, tool_response_texts) diff --git a/tests/common/vllm_test.py b/tests/common/vllm_test.py index e7af9a721e4..c71892b4de2 100644 --- a/tests/common/vllm_test.py +++ b/tests/common/vllm_test.py @@ -179,15 +179,6 @@ async def asyncSetUp(self): self.engines, self.auxiliary_engines = await create_test_models(self.config) self.model_wrapper = self.engines[0] - def _assert_openai_response_routed_experts(self, response, expected_choices: int): - self.assertEqual(len(response.choices), expected_choices) - if not self.enable_return_routed_experts: - return - for choice in response.choices: - self.assertTrue(hasattr(choice, "routed_experts")) - self.assertIsInstance(choice.routed_experts, str) - self.assertGreater(len(choice.routed_experts), 0) - async def test_generate(self): # noqa: C901 self.assertEqual(self.model_wrapper.model_path, self.config.model.model_path) prompts = ["Hello, world!", "Hello, my name is"] @@ -329,7 +320,7 @@ async def test_generate(self): # noqa: C901 max_tokens=32, ) - self._assert_openai_response_routed_experts(openai_response, n) + self.assertEqual(len(openai_response.choices), n) history_experiences = self.model_wrapper.extract_experience_from_history() self.assertEqual(len(history_experiences), n) @@ -651,6 +642,7 @@ async def asyncSetUp(self): self.config.check_and_update() self.engines, self.auxiliary_engines = await create_test_models(self.config) self.model_wrapper = self.engines[0] + self.model_wrapper.set_api_key("0/vllm_api_server/0") self.model_wrapper_no_history = clone_wrapper(self.model_wrapper, enable_history=False) async def test_api(self): @@ -681,13 +673,13 @@ async def test_api(self): self.assertEqual(0, len(response.choices[0].logprobs.content[2].top_logprobs)) # here we check the 3rd token logprob, because the first two tokens (``,`\n` usually have zero logprob) self.assertTrue(response.choices[0].logprobs.content[2].logprob < 0) - self.assertTrue(hasattr(response, "prompt_token_ids")) - self.assertTrue(len(response.prompt_token_ids) > 0) - self.assertTrue(hasattr(response.choices[0], "token_ids")) - self.assertTrue(len(response.choices[0].token_ids) > 0) exps = self.model_wrapper.extract_experience_from_history() self.assertEqual(len(exps), 3) self.assertEqual(exps[0].response_text, content) + for exp in exps: + self.assertTrue(len(exp.tokens) > 0) + self.assertTrue(len(exp.logprobs) > 0) + self.assertTrue(exp.prompt_length + len(exp.logprobs) == len(exp.tokens)) response = openai_client.chat.completions.create( model=model_id, messages=messages, @@ -1111,6 +1103,7 @@ def _update_config(self): async def _setup_engines(self): self.engines, self.auxiliary_engines = await create_test_models(self.config) self.model_wrapper = self.engines[0] + self.model_wrapper.set_api_key("0/vllm_async_api_server/0") self.model_wrapper_no_history = clone_wrapper(self.model_wrapper, enable_history=False) async def test_api_async(self): @@ -1139,12 +1132,12 @@ async def test_api_async(self): # here we check the 3rd token logprob, because the first two tokens (``,`\n` usually have zero logprob) if "Instruct" not in self.model_path: self.assertTrue(response.choices[0].logprobs.content[2].logprob < 0) - self.assertTrue(hasattr(response, "prompt_token_ids")) - self.assertTrue(len(response.prompt_token_ids) > 0) - self.assertTrue(hasattr(response.choices[0], "token_ids")) - self.assertTrue(len(response.choices[0].token_ids) > 0) exps = self.model_wrapper.extract_experience_from_history() self.assertEqual(len(exps), 3) + for exp in exps: + self.assertTrue(len(exp.tokens) > 0) + self.assertTrue(len(exp.logprobs) > 0) + self.assertTrue(exp.prompt_length + len(exp.logprobs) == len(exp.tokens)) response = await openai_client.chat.completions.create( model=model_id, messages=messages, From b8a1c346c3ab2800836040b5bb2718a3698e13d1 Mon Sep 17 00:00:00 2001 From: pxc Date: Wed, 1 Jul 2026 10:40:57 +0800 Subject: [PATCH 60/84] fix vllm tests --- tests/common/vllm_test.py | 23 ++++++++++++++--------- 1 file changed, 14 insertions(+), 9 deletions(-) diff --git a/tests/common/vllm_test.py b/tests/common/vllm_test.py index c71892b4de2..3b84db3fed8 100644 --- a/tests/common/vllm_test.py +++ b/tests/common/vllm_test.py @@ -178,6 +178,7 @@ async def asyncSetUp(self): self.engines, self.auxiliary_engines = await create_test_models(self.config) self.model_wrapper = self.engines[0] + self.model_wrapper.set_api_key("model_wrapper/0/0") async def test_generate(self): # noqa: C901 self.assertEqual(self.model_wrapper.model_path, self.config.model.model_path) @@ -185,10 +186,12 @@ async def test_generate(self): # noqa: C901 n = self.config.algorithm.repeat_times if self.use_async: generate_results = await self.model_wrapper.generate_async( - prompts, n=n, temperature=1.0 + prompts, n=n, temperature=1.0, enable_recording=True, ) else: - generate_results = self.model_wrapper.generate(prompts, n=n, temperature=1.0) + generate_results = self.model_wrapper.generate( + prompts, n=n, temperature=1.0, enable_recording=True, + ) self.assertEqual(len(generate_results), len(prompts) * n) if self.enable_return_routed_experts: for exp in generate_results: @@ -228,9 +231,9 @@ async def test_generate(self): # noqa: C901 {"role": "user", "content": "OK, thanks!"}, ] if self.use_async: - results = await self.model_wrapper.chat_async(messages, n=n, temperature=1.0) + results = await self.model_wrapper.chat_async(messages, n=n, temperature=1.0, enable_recording=True) else: - results = self.model_wrapper.chat(messages, n=n, temperature=1.0) + results = self.model_wrapper.chat(messages, n=n, temperature=1.0, enable_recording=True) self.assertEqual(len(results), n) if self.enable_return_routed_experts: for exp in results: @@ -1499,11 +1502,13 @@ async def submitter(): self.assertTrue( logprobs_similar, - f"Logprobs for interrupted request {idx + 1} are not consistent " - f"after weight sync (mean_diff={mean_diff:.6f}, max_diff={max_diff:.6f}, " - f"num_mismatched={len(mismatch_indices) if not logprobs_similar else 0})" - if not logprobs_similar - else "", + ( + f"Logprobs for interrupted request {idx + 1} are not consistent " + f"after weight sync (mean_diff={mean_diff:.6f}, max_diff={max_diff:.6f}, " + f"num_mismatched={len(mismatch_indices) if not logprobs_similar else 0})" + if not logprobs_similar + else "" + ), ) else: print(" [WARNING] No matching experience found in history") From 0084f0fd2f75be9c648f96a707daa48275810a97 Mon Sep 17 00:00:00 2001 From: "panxuchen.pxc" Date: Wed, 1 Jul 2026 10:46:45 +0800 Subject: [PATCH 61/84] fix vllm tests --- tests/common/vllm_test.py | 81 ++++++++++++++++++++++++++------------- 1 file changed, 55 insertions(+), 26 deletions(-) diff --git a/tests/common/vllm_test.py b/tests/common/vllm_test.py index 3b84db3fed8..64ee48c1731 100644 --- a/tests/common/vllm_test.py +++ b/tests/common/vllm_test.py @@ -88,6 +88,44 @@ def _assert_routed_experts_shape(test_case, exp, expected_layers: int, expected_ ) +def _assert_recorded_experiences_match_unordered( + test_case, + expected_exps, + recorded_exps, + *, + enable_return_routed_experts: bool, + expected_layers: int, + expected_topk: int, +): + test_case.assertEqual(len(recorded_exps), len(expected_exps)) + unmatched_recorded = list(recorded_exps) + for exp in expected_exps: + exp_tokens = exp.tokens.tolist() + match_index = next( + ( + i + for i, recorded_exp in enumerate(unmatched_recorded) + if recorded_exp.tokens.tolist() == exp_tokens + ), + None, + ) + test_case.assertIsNotNone( + match_index, + f"Recorded history does not contain expected response: {exp.response_text[:200]}", + ) + recorded_exp = unmatched_recorded.pop(match_index) + test_case.assertEqual(exp.response_text, recorded_exp.response_text) + test_case.assertEqual(exp.prompt_length, recorded_exp.prompt_length) + test_case.assertEqual(exp.logprobs.tolist(), recorded_exp.logprobs.tolist()) + if enable_return_routed_experts: + _assert_routed_experts_shape( + test_case, + recorded_exp, + expected_layers, + expected_topk, + ) + + def _load_gsm8k_questions() -> list[str]: """Load the diverse math questions from the GSM8K training set.""" path = os.path.join(os.path.dirname(__file__), "..", "template", "data", "gsm8k", "train.jsonl") @@ -205,19 +243,14 @@ async def test_generate(self): # noqa: C901 history_experiences = self.model_wrapper.extract_experience_from_history( clear_history=False ) - self.assertEqual(len(history_experiences), len(generate_results)) - for exp, history_exp in zip(generate_results, history_experiences): - self.assertEqual(exp.response_text, history_exp.response_text) - self.assertEqual(exp.tokens.tolist(), history_exp.tokens.tolist()) - self.assertEqual(exp.prompt_length, history_exp.prompt_length) - self.assertEqual(exp.logprobs.tolist(), history_exp.logprobs.tolist()) - if self.enable_return_routed_experts: - _assert_routed_experts_shape( - self, - history_exp, - self.expected_routed_experts_layers, - self.expected_routed_experts_topk, - ) + _assert_recorded_experiences_match_unordered( + self, + generate_results, + history_experiences, + enable_return_routed_experts=self.enable_return_routed_experts, + expected_layers=self.expected_routed_experts_layers, + expected_topk=self.expected_routed_experts_topk, + ) else: with self.assertRaises(ValueError): self.model_wrapper.extract_experience_from_history(clear_history=False) @@ -245,19 +278,15 @@ async def test_generate(self): # noqa: C901 ) if self.config.explorer.rollout_model.enable_history: history_experiences = self.model_wrapper.extract_experience_from_history() - self.assertEqual(len(history_experiences) - len(generate_results), len(results)) - for exp, history_exp in zip(results, history_experiences[len(generate_results) :]): - self.assertEqual(exp.response_text, history_exp.response_text) - self.assertEqual(exp.tokens.tolist(), history_exp.tokens.tolist()) - self.assertEqual(exp.prompt_length, history_exp.prompt_length) - self.assertEqual(exp.logprobs.tolist(), history_exp.logprobs.tolist()) - if self.enable_return_routed_experts: - _assert_routed_experts_shape( - self, - history_exp, - self.expected_routed_experts_layers, - self.expected_routed_experts_topk, - ) + self.assertEqual(len(history_experiences), len(generate_results) + len(results)) + _assert_recorded_experiences_match_unordered( + self, + results, + history_experiences[len(generate_results) :], + enable_return_routed_experts=self.enable_return_routed_experts, + expected_layers=self.expected_routed_experts_layers, + expected_topk=self.expected_routed_experts_topk, + ) for result in results: self.assertTrue(torch.any(result.logprobs != 0)) if self.use_async: From 52c22968df8c20c72b13f86f1ef38c9d2aae868d Mon Sep 17 00:00:00 2001 From: pxc Date: Wed, 1 Jul 2026 11:16:50 +0800 Subject: [PATCH 62/84] record --- tests/common/vllm_test.py | 29 +++++++++++++++++++++-------- 1 file changed, 21 insertions(+), 8 deletions(-) diff --git a/tests/common/vllm_test.py b/tests/common/vllm_test.py index 64ee48c1731..6da893fd4f8 100644 --- a/tests/common/vllm_test.py +++ b/tests/common/vllm_test.py @@ -113,7 +113,7 @@ def _assert_recorded_experiences_match_unordered( match_index, f"Recorded history does not contain expected response: {exp.response_text[:200]}", ) - recorded_exp = unmatched_recorded.pop(match_index) + recorded_exp = unmatched_recorded.pop(match_index) # type: ignore [arg-type] test_case.assertEqual(exp.response_text, recorded_exp.response_text) test_case.assertEqual(exp.prompt_length, recorded_exp.prompt_length) test_case.assertEqual(exp.logprobs.tolist(), recorded_exp.logprobs.tolist()) @@ -224,11 +224,17 @@ async def test_generate(self): # noqa: C901 n = self.config.algorithm.repeat_times if self.use_async: generate_results = await self.model_wrapper.generate_async( - prompts, n=n, temperature=1.0, enable_recording=True, + prompts, + n=n, + temperature=1.0, + enable_recording=True, ) else: generate_results = self.model_wrapper.generate( - prompts, n=n, temperature=1.0, enable_recording=True, + prompts, + n=n, + temperature=1.0, + enable_recording=True, ) self.assertEqual(len(generate_results), len(prompts) * n) if self.enable_return_routed_experts: @@ -264,7 +270,9 @@ async def test_generate(self): # noqa: C901 {"role": "user", "content": "OK, thanks!"}, ] if self.use_async: - results = await self.model_wrapper.chat_async(messages, n=n, temperature=1.0, enable_recording=True) + results = await self.model_wrapper.chat_async( + messages, n=n, temperature=1.0, enable_recording=True + ) else: results = self.model_wrapper.chat(messages, n=n, temperature=1.0, enable_recording=True) self.assertEqual(len(results), n) @@ -322,7 +330,7 @@ async def test_generate(self): # noqa: C901 ) self.assertTrue(exp.logprobs.shape[0] == exp.tokens.shape[0] - prompt_length) self.assertTrue(torch.equal(result_dict["input_ids"][0], exp.tokens)) - if self.enable_return_routed_experts: + if self.model_wrapper.config.enable_openai_api: self.assertIsNotNone(self.model_wrapper.get_openai_client()) else: self.assertRaises(ValueError, self.model_wrapper.get_openai_client) @@ -455,6 +463,7 @@ async def asyncSetUp(self): self.engines, self.auxiliary_engines = await create_test_models(self.config) self.model_wrapper = self.engines[0] + self.model_wrapper.set_api_key("model_len/0/0") self.tokenizer = AutoTokenizer.from_pretrained(self.config.model.model_path) async def test_model_len(self): @@ -489,7 +498,7 @@ def _check_experience(exp): self.assertLessEqual(len(exp.tokens), self.config.model.max_model_len) # For vllm engine, max_prompt_tokens and max_response_tokens work - response = self.model_wrapper.chat(messages) + response = self.model_wrapper.chat(messages, enable_recording=True) self.assertEqual(len(response), 1) if self.max_prompt_tokens == 5: self.assertEqual(response[0].truncate_status, "prompt_truncated") @@ -526,7 +535,7 @@ def _check_experience(exp): ][0].tolist() self.assertGreater(len(prompt_token_ids), self.config.model.max_prompt_tokens) - responses = self.model_wrapper.generate([prompt], n=2) + responses = self.model_wrapper.generate([prompt], n=2, enable_recording=True) self.assertEqual(len(responses), 2) for response in responses: @@ -555,6 +564,7 @@ async def asyncSetUp(self): self.engines, self.auxiliary_engines = await create_test_models(self.config) self.model_wrapper = self.engines[0] + self.model_wrapper.set_api_key("model_len_no_truncation/0/0") async def test_model_len(self): messages = [ @@ -562,7 +572,7 @@ async def test_model_len(self): ] # For vllm engine, max_prompt_tokens and max_response_tokens work - response = self.model_wrapper.chat(messages) + response = self.model_wrapper.chat(messages, enable_recording=True) self.assertEqual(len(response), 1) self.assertLessEqual( len(response[0].tokens) - response[0].prompt_length, @@ -771,6 +781,7 @@ async def asyncSetUp(self): self.config.check_and_update() self.engines, self.auxiliary_engines = await create_test_models(self.config) self.model_wrapper = self.engines[0] + self.model_wrapper.set_api_key("qwen35_reasoning/0/0") async def test_reasoning_content(self): openai_client = self.model_wrapper.get_openai_client() @@ -848,6 +859,7 @@ async def asyncSetUp(self): self.config.check_and_update() self.engines, self.auxiliary_engines = await create_test_models(self.config) self.model_wrapper = self.engines[0] + self.model_wrapper.set_api_key("qwen35_mm/0/0") async def test_multi_modal_content(self): openai_client = self.model_wrapper.get_openai_client() @@ -942,6 +954,7 @@ async def asyncSetUp(self): self.config.check_and_update() self.engines, self.auxiliary_engines = await create_test_models(self.config) self.model_wrapper = self.engines[0] + self.model_wrapper.set_api_key("logprobs/0/0") async def test_logprobs_api(self): messages = [ From 9c5727760fd36c027658dace4036879369d8e6bb Mon Sep 17 00:00:00 2001 From: pxc Date: Wed, 1 Jul 2026 11:29:30 +0800 Subject: [PATCH 63/84] simplify pipeline --- trinity/buffer/pipelines/experience_pipeline.py | 8 +++----- 1 file changed, 3 insertions(+), 5 deletions(-) diff --git a/trinity/buffer/pipelines/experience_pipeline.py b/trinity/buffer/pipelines/experience_pipeline.py index 0839c5fb4c1..c4dc1141c41 100644 --- a/trinity/buffer/pipelines/experience_pipeline.py +++ b/trinity/buffer/pipelines/experience_pipeline.py @@ -169,7 +169,7 @@ async def process(self, exp_bytes: bytes) -> Dict: Dict: A dictionary containing metrics collected during the processing of experiences. """ exps = Experience.deserialize_many(exp_bytes) - return await self._process_experiences(exps) + return await self.process_experiences(exps) async def process_serialized_chunks(self, exp_chunks: list[bytes]) -> Dict: """Process a batch assembled from multiple serialized task payloads.""" @@ -178,7 +178,8 @@ async def process_serialized_chunks(self, exp_chunks: list[bytes]) -> Dict: if not exp_bytes: continue exps.extend(Experience.deserialize_many(exp_bytes)) - return await self._process_experiences(exps) + return await self.process_experiences(exps) + async def process_experiences(self, exps: list[Experience]) -> Dict: """Process already-assembled experiences (objects, not serialized bytes). @@ -188,9 +189,6 @@ async def process_experiences(self, exps: list[Experience]) -> Dict: over directly — avoiding a serialize/deserialize round-trip for the heavy tensor payload. """ - return await self._process_experiences(exps) - - async def _process_experiences(self, exps: list[Experience]) -> Dict: st = time.time() if self.input_store is not None: await self.input_store.write(exps) From ed4d2d88caa962d0fa1d300c77355b90d84df7c7 Mon Sep 17 00:00:00 2001 From: "panxuchen.pxc" Date: Wed, 1 Jul 2026 11:47:10 +0800 Subject: [PATCH 64/84] fix logprobs --- tests/common/vllm_test.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/tests/common/vllm_test.py b/tests/common/vllm_test.py index 6da893fd4f8..4f44175d819 100644 --- a/tests/common/vllm_test.py +++ b/tests/common/vllm_test.py @@ -751,9 +751,11 @@ async def test_api(self): messages=messages, logprobs=False, ) + self.assertIsNone(response.choices[0].logprobs) exps = self.model_wrapper.extract_experience_from_history() self.assertEqual(len(exps), 1) - self.assertTrue(len(exps[0].logprobs) == 0) + self.assertTrue(len(exps[0].logprobs) > 0) + self.assertTrue(exps[0].prompt_length + len(exps[0].logprobs) == len(exps[0].tokens)) response = self.model_wrapper_no_history.get_openai_client().chat.completions.create( model=model_id, messages=messages, n=2 ) @@ -1219,9 +1221,11 @@ async def test_api_async(self): messages=messages, logprobs=False, ) + self.assertIsNone(response.choices[0].logprobs) exps = self.model_wrapper.extract_experience_from_history() self.assertEqual(len(exps), 1) - self.assertTrue(len(exps[0].logprobs) == 0) + self.assertTrue(len(exps[0].logprobs) > 0) + self.assertTrue(exps[0].prompt_length + len(exps[0].logprobs) == len(exps[0].tokens)) response = ( await self.model_wrapper_no_history.get_openai_async_client().chat.completions.create( model=model_id, messages=messages, n=2 From 8864c36fd01762893b0f111dae0c699ec9a1db4b Mon Sep 17 00:00:00 2001 From: "panxuchen.pxc" Date: Wed, 1 Jul 2026 12:03:50 +0800 Subject: [PATCH 65/84] fix reasoning parser --- tests/common/vllm_test.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/tests/common/vllm_test.py b/tests/common/vllm_test.py index 4f44175d819..e625c044f6e 100644 --- a/tests/common/vllm_test.py +++ b/tests/common/vllm_test.py @@ -1807,7 +1807,10 @@ async def test_api_tool_calls(self): final_exps = self.model_wrapper.extract_experience_from_history() self.assertEqual(len(final_exps), 1) print_debug(f" > Final recorded experience response_text: {final_exps[0].response_text}") - self.assertEqual(final_exps[0].response_text, final_choice.message.content) + if self.reasoning_parser: + self.assertIn(final_choice.message.content.strip(), final_exps[0].response_text) + else: + self.assertEqual(final_exps[0].response_text, final_choice.message.content) print_debug(f"[{time.time() - start_time:.2f}s] Final experience history check passed.") exp = final_exps[0] From 0d359a4cbc1bcce3a7d9fcade89aa1c597fc6ed4 Mon Sep 17 00:00:00 2001 From: pxc Date: Wed, 1 Jul 2026 11:58:36 +0800 Subject: [PATCH 66/84] fix toolcall --- tests/common/vllm_test.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/common/vllm_test.py b/tests/common/vllm_test.py index e625c044f6e..821f2de582a 100644 --- a/tests/common/vllm_test.py +++ b/tests/common/vllm_test.py @@ -1600,6 +1600,7 @@ async def asyncSetUp(self): self.config.check_and_update() self.engines, self.auxiliary_engines = await create_test_models(self.config) self.model_wrapper = self.engines[0] + self.model_wrapper.set_api_key("tool_call/0/0") self.model_wrapper_no_history = clone_wrapper(self.model_wrapper, enable_history=False) async def test_api_tool_calls(self): From 14b726282c2f9a7e3dbb1b18bfd1d30db2d729bc Mon Sep 17 00:00:00 2001 From: "panxuchen.pxc" Date: Wed, 1 Jul 2026 12:23:48 +0800 Subject: [PATCH 67/84] update recording context --- tests/common/vllm_test.py | 4 +- trinity/common/models/recording/__init__.py | 10 +++- trinity/common/models/recording/context.py | 54 ++++++++++++++++--- .../models/sglang_patch/recording/__init__.py | 8 ++- .../models/sglang_patch/recording/recorder.py | 4 +- trinity/common/models/vllm_model.py | 14 +++-- .../models/vllm_patch/recording/__init__.py | 10 +++- .../models/vllm_patch/recording/recorder.py | 33 +++++++----- 8 files changed, 103 insertions(+), 34 deletions(-) diff --git a/tests/common/vllm_test.py b/tests/common/vllm_test.py index 821f2de582a..4989ad026f2 100644 --- a/tests/common/vllm_test.py +++ b/tests/common/vllm_test.py @@ -1870,7 +1870,7 @@ class TestRecording(VLLMTestBase): Paths covered (all async): * Ray-direct ``generate`` / ``chat`` — record_key propagated via - ``record_key_ctx`` (set inside the actor by ``VLLMModel``). + ``recording_ctx`` (set inside the actor by ``VLLMModel``). * OpenAI HTTP regular / streaming / tool-call — record_key propagated via the ``Authorization: Bearer `` header, captured by ``RecordingIdentityMiddleware``. @@ -1983,7 +1983,7 @@ async def test_record(self): # noqa: C901 ] no_think = {"chat_template_kwargs": {"enable_thinking": False}} - # ===== 1. Ray-direct generate (record_key via record_key_ctx) ===== + # ===== 1. Ray-direct generate (record_key via recording_ctx) ===== rk_gen = "0/t_gen/1" await self.model_wrapper.generate_async( ["Hello, world!"], n=1, temperature=1.0, max_tokens=16, key=rk_gen diff --git a/trinity/common/models/recording/__init__.py b/trinity/common/models/recording/__init__.py index 2d849090c23..8b8200cf999 100644 --- a/trinity/common/models/recording/__init__.py +++ b/trinity/common/models/recording/__init__.py @@ -2,10 +2,13 @@ from trinity.buffer.store import MemoryStore, RecordStore from trinity.common.models.recording.context import ( + RecordingContext, RecordingIdentityMiddleware, extract_bearer_token, get_recording_record_key, - record_key_ctx, + get_recording_record_key_from_context, + get_recording_request_from_context, + recording_ctx, skip_recording_ctx, ) from trinity.common.models.recording.recorder import ( @@ -21,6 +24,7 @@ __all__ = [ "MemoryStore", "Recorder", + "RecordingContext", "RecordingIdentityMiddleware", "RecordStore", "TRINITY_RECORD_STORE_ATTR", @@ -28,7 +32,9 @@ "add_recording_middleware", "extract_bearer_token", "get_recording_record_key", + "get_recording_record_key_from_context", + "get_recording_request_from_context", "mount_recording_api", - "record_key_ctx", + "recording_ctx", "skip_recording_ctx", ] diff --git a/trinity/common/models/recording/context.py b/trinity/common/models/recording/context.py index db55a1af054..e149d11c086 100644 --- a/trinity/common/models/recording/context.py +++ b/trinity/common/models/recording/context.py @@ -1,6 +1,7 @@ -"""Per-request recording identity propagation shared by model engines.""" +"""Per-request recording context propagation shared by model engines.""" from contextvars import ContextVar +from dataclasses import dataclass from typing import Any, Optional try: @@ -8,9 +9,17 @@ except ModuleNotFoundError: BaseHTTPMiddleware = object # type: ignore -# The record key for the in-flight request (the MemoryStore group key), or None -# when the client did not send a supported identity header. -record_key_ctx: ContextVar[Optional[str]] = ContextVar("trinity_recording_record_key", default=None) +@dataclass(frozen=True) +class RecordingContext: + """Per-request recording metadata propagated to engine-boundary recorders.""" + + record_key: Optional[str] = None + request: Optional[dict[str, Any]] = None + + +recording_ctx: ContextVar[Optional[RecordingContext]] = ContextVar( + "trinity_recording_context", default=None +) # Set around auxiliary generate calls (logprobs recomputation, message # conversion) so recorders skip them. @@ -37,13 +46,44 @@ def get_recording_record_key(request: Any) -> Optional[str]: return extract_bearer_token(request.headers.get(AUTHORIZATION_HEADER)) +def get_recording_record_key_from_context() -> Optional[str]: + """Return the current in-flight recording identity, if any.""" + ctx = recording_ctx.get() + return None if ctx is None else ctx.record_key + + +def get_recording_request_from_context() -> Optional[dict[str, Any]]: + """Return selected raw OpenAI request fields captured for recording.""" + ctx = recording_ctx.get() + return None if ctx is None else ctx.request + + class RecordingIdentityMiddleware(BaseHTTPMiddleware): - """Capture request identity into ``record_key_ctx`` for the request lifetime.""" + """Capture request identity and selected raw request fields.""" + + async def _get_recording_request(self, request: Any, record_key: Optional[str]): + if record_key is None: + return None + try: + body = await request.json() + except Exception: + return None + if not isinstance(body, dict): + return None + recording_request = {} + for field in ("messages", "tools"): + value = body.get(field) + if value is not None: + recording_request[field] = value + return recording_request or None async def dispatch(self, request: Any, call_next): record_key = get_recording_record_key(request) - token = record_key_ctx.set(record_key) + request_info = await self._get_recording_request(request, record_key) + token = recording_ctx.set( + RecordingContext(record_key=record_key, request=request_info) + ) try: return await call_next(request) finally: - record_key_ctx.reset(token) + recording_ctx.reset(token) diff --git a/trinity/common/models/sglang_patch/recording/__init__.py b/trinity/common/models/sglang_patch/recording/__init__.py index 4811f6e518c..54b2dede670 100644 --- a/trinity/common/models/sglang_patch/recording/__init__.py +++ b/trinity/common/models/sglang_patch/recording/__init__.py @@ -8,9 +8,11 @@ from trinity.buffer.store import MemoryStore, RecordStore # noqa: F401 from trinity.common.models.recording.context import ( # noqa: F401 + RecordingContext, RecordingIdentityMiddleware, get_recording_record_key, - record_key_ctx, + get_recording_record_key_from_context, + recording_ctx, skip_recording_ctx, ) from trinity.common.models.recording.recorder import Recorder # noqa: F401 @@ -29,12 +31,14 @@ "MemoryStore", "RecordStore", "Recorder", + "RecordingContext", "RecordingIdentityMiddleware", "build_sglang_experience", "create_sglang_recorder", "get_recording_record_key", + "get_recording_record_key_from_context", "patch_tokenizer_manager_for_recording", - "record_key_ctx", + "recording_ctx", "setup_sglang_recording", "skip_recording_ctx", ] diff --git a/trinity/common/models/sglang_patch/recording/recorder.py b/trinity/common/models/sglang_patch/recording/recorder.py index 747fcffb82f..1d48ee537e3 100644 --- a/trinity/common/models/sglang_patch/recording/recorder.py +++ b/trinity/common/models/sglang_patch/recording/recorder.py @@ -29,7 +29,7 @@ from typing import Any, List, Optional, Tuple from trinity.buffer.store import MemoryStore, RecordStore -from trinity.common.models.recording.context import record_key_ctx +from trinity.common.models.recording.context import get_recording_record_key_from_context from trinity.common.models.recording.recorder import ( TRINITY_RECORD_STORE_ATTR, TRINITY_RECORDER_ATTR, @@ -273,7 +273,7 @@ async def _patched_generate_request(*args, **kwargs): # Trigger on the finished yield (not on generator exhaustion): the # non-stream /generate consumer pulls only once via __anext__(). if recorder.enabled and _is_finished(out): - record_key = record_key_ctx.get() + record_key = get_recording_record_key_from_context() if record_key is not None and state: reconstructed = _build_ret(state, order) recorder.schedule_record( diff --git a/trinity/common/models/vllm_model.py b/trinity/common/models/vllm_model.py index 896db6b80f6..b3935344aa7 100644 --- a/trinity/common/models/vllm_model.py +++ b/trinity/common/models/vllm_model.py @@ -16,7 +16,11 @@ from trinity.common.experience import Experience from trinity.common.models.mm_utils import vLLMMultiModalRender from trinity.common.models.model import BaseInferenceModel -from trinity.common.models.recording.context import record_key_ctx, skip_recording_ctx +from trinity.common.models.recording.context import ( + RecordingContext, + recording_ctx, + skip_recording_ctx, +) from trinity.common.models.vllm_patch import get_vllm_version from trinity.common.models.vllm_patch.recording.models import build_experience @@ -229,7 +233,7 @@ async def chat( messages (List[dict]): The input history messages. key (Optional[str]): Recording identity for the in-vLLM recorder (the MemoryStore group key). Propagated to - ``generate`` via ``record_key_ctx`` so the recorder stamps it + ``generate`` via ``recording_ctx`` so the recorder stamps it into ``Experience.eid`` without an HTTP hop. None skips recording. kwargs (dict): A dictionary of sampling parameters. @@ -271,7 +275,7 @@ async def generate( Args: prompt (str): The input prompt. key (Optional[str]): Recording identity propagated to the - in-vLLM recorder via ``record_key_ctx`` (see ``chat``). + in-vLLM recorder via ``recording_ctx`` (see ``chat``). kwargs (dict): A dictionary of sampling parameters. Returns: @@ -308,13 +312,13 @@ async def generate( # Propagate the recording identity to the engine-level recorder (same # async task, same process) so the recorded experience is grouped under # this record key in the MemoryStore. - record_key_token = record_key_ctx.set(key) + record_key_token = recording_ctx.set(RecordingContext(record_key=key)) try: output = await self._generate_internal( prompt=prompt, lora_request=lora_request, **kwargs ) finally: - record_key_ctx.reset(record_key_token) + recording_ctx.reset(record_key_token) if is_mm_prompt: if self.mm_render is None: self.mm_render = vLLMMultiModalRender( diff --git a/trinity/common/models/vllm_patch/recording/__init__.py b/trinity/common/models/vllm_patch/recording/__init__.py index 857d865bb54..7dd75f61a51 100644 --- a/trinity/common/models/vllm_patch/recording/__init__.py +++ b/trinity/common/models/vllm_patch/recording/__init__.py @@ -4,8 +4,11 @@ from trinity.buffer.store import MemoryStore, RecordStore from trinity.common.models.recording.context import ( + RecordingContext, RecordingIdentityMiddleware, - record_key_ctx, + get_recording_record_key_from_context, + get_recording_request_from_context, + recording_ctx, skip_recording_ctx, ) from trinity.common.models.recording.recorder import Recorder @@ -21,12 +24,15 @@ __all__ = [ "MemoryStore", "RecordStore", + "RecordingContext", "RecordingIdentityMiddleware", "Recorder", "build_experience", "create_vllm_recorder", + "get_recording_record_key_from_context", + "get_recording_request_from_context", "patch_engine_for_recording", - "record_key_ctx", + "recording_ctx", "run_api_server_with_recording", "skip_recording_ctx", ] diff --git a/trinity/common/models/vllm_patch/recording/recorder.py b/trinity/common/models/vllm_patch/recording/recorder.py index 7770ff7ff07..8ae841bff5b 100644 --- a/trinity/common/models/vllm_patch/recording/recorder.py +++ b/trinity/common/models/vllm_patch/recording/recorder.py @@ -24,7 +24,10 @@ from typing import Optional from trinity.buffer.store import MemoryStore, RecordStore -from trinity.common.models.recording.context import record_key_ctx +from trinity.common.models.recording.context import ( + get_recording_record_key_from_context, + get_recording_request_from_context, +) from trinity.common.models.recording.recorder import ( TRINITY_RECORD_STORE_ATTR, TRINITY_RECORDER_ATTR, @@ -85,24 +88,30 @@ def _get_prompt_arg(args, kwargs): def _build_multi_modal_inputs(engine_client, prompt, output, logger: logging.Logger): - if not isinstance(prompt, dict): - return None - multi_modal_data = prompt.get("multi_modal_data") - if not multi_modal_data: - return None mm_render = getattr(engine_client, TRINITY_MM_RENDER_ATTR, None) if mm_render is None: logger.warning( - "Recording saw a multimodal vLLM prompt but no %s is attached to engine_client; " + "Recording saw a possible multimodal vLLM prompt but no %s is attached to engine_client; " "recorded Experience will not include multi_modal_inputs.", TRINITY_MM_RENDER_ATTR, ) return None try: - return mm_render.build_mm_input_for_training( - input_ids=output.prompt_token_ids, - multi_modal_data=multi_modal_data, - ) + if isinstance(prompt, dict): + multi_modal_data = prompt.get("multi_modal_data") + if multi_modal_data: + return mm_render.build_mm_input_for_training( + input_ids=output.prompt_token_ids, + multi_modal_data=multi_modal_data, + ) + request_info = get_recording_request_from_context() + if request_info and request_info.get("messages") is not None: + return mm_render.build_mm_input_for_training( + input_ids=output.prompt_token_ids, + messages=request_info["messages"], + tools=request_info.get("tools"), + ) + return None except Exception: logger.exception("Failed to build multi_modal_inputs for recorded vLLM Experience") return None @@ -300,7 +309,7 @@ async def _patched_generate(*args, **kwargs): # RecordingIdentityMiddleware on the HTTP path, or by VLLMModel.chat # on the Ray-direct path). A missing key means the caller did not # opt into grouping this turn, so skip recording entirely. - record_key = record_key_ctx.get() + record_key = get_recording_record_key_from_context() if record_key is not None: record_output = _build_record_output(accumulated, last) multi_modal_inputs = _build_multi_modal_inputs( From 892cefcd654422fcac3a762ad7c80709586e413a Mon Sep 17 00:00:00 2001 From: "panxuchen.pxc" Date: Wed, 1 Jul 2026 12:28:33 +0800 Subject: [PATCH 68/84] fix pre-commit --- trinity/buffer/pipelines/experience_pipeline.py | 1 - trinity/common/models/recording/context.py | 5 ++--- trinity/common/models/sglang_patch/recording/recorder.py | 4 +++- 3 files changed, 5 insertions(+), 5 deletions(-) diff --git a/trinity/buffer/pipelines/experience_pipeline.py b/trinity/buffer/pipelines/experience_pipeline.py index c4dc1141c41..fe8dfadce62 100644 --- a/trinity/buffer/pipelines/experience_pipeline.py +++ b/trinity/buffer/pipelines/experience_pipeline.py @@ -180,7 +180,6 @@ async def process_serialized_chunks(self, exp_chunks: list[bytes]) -> Dict: exps.extend(Experience.deserialize_many(exp_bytes)) return await self.process_experiences(exps) - async def process_experiences(self, exps: list[Experience]) -> Dict: """Process already-assembled experiences (objects, not serialized bytes). diff --git a/trinity/common/models/recording/context.py b/trinity/common/models/recording/context.py index e149d11c086..177c1b3c12d 100644 --- a/trinity/common/models/recording/context.py +++ b/trinity/common/models/recording/context.py @@ -9,6 +9,7 @@ except ModuleNotFoundError: BaseHTTPMiddleware = object # type: ignore + @dataclass(frozen=True) class RecordingContext: """Per-request recording metadata propagated to engine-boundary recorders.""" @@ -80,9 +81,7 @@ async def _get_recording_request(self, request: Any, record_key: Optional[str]): async def dispatch(self, request: Any, call_next): record_key = get_recording_record_key(request) request_info = await self._get_recording_request(request, record_key) - token = recording_ctx.set( - RecordingContext(record_key=record_key, request=request_info) - ) + token = recording_ctx.set(RecordingContext(record_key=record_key, request=request_info)) try: return await call_next(request) finally: diff --git a/trinity/common/models/sglang_patch/recording/recorder.py b/trinity/common/models/sglang_patch/recording/recorder.py index 1d48ee537e3..c602a707757 100644 --- a/trinity/common/models/sglang_patch/recording/recorder.py +++ b/trinity/common/models/sglang_patch/recording/recorder.py @@ -29,7 +29,9 @@ from typing import Any, List, Optional, Tuple from trinity.buffer.store import MemoryStore, RecordStore -from trinity.common.models.recording.context import get_recording_record_key_from_context +from trinity.common.models.recording.context import ( + get_recording_record_key_from_context, +) from trinity.common.models.recording.recorder import ( TRINITY_RECORD_STORE_ATTR, TRINITY_RECORDER_ATTR, From da9b0f1563cdf2b214001b5cf836db161a00699e Mon Sep 17 00:00:00 2001 From: "panxuchen.pxc" Date: Wed, 1 Jul 2026 13:31:34 +0800 Subject: [PATCH 69/84] clean scheduler --- tests/explorer/rollout_coordinator_test.py | 12 ++++++------ tests/explorer/scheduler_test.py | 4 ++-- trinity/explorer/rollout_coordinator.py | 4 ++-- trinity/explorer/scheduler.py | 19 +++---------------- 4 files changed, 13 insertions(+), 26 deletions(-) diff --git a/tests/explorer/rollout_coordinator_test.py b/tests/explorer/rollout_coordinator_test.py index d90d5c70b21..d649e6d6772 100644 --- a/tests/explorer/rollout_coordinator_test.py +++ b/tests/explorer/rollout_coordinator_test.py @@ -46,7 +46,7 @@ def __init__(self): self.stopped = False self.schedule_calls = [] self.scheduled_task_counts = {} - self.abort_calls = [] + self.cleanup_calls = [] self.batch_results = {} self.get_statuses_calls = [] @@ -99,10 +99,10 @@ async def get_payload_results( _ = min_num, timeout, clear_timeout_tasks, return_partial_tasks return self.batch_results.pop(batch_id, ([], [])) - async def abort_batch(self, batch_id, return_partial_tasks=False, restart_runners=True): - """Record one scheduler abort request.""" + async def cleanup_batch(self, batch_id, return_partial_tasks=False, restart_runners=True): + """Record one scheduler cleanup request.""" - self.abort_calls.append( + self.cleanup_calls.append( { "batch_id": batch_id, "return_partial_tasks": return_partial_tasks, @@ -230,7 +230,7 @@ async def test_finalize_train_batch_supports_partial_finalize(self): self.assertEqual(result["finished_task_count"], 1) self.assertEqual(self.pipeline.process_chunk_calls[-1], [b"payload-0"]) - self.assertEqual(self.scheduler.abort_calls[-1]["batch_id"], 2) + self.assertEqual(self.scheduler.cleanup_calls[-1]["batch_id"], 2) self.assertIn("2", self.coordinator.discard_recorded_prefixes) self.assertNotIn(2, self.coordinator.pending_batches) @@ -326,7 +326,7 @@ async def test_abort_batch_marks_batch_aborted_and_evicts_it(self): await self.coordinator.abort_batch(4, reason="shutdown") - self.assertEqual(self.scheduler.abort_calls[0]["batch_id"], 4) + self.assertEqual(self.scheduler.cleanup_calls[0]["batch_id"], 4) self.assertEqual(self.coordinator.discard_recorded_prefixes[-1], "4") self.assertNotIn(4, self.coordinator.pending_batches) diff --git a/tests/explorer/scheduler_test.py b/tests/explorer/scheduler_test.py index e34dd6dfb47..6cb1b15218c 100644 --- a/tests/explorer/scheduler_test.py +++ b/tests/explorer/scheduler_test.py @@ -659,7 +659,7 @@ async def test_get_payload_results(self): _, exps = await collect_results(scheduler, batch_id=1, min_num=1, timeout=1) self.assertEqual(len(exps), 0) - # test _cleanup_batch_and_restart_runners: part I, no clear + # test cleanup_batch and runner restart: part I, no clear tasks = generate_tasks(3, timeout_num=1, timeout_seconds=3) scheduler.schedule(tasks, batch_id=2) statuses, exps = await collect_results( @@ -672,7 +672,7 @@ async def test_get_payload_results(self): ) self.assertEqual(len(statuses), 1) self.assertEqual(len(exps), 1) - # test _cleanup_batch_and_restart_runners: part II, clear + # test cleanup_batch and runner restart: part II, clear tasks = generate_tasks(3, timeout_num=1, timeout_seconds=3) scheduler.schedule(tasks, batch_id=3) statuses, exps = await collect_results(scheduler, batch_id=3, timeout=2) diff --git a/trinity/explorer/rollout_coordinator.py b/trinity/explorer/rollout_coordinator.py index 8891632caa6..ce2a460ad54 100644 --- a/trinity/explorer/rollout_coordinator.py +++ b/trinity/explorer/rollout_coordinator.py @@ -223,7 +223,7 @@ async def abort_batch( return self.logger.warning("Abort batch %s: %s", batch_id, reason) - await scheduler.abort_batch( + await scheduler.cleanup_batch( batch_id, return_partial_tasks=keep_partial_results, restart_runners=True, @@ -349,7 +349,7 @@ def _get_active_batch_state(self, batch_state: BatchState) -> BatchLifecycleStat async def _cleanup_train_batch_runtime(self, batch_state: BatchState) -> None: """Drop unfinished train work after a non-complete finalize result.""" scheduler = self._require_scheduler() - await scheduler.abort_batch( + await scheduler.cleanup_batch( batch_state.batch_id, return_partial_tasks=False, restart_runners=True, diff --git a/trinity/explorer/scheduler.py b/trinity/explorer/scheduler.py index 72e68423b9f..be3dc6ac931 100644 --- a/trinity/explorer/scheduler.py +++ b/trinity/explorer/scheduler.py @@ -628,7 +628,7 @@ def dynamic_timeout(self, timeout: Optional[float] = None) -> float: avg_time_per_task * self.config.explorer.dynamic_timeout.ratio, ) - async def _cleanup_batch( + async def cleanup_batch( self, batch_id: Union[int, str], return_partial_tasks: bool = False, @@ -695,7 +695,7 @@ async def _wait_for_batch_results( >= self.config.explorer.over_rollout.wait_after_min ): if clear_timeout_tasks: - await self._cleanup_batch( + await self.cleanup_batch( batch_id, return_partial_tasks=return_partial_tasks, restart_runners=False, @@ -757,7 +757,7 @@ async def _get_batch_payload_results( f"Timed out waiting for tasks at batch {batch_id} to complete after {timeout} seconds" ) if clear_timeout_tasks: - await self._cleanup_batch( + await self.cleanup_batch( batch_id, return_partial_tasks=return_partial_tasks, restart_runners=True, @@ -809,19 +809,6 @@ async def get_statuses( ) return statuses - async def abort_batch( - self, - batch_id: Union[int, str], - return_partial_tasks: bool = False, - restart_runners: bool = True, - ) -> None: - """Abort one batch and cleanup unfinished scheduler state.""" - await self._cleanup_batch( - batch_id, - return_partial_tasks=return_partial_tasks, - restart_runners=restart_runners, - ) - def has_step(self, batch_id: Union[int, str]) -> bool: return ( batch_id in self.completed_tasks From 82b1544f2e53dca5058d21f131a26db0b2f0989f Mon Sep 17 00:00:00 2001 From: "panxuchen.pxc" Date: Wed, 1 Jul 2026 13:45:55 +0800 Subject: [PATCH 70/84] clean coordinator --- tests/explorer/rollout_coordinator_test.py | 5 +++++ trinity/explorer/rollout_coordinator.py | 25 ++++++++++++---------- 2 files changed, 19 insertions(+), 11 deletions(-) diff --git a/tests/explorer/rollout_coordinator_test.py b/tests/explorer/rollout_coordinator_test.py index d649e6d6772..d54881b18b8 100644 --- a/tests/explorer/rollout_coordinator_test.py +++ b/tests/explorer/rollout_coordinator_test.py @@ -159,6 +159,11 @@ async def _init_scheduler(self): self.scheduler = self._test_scheduler + def _init_rollout_actors(self): + """Skip Ray actor resolution in unit tests.""" + + self._rollout_actors = {} + async def _discard_recorded_experiences(self, prefix: str) -> None: """Record cleanup requests without resolving real rollout actors.""" diff --git a/trinity/explorer/rollout_coordinator.py b/trinity/explorer/rollout_coordinator.py index ce2a460ad54..07ea4de6da5 100644 --- a/trinity/explorer/rollout_coordinator.py +++ b/trinity/explorer/rollout_coordinator.py @@ -64,19 +64,19 @@ def __init__( self.pending_batches: Dict[BatchId, BatchState] = {} self.running = False self.detailed_stats = getattr(getattr(config, "monitor", None), "detailed_stats", False) - # Lazily-resolved map of rollout engine_id -> rollout actor handle, for - # recording residual cleanup. + # Prepared map of rollout engine_id -> rollout actor handle, for + # scheduler construction and recording residual cleanup. self._rollout_actors: Dict[int, ActorHandle] = {} - def _resolve_rollout_actors(self) -> Dict[int, ActorHandle]: + def _init_rollout_actors(self) -> None: """Resolve each rollout engine's actor handle via named Ray actors. Mirrors ``Allocator.get_actor_name`` + ``ray.get_actor``: rollout model actors are named ``f"{explorer.name}_rollout_model_{engine_id}_0"`` - (node_id 0 holds the recording store). Cached after first resolution. + (node_id 0 holds the recording store). """ if self._rollout_actors: - return self._rollout_actors + return rollout_cfg = self.config.explorer.rollout_model name = self.config.explorer.name namespace = rollout_cfg.ray_namespace @@ -92,12 +92,13 @@ def _resolve_rollout_actors(self) -> Dict[int, ActorHandle]: % (actor_name, namespace) ) from exc self._rollout_actors = actors - return actors async def prepare(self) -> None: """Initialize the owned pipeline and scheduler.""" if self.running: return + if not self._rollout_actors and getattr(self.config, "mode", None) != "serve": + self._init_rollout_actors() if self.experience_pipeline is None: await self._init_experience_pipeline() if self.scheduler is None: @@ -116,18 +117,18 @@ async def shutdown(self) -> None: async def _init_experience_pipeline(self): """Create the experience pipeline owned by this coordinator actor.""" - if self.config.mode == "bench": + if getattr(self.config, "mode", None) == "bench": return None self.experience_pipeline = ExperiencePipeline(self.config) await self.experience_pipeline.prepare() async def _init_scheduler(self): """Create the scheduler owned by this coordinator.""" - if self.config.mode == "serve": + if getattr(self.config, "mode", None) == "serve": return self.scheduler = Scheduler( self.config, - rollout_actors=self._resolve_rollout_actors(), + rollout_actors=self._rollout_actors, ) await self.scheduler.start() @@ -319,9 +320,11 @@ async def _finalize_train_batch( async def _discard_recorded_experiences(self, prefix: str) -> None: """Delete recorded experiences matching a prefix from all rollout ranks.""" - actors = self._resolve_rollout_actors() results = await asyncio.gather( - *[actor.delete_experience_records.remote(prefix=prefix) for actor in actors.values()], + *[ + actor.delete_experience_records.remote(prefix=prefix) + for actor in self._rollout_actors.values() + ], return_exceptions=True, ) for result in results: From 98430b4087ab75cb63ebf7025d721de2aee79f6f Mon Sep 17 00:00:00 2001 From: pxc Date: Wed, 1 Jul 2026 14:34:39 +0800 Subject: [PATCH 71/84] block finished batch --- tests/buffer/memory_store_test.py | 39 +++++++++++++++++++++++++ trinity/buffer/store/base_store.py | 15 ++++++++++ trinity/buffer/store/memory_store.py | 23 +++++++++++++++ trinity/common/models/model.py | 22 ++++++++++++++ trinity/explorer/rollout_coordinator.py | 19 +++++++++--- 5 files changed, 114 insertions(+), 4 deletions(-) diff --git a/tests/buffer/memory_store_test.py b/tests/buffer/memory_store_test.py index e5fadef10e2..13e1502111b 100644 --- a/tests/buffer/memory_store_test.py +++ b/tests/buffer/memory_store_test.py @@ -112,6 +112,45 @@ def test_duplicate_sample_id_is_rejected(self): with self.assertRaises(ValueError): store.add("0/task_a/1", exp) + def test_blocked_prefix_drops_add_and_overwrite(self): + store = MemoryStore() + key = "0/task_a/0" + store.add(key, get_dummy_experience(1, request_id="pre")) + self.assertFalse(store.is_prefix_blocked("0")) + + # Real flow: block the batch, then delete its existing records. + store.block_prefix("0") + self.assertTrue(store.is_prefix_blocked("0")) + store.remove(key) + self.assertEqual(store.get(key), []) + + # A late add on a fresh key under the blocked batch is dropped. + store.add("0/task_a/1", get_dummy_experience(2, request_id="post")) + self.assertEqual(store.get("0/task_a/1"), []) + self.assertNotIn("0/task_a/1", store.keys()) + + # A late overwrite is also dropped: _drop_key is a no-op (records were + # already deleted) and add is blocked, so nothing reappears. + store.overwrite(key, get_dummy_experience(1, request_id="overwrite")) + self.assertEqual(store.get(key), []) + self.assertNotIn(key, store.keys()) + + def test_blocked_prefix_does_not_affect_other_batches(self): + store = MemoryStore() + store.block_prefix("0") + store.add("1/task_a/0", get_dummy_experience(1, request_id="other")) + self.assertEqual(len(store.get("1/task_a/0")), 1) + + def test_blocked_prefix_keeps_get_and_remove_working(self): + store = MemoryStore() + key = "0/task_a/0" + store.add(key, get_dummy_experience(2, request_id="keep")) + store.block_prefix("0") + # Reads and removes still work on already-stored records. + self.assertEqual(len(store.get(key)), 2) + self.assertEqual(len(store.remove(key)), 2) + self.assertEqual(store.get(key), []) + if __name__ == "__main__": unittest.main() diff --git a/trinity/buffer/store/base_store.py b/trinity/buffer/store/base_store.py index ed116f903c6..f1c1c760ce9 100644 --- a/trinity/buffer/store/base_store.py +++ b/trinity/buffer/store/base_store.py @@ -50,3 +50,18 @@ def remove(self, key: str) -> List[Experience]: @abstractmethod def keys(self) -> list[str]: """Return complete keys currently stored in insertion order.""" + + @abstractmethod + def block_prefix(self, prefix: str) -> None: + """Mark a batch prefix as blocked. + + Once a prefix is blocked, ``add`` and ``overwrite`` for any complete + key whose batch segment matches the prefix are silently dropped. + ``get`` and ``remove`` are unaffected. This is used to reject writes + that race in after a batch has been aborted/finalized and its records + deleted, so they cannot reappear as orphans. + """ + + @abstractmethod + def is_prefix_blocked(self, prefix: str) -> bool: + """Return whether the given batch prefix is blocked.""" diff --git a/trinity/buffer/store/memory_store.py b/trinity/buffer/store/memory_store.py index 241bfaad2cc..e4b875e7b22 100644 --- a/trinity/buffer/store/memory_store.py +++ b/trinity/buffer/store/memory_store.py @@ -1,11 +1,14 @@ """In-memory implementation of the experience store interface.""" +import logging from collections import OrderedDict from typing import Iterable, List from trinity.buffer.store.base_store import ExperienceUpdate, RecordStore from trinity.common.experience import Experience +_logger = logging.getLogger(__name__) + def parse_record_key(key: str) -> tuple[str, str, int]: """Parse a complete ``//`` store key. @@ -56,12 +59,24 @@ def __init__(self) -> None: self._batch_keys: dict[str, OrderedDict[str, None]] = {} self._task_keys: dict[tuple[str, str], OrderedDict[str, None]] = {} self._sample_to_key: dict[str, str] = {} + # batch prefixes whose writes should be silently dropped (aborted/ + # finalized batches); see ``block_prefix``. Only grows since batch_id + # is never reused. + self._blocked_batches: set[str] = set() def __len__(self) -> int: return sum(len(exps) for exps in self._records.values()) def add(self, key: str, exps: List[Experience]) -> None: batch, task, _ = self._parse_complete_key(key) # validate key format + if batch in self._blocked_batches: + _logger.debug( + "Dropping write to blocked batch '%s' (key=%s, %d exps).", + batch, + key, + len(exps), + ) + return if not exps: return @@ -151,6 +166,14 @@ def remove(self, key: str) -> List[Experience]: def keys(self) -> list[str]: return list(self._records.keys()) + def block_prefix(self, prefix: str) -> None: + """Mark a batch prefix as blocked; future ``add``/``overwrite`` are dropped.""" + self._blocked_batches.add(prefix) + + def is_prefix_blocked(self, prefix: str) -> bool: + """Return whether the given batch prefix is blocked.""" + return prefix in self._blocked_batches + @staticmethod def _parse_complete_key(key: str) -> tuple[str, str, int]: """Parse a complete store key; also usable as a key-format validator.""" diff --git a/trinity/common/models/model.py b/trinity/common/models/model.py index 7eade887ab1..4a1cc304184 100644 --- a/trinity/common/models/model.py +++ b/trinity/common/models/model.py @@ -231,6 +231,20 @@ async def delete_experience_records(self, prefix: str) -> None: """Remove recorded experiences matching a key or prefix.""" await self._drain_experience_records(prefix) + async def block_experience_records(self, prefix: str) -> None: + """Block future writes for the given batch prefix on this rollout rank. + + Sets the block flag before flushing the recorder so that any in-flight + experiences still queued in the recorder are dropped by ``MemoryStore`` + rather than written back as orphans. ``prefix`` is the batch segment + of the store key (``str(batch_id)``). + """ + recorder = getattr(self, "recorder", None) + if recorder is None: + return + recorder.store.block_prefix(prefix) + await recorder.flush() + def get_model_config(self) -> InferenceModelConfig: """Get the model configuration.""" return self.config @@ -942,6 +956,14 @@ async def delete_experience_records_async(self, prefix: str) -> None: raise ValueError("Recording delete requires an inference model actor.") await self.model.delete_experience_records.remote(prefix=prefix) + async def block_experience_records_async(self, prefix: str) -> None: + """Block future writes for the given batch prefix on the rollout actor.""" + if not self.enable_history: + raise ValueError("History recording is not enabled.") + if self.model is None: + raise ValueError("Recording block requires an inference model actor.") + await self.model.block_experience_records.remote(prefix=prefix) + async def shutdown(self) -> None: """Shutdown all underlying model actors cleanly.""" try: diff --git a/trinity/explorer/rollout_coordinator.py b/trinity/explorer/rollout_coordinator.py index ce2a460ad54..685e646a5ac 100644 --- a/trinity/explorer/rollout_coordinator.py +++ b/trinity/explorer/rollout_coordinator.py @@ -318,15 +318,26 @@ async def _finalize_train_batch( return self._finish_batch(batch_state, pipeline_metrics=pipeline_metrics) async def _discard_recorded_experiences(self, prefix: str) -> None: - """Delete recorded experiences matching a prefix from all rollout ranks.""" + """Block future writes and delete recorded experiences for a prefix. + + Blocking happens before deleting across all rollout ranks so that any + in-flight write that lands after the delete is dropped by the store + instead of reappearing as an orphan. The block flag persists on each + rollout actor (batch_id is never reused), so the prefix stays + unwritable for the lifetime of the process. + """ actors = self._resolve_rollout_actors() - results = await asyncio.gather( + block_results = await asyncio.gather( + *[actor.block_experience_records.remote(prefix=prefix) for actor in actors.values()], + return_exceptions=True, + ) + delete_results = await asyncio.gather( *[actor.delete_experience_records.remote(prefix=prefix) for actor in actors.values()], return_exceptions=True, ) - for result in results: + for result in [*block_results, *delete_results]: if isinstance(result, Exception): - self.logger.error("records delete from rollout actor failed: %s", result) + self.logger.error("records cleanup on rollout actor failed: %s", result) def _finish_batch( self, From 1146972c33ab4417d0c6b15abe50d75d3136ee79 Mon Sep 17 00:00:00 2001 From: pxc Date: Wed, 1 Jul 2026 19:56:31 +0800 Subject: [PATCH 72/84] update workflow doc --- .../source/tutorial/develop_workflow.md | 634 ++++++++++-------- .../source/tutorial/trinity_configs.md | 4 - .../source_zh/tutorial/develop_workflow.md | 530 ++++++++------- .../source_zh/tutorial/trinity_configs.md | 4 - tests/common/vllm_test.py | 8 +- tests/explorer/workflow_test.py | 1 - trinity/common/config.py | 36 +- trinity/common/config_validator.py | 35 +- trinity/common/models/allocator.py | 20 +- trinity/common/models/model.py | 14 +- trinity/common/models/sglang_model.py | 6 +- trinity/common/models/vllm_model.py | 4 - trinity/common/workflows/workflow.py | 13 +- trinity/explorer/proxy/service.py | 3 +- trinity/explorer/workflow_runner.py | 2 +- 15 files changed, 709 insertions(+), 605 deletions(-) diff --git a/docs/sphinx_doc/source/tutorial/develop_workflow.md b/docs/sphinx_doc/source/tutorial/develop_workflow.md index 11c2d01d7c8..e77694a1dda 100644 --- a/docs/sphinx_doc/source/tutorial/develop_workflow.md +++ b/docs/sphinx_doc/source/tutorial/develop_workflow.md @@ -1,58 +1,58 @@ (Workflows)= ## Workflow Development Guide -In Trinity-RFT, workflows are the core components that define the interaction between Agents and Environments. -A qualified workflow needs to use a model to complete the specified task and obtain feedback information (reward) from the environment. Below are the steps to create a new workflow: +In Trinity-RFT, a workflow (Workflow) is the core component that defines the interaction between an Agent and its Environment. +A qualified workflow needs to use the model being trained to complete a specified task and obtain feedback (reward) from the environment. This section introduces how to develop a new workflow. --- ### Step 0: Basic Concepts -Before starting development, it's important to understand several core concepts: +Before starting development, it is important to understand the following core concepts: ```{mermaid} flowchart LR - A([Task]) & B([Model]) --> C[Workflow] - C --> D([Experience]) + A([Task]) --> C[Workflow] + C -- "call OpenAI API" --> B([Rollout Model]) + B -- "auto recording" --> D([Experience]) + C -- "update_reward" --> D ``` -- **Task** ({class}`trinity.common.workflows.Task`): Represents a data structure that contains all the information needed for a single run of the workflow. Commonly provided by the training dataset, each sample in the dataset is converted into a `Task` instance. The content of the `Task` varies depending on the task type: - - **Math problems**: A `Task` contains the problem description and the golden answer. - - **Programming scenarios**: A `Task` includes the problem description, test cases, runtime environment, and other complex information. +- **Task** ({class}`trinity.common.workflows.Task`): A structured data instance that contains the information needed for a single run of the workflow. It is usually provided by the training dataset, where each sample is converted into a `Task` instance. The contents of a `Task` vary by task type: + - **Math problem**: contains the question and the answer. + - **Programming scenario**: contains complex information such as the problem description, test cases, and the execution environment. -- **Model** ({class}`trinity.common.models.model.ModelWrapper`): The model being trained. The workflow uses this model to generate responses based on the task. Trinity-RFT will provide the model instance to initialize the workflow. +- **Rollout Model** ({class}`trinity.common.models.model.ModelWrapper`): The model being trained. The workflow creates its own OpenAI client from the `base_url` and `api_key` exposed by the model to call the inference API; while responding, the model **automatically records** the generation process and turns it into trainable `Experience` objects, so the workflow does not need to construct them manually. -- **Workflow** ({class}`trinity.common.workflows.Workflow`): It defines the interaction flow between Agents and Environments. It uses the `Task` to initialize itself and uses the `Model` to generate responses. Different from general Agent Applications, a `Workflow` also needs to calculate rewards based on the environment's feedback. Trinity-RFT provides several built-in workflows, including: - - `MathWorkflow` ({class}`trinity.common.workflows.MathWorkflow`): For math scenarios, submits problems to LLM, parses LLM responses, and calculates scores (rewards). - - `WebShopWorkflow` ({class}`trinity.common.workflows.WebShopWorkflow`): For webshop scenarios, it contains multi-turn interaction with environment. - - `AgentScopeReActWorkflow` ({class}`trinity.common.workflows.AgentScopeReActWorkflow`): It directly uses a pre-implemented ReActAgent (based on AgentScope) to solve tasks. - -- **Experience** ({class}`trinity.common.experience.Experience`): The output of running a `Workflow`. The number and structure of `Experience` depend on the specific workflow. For example, for common PPO/GRPO algorithms, `Experience` includes lists of token IDs, action masks (identifying which tokens were generated by the LLM), log probabilities, rewards, etc. +- **Workflow** ({class}`trinity.common.workflows.WorkflowBase`): Defines the interaction flow between Agent and Environment. A `Workflow` initializes itself from the information provided by the `Task` and uses the Rollout Model to execute the predefined interaction flow. Unlike a regular agent application, a workflow must also compute a reward signal to guide training, and writes the reward back onto the automatically recorded `Experience` via the `update_reward` method. +- **Experience** ({class}`trinity.common.experience.Experience`): The data unit needed for training. `Experience` objects are produced automatically by the Rollout Model during inference; their number and internal format depend on the training algorithm used. For example, for common PPO/GRPO algorithms, an `Experience` contains a token ID list, an action mask (indicating which tokens were generated by the LLM), per-token log probabilities (logprobs), a reward signal, etc. A workflow does not need to, and should not, construct `Experience` objects manually. --- -### Step 1: Prepare Task Dataset +### Step 1: Prepare the Task Dataset + +The task dataset is loaded through the `buffer.explorer_input.taskset` field in the YAML config file. +To handle differences in `Task` contents, Trinity-RFT provides a unified `Task` interface with the following fields: -The task dataset is loaded via the `buffer.explorer_input.taskset` configuration entry in your YAML config file. -To handle differences in `Task` contents, Trinity-RFT provides a unified `Task` interface containing the following fields. +- **`workflow`** (`str`): The registered name of your workflow class. You can specify it via `buffer.explorer_input.taskset.default_workflow_type` in the YAML config. +- **`raw_task`** (`Dict`): The raw data record, stored as a `Dict`. For highly customized workflows, you can initialize the `Workflow` instance directly from `raw_task` without relying on the fields below. -- **`workflow`** (`str`): The registered name of your workflow class. You can specify it in `buffer.explorer_input.taskset.default_workflow_type` of your YAML config file. -- **`reward_fn`** (`Optional[str]`): The registered name of your reward function. You can specify it in `buffer.explorer_input.taskset.default_reward_fn_type`. Note that some workflows already include built-in reward calculation; in such cases, you can omit this field. -- **`raw_task`** (`Dict`): A record of raw data in `Dict` format. For highly customized workflow, you can directly use `raw_task` to initialize your `Workflow` instance without relying on the following fields. -- **`format_args`** ({class}`trinity.common.config.FormatConfig`): Parameters to facilitate the construction of `Workflow` instances. For example, the `prompt_key` and `response_key` can be used to get the prompt and response from `raw_task`. These settings come from the YAML configuration file and can be set in `buffer.explorer_input.task_set.format`. -- **`rollout_args`** ({class}`trinity.common.config.GenerationConfig`): Parameters that control the rollout process, such as `temperature`. This field also comes from the YAML configuration file and can be set in `buffer.explorer_input.task_set.rollout_args`. -- **`workflow_args`** (`Dict`): A dictionary of parameters to facilitate the construction of `Workflow` instances. Provides more flexibility than `format_args` and `rollout_args` by using a dictionary. This field also comes from the YAML configuration file and can be set in `buffer.explorer_input.task_set.workflow_args`. Normally, you do not need to set this field. +The following fields are all optional and usually do not need to be set: +- **`reward_fn`** (`Optional[str]`): The registered name of your reward function. You can specify it via `buffer.explorer_input.taskset.default_reward_fn_type`. Note that some workflows have built-in reward computation; in that case this field can be omitted. +- **`format_args`** ({class}`trinity.common.config.FormatConfig`): Parameters that help construct the `Workflow` instance. For example, `prompt_key` and `response_key` can be used to extract the prompt and response from `raw_task`. These come from the YAML config and can be set under `buffer.explorer_input.task_set.format`. +- **`rollout_args`** ({class}`trinity.common.config.GenerationConfig`): Parameters controlling the rollout process, such as `temperature`. Also from the YAML config, settable under `buffer.explorer_input.task_set.rollout_args`. +- **`workflow_args`** (`Dict`): A parameter dict for constructing the `Workflow` instance, more flexible than `format_args` and `rollout_args`. Also from the YAML config, settable under `buffer.explorer_input.task_set.workflow_args`. Usually you do not need to set this field. ```{tip} -`workflow`, `workflow_args` and `raw_task` provide different levels of customization. +`workflow`, `workflow_args`, and `raw_task` provide different levels of customization. -- `workflow` provides the global settings for all tasks that uses the same workflow. (Global Level) -- `workflow_args` can be set for each task dataset, allowing different task datasets using the same workflow to behave differently. (Dataset Level) -- `raw_task` provides the ability to customize the behavior of each task, which is most flexible. (Data Sample Level) +- `workflow` provides global settings for all tasks using the same workflow. (Global level) +- `workflow_args` can be set per task dataset, allowing different datasets that share the same workflow to behave differently. (Dataset level) +- `raw_task` provides per-task customization, the most flexible level. (Sample level) ``` -In the math problem scenario, the `Task` dataset can be a `jsonl` file, where each line contains JSON with `question` and `answer` fields representing the problem description and standard answer, respectively. For example: +For a math scenario, the `Task` dataset can be a `jsonl` file where each line is a JSON with `question` and `answer` fields, representing the problem description and the ground-truth answer. For example: ```json {"question": "1+1=", "answer": "2"} @@ -60,14 +60,14 @@ In the math problem scenario, the `Task` dataset can be a `jsonl` file, where ea ... ``` -Example configuration snippet: +Example config snippet: ```yaml # some config buffer: explorer_input: taskset: - default_workflow: "math_workflow" + default_workflow_type: "math_workflow" path: ${oc.env:TRINITY_TASKSET_PATH} format: prompt_key: "question" @@ -77,16 +77,16 @@ buffer: # some other configs ``` -In this example, each task object's `raw_task` is a `Dict` with two keys (`question` and `answer`). The `MathWorkflow` uses the `prompt_key` and `response_key` to extract the question and answer from the `raw_task` and use the `rollout_args` to generate the response. +In this example, the `raw_task` of each task object is a `Dict` with two keys (`question` and `answer`). `MathWorkflow` uses `prompt_key` and `response_key` to extract the question and answer from `raw_task`, and uses `rollout_args` to generate responses. --- -### Step 2: Implement a New Workflow +### Step 2: Implement the Workflow -The `Workflow` base class interface is as follows: +To implement a new workflow, you need to subclass the `WorkflowWithRecording` base class: ```python -class Workflow(ABC): +class WorkflowWithRecording(WorkflowBase): def __init__( self, @@ -95,55 +95,68 @@ class Workflow(ABC): model: ModelWrapper, auxiliary_models: Optional[List[ModelWrapper]] = None, ): - self.task = task - self.model = model - self.auxiliary_model_wrappers = auxiliary_models - self.auxiliary_models = ... # OpenAI clients auto-derived from ModelWrapper - self.logger = get_logger(__name__) # built-in logger for runtime monitoring + """Initialize the workflow.""" - @abstractmethod - def run(self) -> List[Experience]: - """Run the workflow and return a list of Experiences.""" -``` + async def run_async(self) -> Metrics: + """Run the workflow and return a Metric dict.""" + # you need to implement this method -#### Initialize Your Workflow + @property + def base_url(self) -> str: + """Return the base_url of the rollout model.""" -During initialization, `Workflow` receives the following parameters: + @property + def api_key(self) -> str: + """Return the api_key of the rollout model.""" -- `task`({class}`trinity.common.workflows.Task`): A single data item from the task dataset. -- `model`({class}`trinity.common.models.model.ModelWrapper`): The model being trained, which provides an interface similar to OpenAI, capable of receiving a list of conversation messages and returning content generated by the LLM (including reply text `response_text`, full sequence token ids `tokens`, prompt part token length `prompt_length`, and a list of output token logprobs `logprobs`). -- `auxiliary_models`(`List[ModelWrapper]`): A list of auxiliary model wrappers. You can access OpenAI clients via `self.auxiliary_models` (auto-derived based on workflow's `is_async`). + @property + def model_name(self) -> str: + """Return the model_name of the rollout model.""" + + async def update_reward( + self, + reward: float, + info: Optional[Dict] = None, + ): + """Write the reward back onto the model's automatically recorded Experience, with optional extra info.""" -```{tip} -You can switch to using the OpenAI API by setting `explorer.rollout_model.enable_openai_api` to `true` in your config file and calling `model.get_openai_client()` to get an `openai.OpenAI` instance in your workflow. -And the `model` field when calling openai API can be obtained via `openai_client.models.list().data[0].id` or `openai_client.model_path`. ``` -Here's an example of initializing a simple workflow using only `raw_task` and `rollout_args`. In more complex cases, you can use the `format_args` for further customization. +#### Initialize your workflow + +`WorkflowWithRecording` accepts the following initialization parameters: + +- `task` ({class}`trinity.common.workflows.Task`): A single task from the dataset. +- `model` ({class}`trinity.common.models.model.ModelWrapper`): The rollout model being trained. You can directly use the `base_url`, `api_key`, and `model_name` properties of `WorkflowWithRecording` to create an OpenAI client and call the model's inference API. +- `auxiliary_models` (`List[ModelWrapper]`): A list of `ModelWrapper` instances for auxiliary models. Each element also exposes `base_url`, `api_key`, and `model_name`, and can be used directly to create an OpenAI client (see [LLM-as-a-judge support](#llm-as-a-judge-support)). + +Here is an initialization example for a simple workflow. In `__init__` we create an async OpenAI client from `base_url` and `api_key`, and read out the model name: ```python -class ExampleWorkflow(Workflow): +import openai +from trinity.common.workflows import WorkflowWithRecording - def __init__(self, task: Task, model: ModelWrapper, auxiliary_models: List): +class ExampleWorkflow(WorkflowWithRecording): + + def __init__(self, *, task: Task, model: ModelWrapper, auxiliary_models: List = None): super().__init__(task=task, model=model, auxiliary_models=auxiliary_models) self.question = task.raw_task.get("question") self.answer = task.raw_task.get("answer") self.rollout_args = task.rollout_args - # Optional: If you want to use OpenAI API in your workflow - # self.openai_client = self.model.get_openai_client() + # create the OpenAI client from base_url and api_key + self.client = openai.AsyncOpenAI(base_url=self.base_url, api_key=self.api_key) ``` -#### Implementing the `run` method +#### Implement the `run_async` method + +`run_async` is the core method of the workflow. It takes no input parameters and returns a `Metrics` dict. -The `run` method is the core of your workflow. It returns a list of `Experience`. -Below is a simple implementation for a math workflow. +The workflow's responsibility is: call the model to complete the agent task, compute the reward, write the reward back onto the automatically recorded `Experience` via `update_reward`, and finally return a metric for monitoring. -We first call the model to generate a response using the provided question and rollout arguments. -Then we calculate the reward for each response using the `calculate_reward` function. -Finally, we construct a list of `Experience` with the responses and rewards and return it. +Here is a simple implementation of a math workflow. We first use the OpenAI client to generate an answer, then compute the reward and write it back: ```python -class ExampleWorkflow(Workflow): +class ExampleWorkflow(WorkflowWithRecording): # the __init__ function @@ -153,10 +166,11 @@ class ExampleWorkflow(Workflow): else: return 0.0 - def run(self) -> List[Experience]: - # call the model to generate multiple responses - responses = self.model.chat( - [ + async def run_async(self) -> Metrics: + # call the model to generate a response + resp = await self.client.chat.completions.create( + model=self.model_name, + messages=[ { "role": "user", "content": f"Question:\n{self.question}", @@ -164,114 +178,75 @@ class ExampleWorkflow(Workflow): ], temperature=self.rollout_args.temperature, ) - response = responses[0] # there is only one response - reward: float = self.calculate_reward(response.response_text, self.answer) - return [ - Experience( - tokens=response.tokens, - prompt_length=response.prompt_length, - reward=reward, - logprobs=response.logprobs, - ) - ] + response_text = resp.choices[0].message.content + # compute the reward and write it back onto the automatically recorded Experience + reward: float = self.calculate_reward(response_text, self.answer) + await self.update_reward(reward) + # return the metric to monitor + return {"example/reward": reward} +``` + +```{note} +1. The rollout model automatically records the training data produced by each `chat.completions.create` call and turns it into `Experience` objects. `update_reward` writes the reward precisely onto the `Experience` objects produced by this run. +2. For workflows with multi-turn interactions, `update_reward` writes the reward onto all `Experience` objects produced by this run. +3. The `Metrics` dict returned by `run_async` is only used for runtime monitoring and log display. ``` -#### Registering Your Workflow +#### Register your workflow -Register your workflow using the `default_mapping` in `trinity/common/workflows/__init__.py`. -Ensure the name does not conflict with existing workflows. +So that Trinity-RFT can find your workflow by name from the config file, you need to register it in the `WORKFLOWS` registry. The recommended way is to register with a decorator: ```python -WORKFLOWS = Registry( - "workflows", - default_mapping={ - "example_workflow": "trinity.common.workflows.workflow.ExampleWorkflow", - }, -) +from trinity.common.workflows import WORKFLOWS, WorkflowWithRecording + +@WORKFLOWS.register_module(name="example_workflow") +class ExampleWorkflow(WorkflowWithRecording): + ... ``` -#### Performance Optimization +You can also register it directly, or add an entry `"example_workflow": "path.to.module.ExampleWorkflow"` to the `default_mapping` in `trinity/common/workflows/__init__.py`. + +#### Performance tuning -##### Avoid Re-initialization +For more complex workflows, re-initializing each time brings extra overhead. In that case, you can set the `can_reset` class attribute and implement the `reset` method to avoid repeated initialization. -For heavy workflows, re-initializing every time can incurs extra computational costs. -In this case, you can set the `can_reset` property and implement `reset` method to avoid re-initialization. +Note that in the `reset` method you must overwrite the workflow's `task` attribute with the input `task`, and update the API key of the model and the client with `task.api_key`. -The `can_reset` is a class property that indicates whether the workflow supports resetting. +> Trinity-RFT internally uses `api_key` to distinguish experiences produced by different tasks. If you do not update the API key, experiences from different tasks may be misclassified, causing rewards to be written back to the wrong records. -The `reset` method accepts a `Task` parameter and resets the workflow's internal state based on the new task. +Here is a simple example: ```python -class ExampleWorkflow(Workflow): +class ExampleWorkflow(WorkflowWithRecording): can_reset: bool = True # some code # ... def reset(self, task: Task): + self.task = task + self.model.set_api_key(task.api_key) + self.client.api_key = task.api_key self.question = task.raw_task.get("question") self.answer = task.raw_task.get("answer") ``` -##### Support Batch Inference - -In many popular RL algorithms, multiple runs of the same task are required (e.g., GRPO). In such scenarios, you can directly use batch inference to obtain multiple responses for a single question to improve efficiency. -For this case, you can implement the `can_repeat` property and `set_repeat_times` method. - -The `can_repeat` is a class property that indicates whether the workflow supports multiple executions within the `run` method. - -The `set_repeat_times` method accepts two parameters: `repeat_times` specifies the number of times to execute within the `run` method, and `run_id_base` is an integer used to identify the first run ID in multiple runs (this parameter is used in multi-turn interaction scenarios; for tasks that can be completed with a single model call, this can be ignored). +#### Complete code example ```python -class ExampleWorkflow(Workflow): - can_repeat: bool = True - # some code - - def set_repeat_times(self, repeat_times, run_id_base): - self.repeat_times = repeat_times - self.run_id_base = run_id_base - - def run(self) -> List[Experience]: - # call the model to generate multiple responses - responses = self.model.chat( - [ - { - "role": "user", - "content": f"Question:\n{self.question}", - } - ], - n=self.repeat_times, # run multiple times in one call - temperature=self.rollout_args.temperature, - ) - experiences = [] - for response in responses: - # calculate reward - reward: float = self.calculate_reward(response.response_text, self.answer) - # construct Experience - experiences.append( - Experience( - tokens=response.tokens, - prompt_length=response.prompt_length, - reward=reward, - logprobs=response.logprobs, - ) - ) - return experiences -``` - - -#### Full Code Example +import openai +from trinity.common.workflows import WORKFLOWS, WorkflowWithRecording -```python -class ExampleWorkflow(Workflow): +@WORKFLOWS.register_module(name="example_workflow") +class ExampleWorkflow(WorkflowWithRecording): can_reset: bool = True - can_repeat: bool = True - def __init__(self, task: Task, model: ModelWrapper, auxiliary_models: List): + def __init__(self, *, task: Task, model: ModelWrapper, auxiliary_models: List = None): super().__init__(task=task, model=model, auxiliary_models=auxiliary_models) self.question = task.raw_task.get("question") self.answer = task.raw_task.get("answer") self.rollout_args = task.rollout_args + self.client = openai.AsyncOpenAI(base_url=self.base_url, api_key=self.api_key) def calculate_reward(self, response: str, truth: str) -> float: if response == truth: @@ -279,47 +254,35 @@ class ExampleWorkflow(Workflow): else: return 0.0 - def run(self) -> List[Experience]: - # call the model to generate multiple responses - responses = self.model.chat( - [ + async def run_async(self) -> Metrics: + resp = await self.client.chat.completions.create( + model=self.model_name, + messages=[ { "role": "user", "content": f"Question:\n{self.question}", } ], - n=self.rollout_args.n, temperature=self.rollout_args.temperature, ) - experiences = [] - for response in responses: - # calulcate reward - reward: float = self.calculate_reward(response.response_text, self.answer) - # construct Experience - experiences.append( - Experience( - tokens=response.tokens, - prompt_length=response.prompt_length, - reward=reward, - logprobs=response.logprobs, - ) - ) - return experiences + response_text = resp.choices[0].message.content + reward: float = self.calculate_reward(response_text, self.answer) + await self.update_reward(reward) + return {"example/reward": reward} def reset(self, task: Task): + self.task = task + self.model.set_api_key(task.api_key) + self.client.api_key = task.api_key self.question = task.raw_task.get("question") self.answer = task.raw_task.get("answer") - - def set_repeat_times(self, repeat_times, run_id_base): - self.repeat_times = repeat_times - self.run_id_base = run_id_base ``` --- -### Step 3: Use Your Workflow +### Step 3: Use your workflow -After implementing and registering your workflow, you need to update the configuration file to set the `default_workflow_type` in the `buffer.explorer_input.taskset` domain to the newly registered `Workflow` name. +After implementing and registering the workflow, you can use it by setting `default_workflow_type` under `buffer.explorer_input.taskset` in the config file to your workflow name. For example: ```yaml buffer: @@ -331,86 +294,21 @@ buffer: # Other fields ``` -Now you can run your workflow in Trinity-RFT using the command: +Now you can run your workflow in Trinity-RFT with: -``` +```bash trinity run --config ``` --- -### Advanced Features - -#### Async Support - -The example above mainly targets synchronous mode. If your workflow needs to use asynchronous methods (e.g., asynchronous API), you can set `is_async` to `True`, then implement the `run_async` method. In this case, you no longer need to implement the `run` method, and the initialization parameter `auxiliary_models` will also change to `List[openai.AsyncOpenAI]`, while other methods and properties remain changed. - -```python -class ExampleWorkflowAsync(Workflow): - - is_async: bool = True +### LLM-as-a-judge Support - async def run_async(self) -> List[Experience]: - # your async code here +LLM-as-a-judge is a common reward computation method, especially suited for open-ended tasks (such as programming, writing, etc.). In such scenarios, the workflow needs an extra LLM to evaluate answer quality and compute the reward signal. - # no need to implement run() method -``` - -#### Using OpenAI API - -Trinity-RFT provides an option to use the OpenAI API for model inference. You can enable this feature by setting `explorer.rollout_model.enable_openai_api` to `true` in your configuration file. This allows you to obtain an `openai.OpenAI` instance via the `get_openai_client` method of the model instance provided by Trinity-RFT. +To support this, Trinity-RFT provides the Auxiliary Models mechanism. Auxiliary models are a set of models not involved in training; the workflow can use them to assist with the task, for example as a judge that computes the reward. -Additionally, since the OpenAI API does not provide all the data required for training, you also need to set `explorer.rollout_model.enable_history` to `true`. This lets the framework automatically record data that can be used for training and convert it into a list of `Experience`. You can extract these experiences using the `extract_experience_from_history` method. - -```yaml -# example config snippet -explorer: - rollout_model: - enable_openai_api: true - enable_history: true - # Other fields -``` - -```python -class ExampleWorkflow(Workflow): - - def __init__(self, task: Task, model: ModelWrapper, auxiliary_models: List): - super().__init__(task=task, model=model, auxiliary_models=auxiliary_models) - self.model = model - self.client: openai.OpenAI = self.model.get_openai_client() - # or async client - # self.client: openai.AsyncOpenAI = self.model.get_openai_async_client() - self.agent = MyAgent(openai_client=self.client) - - def calculate_reward(self, response: str) -> float: - # your reward calculation logic - - def run(self) -> List[Experience]: - # run your agent - response = self.agent.run() - # calculate reward - reward = self.calculate_reward(response) - # extract experiences from history recorded in self.model - experiences = self.model.extract_experience_from_history() - for exp in experiences: - exp.reward = reward - return experiences -``` - - -```{tip} -1. Currently, the OpenAI API will only automatically record calls to `openai.OpenAI.chat.completions.create` and `openai.AsyncOpenAI.chat.completions.create`, and convert them into `Experience` objects. Streaming output is not supported. -2. When calling `chat.completions.create`, the `model` field can be obtained via `openai_client.models.list().data[0].id` or `openai_client.model_path`. -3. For more complex workflow examples using the OpenAI API, refer to [ReAct Agent Training](./example_react.md). -``` - -#### LLM-as-a-judge Support - -LLM-as-a-judge is a common reward calculation method, especially suitable for open-ended tasks (such as programming, writing, etc.). In these scenarios, the Workflow needs to leverage an additional LLM to evaluate the answer quality and compute the reward signal. - -To support this, Trinity-RFT provides an Auxiliary Models mechanism. Auxiliary models are a set of models not involved in training; the Workflow can use these models to assist with tasks, such as acting as a judge to calculate rewards. - -You can specify one or more auxiliary models in the configuration file via the `explorer.auxiliary_models` field. For example: +You can specify one or more auxiliary models via the `explorer.auxiliary_models` field in the config file. For example: ```yaml explorer: @@ -431,26 +329,23 @@ explorer: max_model_len: 16384 ``` -Note that each auxiliary model will independently occupy `tensor_parallel_size * engine_num` GPUs. Please configure according to your hardware resources. After enabling auxiliary models, the number of GPUs available to the Trainer is the total GPU count minus those occupied by all auxiliary models and the inference model being trained (`rollout_model`). +Note that each auxiliary model independently occupies `tensor_parallel_size * engine_num` GPUs; configure them reasonably according to your hardware. After enabling auxiliary models, the GPUs available to the Trainer equal the total GPUs minus those occupied by all auxiliary models and the rollout model (`rollout_model`). -The auxiliary models specified in the configuration file will automatically activate the OpenAI API and pass the corresponding `openai.OpenAI` or `openai.AsyncOpenAI` instances (depending on the `is_async` setting) to the `auxiliary_models` parameter of the `Workflow` initialization method. For example: +The auxiliary models specified in the config are passed to the `auxiliary_models` parameter of the `Workflow` initializer as a list of `ModelWrapper` instances. Each `ModelWrapper` also exposes `base_url`, `api_key`, and `model_name`; it is recommended to create an OpenAI client directly from them to access the auxiliary model: ```python -class MyWorkflow(Workflow): - def __init__( - self, - *, - task: Task, - model: ModelWrapper, - auxiliary_models: Optional[List[ModelWrapper]] = None, - ): +class MyWorkflow(WorkflowWithRecording): + def __init__(self, *, task, model, auxiliary_models=None): super().__init__(task=task, model=model, auxiliary_models=auxiliary_models) - self.judge_model = self.auxiliary_models[0] # OpenAI client auto-derived from ModelWrapper + self.judge = self.auxiliary_models[0] # ModelWrapper + self.judge_client = openai.AsyncOpenAI( + base_url=self.judge.base_url, api_key=self.judge.api_key + ) - def run(self) -> List[Experience]: - response = self.do_something() - reward_response = self.judge_model.chat.completions.create( - model=self.judge_model.model_path, + async def run_async(self) -> Metrics: + response = await self.do_something() + reward_response = await self.judge_client.chat.completions.create( + model=self.judge.model_name, messages=[ { "role": "system", @@ -458,74 +353,66 @@ class MyWorkflow(Workflow): }, { "role": "user", - "content": f"Question:\n{self.task.raw_task['question']}\nAnswer:\n{response.response_text}\nPlease give a score from 0 to 1.", + "content": f"Question:\n{self.task.raw_task['question']}\nAnswer:\n{response}\nPlease give a score from 0 to 1.", }, ], temperature=0.0, max_tokens=10, ) - # Parse the reward score + # parse the reward score reward = float(reward_response.choices[0].message.content.strip()) - return [ - Experience( - tokens=response.tokens, - prompt_length=response.prompt_length, - reward=reward, - logprobs=response.logprobs, - ) - ] + await self.update_reward(reward, info={"source": "llm_as_a_judge"}) + return {"my_workflow/judge_reward": reward} ``` - #### Debug Mode -During Workflow development, repeatedly launching the full training process for testing is time-consuming and inefficient. To address this, Trinity-RFT provides a Debug Mode for developers. This mode leverages a pre-launched inference model to quickly run specified workflows and obtain results, avoiding repeated model loading and initialization delays, and significantly improving development efficiency. The process is illustrated below: +During workflow development, repeatedly launching the full training pipeline for testing is time-consuming and inefficient. To help, Trinity-RFT provides a debug mode. By starting the inference model in advance, this mode can quickly run a specified workflow and obtain results, avoiding repeated waits for model loading and initialization, greatly improving development efficiency. The flow is as follows: ```{mermaid} flowchart LR - A[Start Inference Model] --> B[Debug Workflow] - B --> C[Check Experiences] + A[Start inference model] --> B[Debug Workflow] + B --> C[Inspect Experience] C --> B ``` -To start the inference model, use the following command: +The command to start the inference model is: ```bash trinity debug --config --module inference_model ``` -Here, `` is the path to a YAML configuration file, which should follow the same format as the one used by the `trinity run` command. The `explorer.rollout_model` and `explorer.auxiliary_models` fields in the config will be loaded to initialize the inference model. +Here `config_file_path` is the path to a YAML config file, in the same format as the one used by `trinity run`. The `explorer.rollout_model` and `explorer.auxiliary_models` fields in the config are loaded to initialize the inference model. -Once started, the model will keep running and wait for debug instructions; it will not exit automatically. You can then run the following command in another terminal to debug your workflow: +After starting, the model keeps running and waits for debug commands; it does not exit automatically. You can then run the following command in another terminal to debug the workflow: ```bash trinity debug --config --module workflow --output-dir [--plugin-dir ] [--enable-profiling] [--disable-overwrite] ``` -- ``: Path to the YAML configuration file, usually the same as used for starting the inference model. -- ``: Directory to save the debug output. If not specified, the output will be saved to the `debug_output` in the current working directory. -- `` (optional): Path to the plugin directory. If your workflow or reward function modules are not built into Trinity-RFT, you can specify this parameter to load custom modules. -- `--enable-profiling` (optional): Enable performance profiling using [viztracer](https://github.com/gaogaotiantian/viztracer). -- `--disable-overwrite` (optional): Disable overwriting the output directory. If the directory is not empty, it will automatically change to a new directory with a timestamp suffix (e.g., `debug_output_20251203211200`) to avoid overwriting existing data. +- ``: Path to the YAML config file, usually the same as the one used to start the inference model. +- ``: Directory for debug output. If not specified, output is saved under `debug_output` in the current working directory. +- `` (optional): Plugin directory path. If your workflow or reward function modules are not built into Trinity-RFT, you can load custom modules via this parameter. +- `--enable-profiling` (optional): Enable profiling, using [viztracer](https://github.com/gaogaotiantian/viztracer) to profile the workflow run. +- `--disable-overwrite` (optional): Disable output directory overwriting. If the specified folder is non-empty, a new directory with a timestamp suffix (e.g. `debug_output_20251203211200`) is created automatically to avoid overwriting existing data. -During debugging, the `buffer.explorer_input.taskset` field in the config will be loaded to initialize the workflow's required task dataset and instance. Note that Debug Mode only reads the first sample in the dataset for testing. After running the above command, the workflow's return experiences will be written to the `experiences.db` file in the specified output directory. Additionally, the metrics will be printed in the terminal for easy inspection. +During debugging, the `buffer.explorer_input.taskset` field in the config is loaded to initialize the task dataset and instances for the workflow. Note that debug mode only reads the first sample of the dataset for testing. After running the command above, the experiences produced by the workflow are written to `experiences.db` in the specified output directory, and the metrics recorded during the run are printed to the terminal for inspection. ```bash trinity debug --config --module viewer --output-dir --port 8502 ``` -This command launches the Experience Viewer at `http://localhost:8502` to visualize the experiences generated during debugging. You can inspect the generated experiences in a user-friendly interface. -Note that the viewer reads experiences from the `experiences.db` file in the specified output directory, so ensure that you have successfully run the workflow debug command beforehand and use the same output directory. +This command launches an Experience Viewer at `http://localhost:8502` to visualize the experiences generated during debugging. You can inspect the generated experiences in a user-friendly interface. Note that the Viewer reads experiences from `experiences.db` in the specified output directory, so make sure you have successfully run the workflow debug command and replaced `` with the actual output directory. -When debugging is complete, you can terminate the inference model by pressing `Ctrl+C` in its terminal. +When debugging is done, press `Ctrl+C` in the inference model terminal to stop the model. -#### Runtime Monitoring +#### Runtime monitoring -The debug mode above provides the ability to quickly test and validate workflow implementations. However, during actual training, you may want to monitor the workflow's runtime behavior in real-time to ensure it operates as expected. To facilitate this, Trinity-RFT offers monitoring capabilities based on the logger system. The base `Workflow` class includes a built-in `logger` that you can use to log important runtime information. +In the debug mode above, you can quickly test and validate your workflow implementation. However, during actual training, you may want to monitor the workflow's running state in real time to ensure it works as expected. To support this, Trinity-RFT provides a log-based monitoring feature. The `WorkflowWithRecording` base class has a built-in logger; you can use it to record important runtime information. ```python -class Workflow(ABC): +class WorkflowWithRecording(WorkflowBase): def __init__( self, @@ -538,21 +425,20 @@ class Workflow(ABC): self.logger = get_logger(__name__) # built-in logger for runtime monitoring ``` -Different from standard Python loggers, this built-in logger is configured to output logs to both the console and a file under the `////log` directory. This allows you to monitor the workflow's runtime status during training conveniently. All workflow subclasses inherit this logger, so you can directly use it in your custom workflow implementations with `self.logger`. +This built-in logger writes logs to the console and to files under `////log`. This makes it convenient to monitor the workflow's state during training. Since all Workflow subclasses inherit this logger, you can use it directly in your custom workflow to record key information. ```python -class ExampleWorkflow(Workflow): - def run(self) -> List[Experience]: +class ExampleWorkflow(WorkflowWithRecording): + async def run_async(self) -> Metrics: self.logger.info(f"Starting workflow for task: {self.task}") # your workflow logic if some_error_condition: self.logger.error("An error occurred during workflow execution.") self.logger.info(f"Completed workflow for task: {self.task}") - return experiences + return {"example/reward": reward} ``` -Trinity-RFT will automatically create a group of workflow runners to execute the workflows in parallel during training. -Each runner will log its output to a separate log file. The log file naming convention is `explorer_runner_.log`, where `` is the unique identifier of the workflow runner. Such design allows you to trace the execution of each workflow runner independently. And the log files are organized as follows: +Trinity-RFT automatically creates a set of Workflow Runners to execute workflows in parallel. Each runner writes its logs to a separate log file. The file naming convention is `explorer_runner_.log`, where `` is the unique identifier of the workflow runner. With this design, you can independently track each parallel workflow instance. The log directory is organized as follows: ``` ////log/ @@ -562,4 +448,170 @@ Each runner will log its output to a separate log file. The log file naming conv └── ... ``` -Trinity-RFT also provide a convenient command `log` to view these logs in real-time. You can use `trinity log --log-dir /path/to/log/dir -k explorer_runner` command to filter and view the logs of all runners at once or use `trinity log --log-dir /path/to/log/dir -k explorer_runner_0` to view the logs of a specific runner. If you encounter errors or blocking issues during training, you can check the corresponding log files for detailed information to help diagnose and resolve the problems. +Trinity-RFT also provides a convenient `log` command to view these logs in real time. You can use `trinity log --log-dir /path/to/log/dir -k explorer_runner` to filter and view logs of all workflow runners, or `trinity log --log-dir /path/to/log/dir -k explorer_runner_0` to view the log of a specific workflow runner. + +--- + +### Appendix: The legacy Workflow interface (compatible) + +For simple single-turn tasks, Trinity-RFT still keeps the legacy `Workflow` interface. Unlike `WorkflowWithRecording`, the legacy interface requires the workflow to **manually construct and return a list of `Experience` objects**, and the model does not record automatically. All built-in workflows (such as `MathWorkflow`) are still based on this interface. If you do not need a complex agent loop, you can continue to use it. + +The legacy `Workflow` base class interface is as follows: + +```python +class Workflow(WorkflowBase): + + def __init__( + self, + *, + task: Task, + model: ModelWrapper, + auxiliary_models: Optional[List[ModelWrapper]] = None, # mainly for LLM-as-a-judge, can also be used as a distillation teacher + ): + self.task = task + self.model = model + self.auxiliary_model_wrappers = auxiliary_models + self.auxiliary_models = ... # OpenAI clients auto-derived from ModelWrapper + self.logger = get_logger(__name__) # built-in logger for runtime monitoring + + @abstractmethod + def run(self) -> List[Experience]: + """Run the workflow and return a list of Experiences.""" +``` + +##### Initialization and the `run` method + +`Workflow` accepts the same initialization parameters as the new interface (`task`, `model`, `auxiliary_models`), but `model` provides synchronous/asynchronous `generate` and `chat` methods, whose return structure contains `response_text`, `tokens`, `prompt_length`, and `logprobs`. `auxiliary_models` is a list of `openai.OpenAI` / `openai.AsyncOpenAI` clients auto-derived by the framework. + +Here is a simple implementation that manually constructs an `Experience`: + +```python +class ExampleWorkflow(Workflow): + + def __init__(self, task: Task, model: ModelWrapper, auxiliary_models: List): + super().__init__(task=task, model=model, auxiliary_models=auxiliary_models) + self.question = task.raw_task.get("question") + self.answer = task.raw_task.get("answer") + self.rollout_args = task.rollout_args + + def calculate_reward(self, response: str, truth: str) -> float: + if response == truth: + return 1.0 + else: + return 0.0 + + def run(self) -> List[Experience]: + responses = self.model.chat( + [ + { + "role": "user", + "content": f"Question:\n{self.question}", + } + ], + temperature=self.rollout_args.temperature, + ) + response = responses[0] + reward: float = self.calculate_reward(response.response_text, self.answer) + return [ + Experience( + tokens=response.tokens, + prompt_length=response.prompt_length, + reward=reward, + logprobs=response.logprobs, + ) + ] +``` + +##### Batch repeat runs + +The legacy `Workflow` supports `can_repeat` and `set_repeat_times`, which can obtain multiple responses to the same question in one `run` via batch inference (suitable for algorithms such as GRPO). `set_repeat_times` takes `repeat_times` (the number of executions) and `run_id_base` (the first run ID, used in multi-turn interaction scenarios): + +```python +class ExampleWorkflow(Workflow): + can_repeat: bool = True + + def set_repeat_times(self, repeat_times, run_id_base): + self.repeat_times = repeat_times + self.run_id_base = run_id_base + + def run(self) -> List[Experience]: + responses = self.model.chat( + [ + { + "role": "user", + "content": f"Question:\n{self.question}", + } + ], + n=self.repeat_times, + temperature=self.rollout_args.temperature, + ) + experiences = [] + for response in responses: + reward: float = self.calculate_reward(response.response_text, self.answer) + experiences.append( + Experience( + tokens=response.tokens, + prompt_length=response.prompt_length, + reward=reward, + logprobs=response.logprobs, + ) + ) + return experiences +``` + +##### Using the OpenAI API and `extract_experience_from_history` + +In the legacy interface, to call the model via the OpenAI API style, you can get a client via `self.model.get_openai_client()` (or `get_openai_async_client()`). Recording and the OpenAI API server are enabled automatically by the framework (no need to manually configure `enable_history` / `enable_openai_api`); the framework records the trainable data automatically, and you can extract it into a list of `Experience` via `extract_experience_from_history`: + +```python +class ExampleWorkflow(Workflow): + + def __init__(self, task: Task, model: ModelWrapper, auxiliary_models: List): + super().__init__(task=task, model=model, auxiliary_models=auxiliary_models) + self.client: openai.OpenAI = self.model.get_openai_client() + self.agent = MyAgent(openai_client=self.client) + + def calculate_reward(self, response: str) -> float: + # your reward calculation logic + + def run(self) -> List[Experience]: + response = self.agent.run() + reward = self.calculate_reward(response) + experiences = self.model.extract_experience_from_history() + for exp in experiences: + exp.reward = reward + return experiences +``` + +```{tip} +1. The legacy OpenAI API only automatically records calls to `openai.OpenAI.chat.completions.create` and `openai.AsyncOpenAI.chat.completions.create`, and does not support streaming output. +2. When calling `chat.completions.create`, the `model` field can be obtained via `openai_client.models.list().data[0].id` or `openai_client.model_path`. +3. For a more complex workflow example using the OpenAI API, see [ReAct Agent training](./example_react.md). +``` + +For LLM-as-a-judge under the legacy interface, `auxiliary_models` is a list of OpenAI clients auto-derived by the framework and can be called directly: + +```python +class MyWorkflow(Workflow): + def __init__(self, *, task, model, auxiliary_models=None): + super().__init__(task=task, model=model, auxiliary_models=auxiliary_models) + self.judge_model = self.auxiliary_models[0] # auto-derived OpenAI client + + def run(self) -> List[Experience]: + response = self.do_something() + reward_response = self.judge_model.chat.completions.create( + model=self.judge_model.model_path, + messages=[...], + temperature=0.0, + max_tokens=10, + ) + reward = float(reward_response.choices[0].message.content.strip()) + return [ + Experience( + tokens=response.tokens, + prompt_length=response.prompt_length, + reward=reward, + logprobs=response.logprobs, + ) + ] +``` diff --git a/docs/sphinx_doc/source/tutorial/trinity_configs.md b/docs/sphinx_doc/source/tutorial/trinity_configs.md index e9793b0933a..54be1063c5c 100644 --- a/docs/sphinx_doc/source/tutorial/trinity_configs.md +++ b/docs/sphinx_doc/source/tutorial/trinity_configs.md @@ -421,8 +421,6 @@ explorer: engine_type: vllm engine_num: 1 tensor_parallel_size: 1 - enable_history: false - enable_openai_api: false nnodes: 1 auxiliary_models: - model_path: Qwen/Qwen2.5-7B-Instruct @@ -460,8 +458,6 @@ explorer: - `external`: Use external API-based model engine. - `rollout_model.engine_num`: Number of inference engines. - `rollout_model.tensor_parallel_size`: Degree of tensor parallelism. -- `rollout_model.enable_history`: Whether to enable model call history recording. If set to `true`, the model wrapper automatically records the return experiences of model calls. Please periodically extract the history via `extract_experience_from_history` to avoid out-of-memory issues. Default is `false`. -- `rollout_model.enable_openai_api`: Whether to enable the openai API provided by Explorer. Default is `false`. - `rollout_model.nnodes`: Number of nodes for each engine. Default is `1`. Only take effect when `rollout_model.engine_type` is `vllm` or `sglang`. When `nnodes` is greater than `1`, each engine instance will exclusively occupy the GPU resources of the full `nnodes` nodes (`nnodes * cluster.gpu_per_node`); sharing nodes with other instances is not supported. - `auxiliary_models`: Additional models used for custom workflows, which has the same configuration options as `rollout_model`. - `eval_interval`: Interval (in steps) for evaluating the model. diff --git a/docs/sphinx_doc/source_zh/tutorial/develop_workflow.md b/docs/sphinx_doc/source_zh/tutorial/develop_workflow.md index 8418a31a539..fb7c3fd470c 100644 --- a/docs/sphinx_doc/source_zh/tutorial/develop_workflow.md +++ b/docs/sphinx_doc/source_zh/tutorial/develop_workflow.md @@ -12,22 +12,21 @@ ```{mermaid} flowchart LR - A([Task]) & B([Model]) --> C[Workflow] - C --> D([Experience]) + A([Task]) --> C[Workflow] + C -- "调用 OpenAI API" --> B([Rollout Model]) + B -- "自动 recording" --> D([Experience]) + C -- "update_reward" --> D ``` - **任务（Task）** ({class}`trinity.common.workflows.Task`)：结构化的数据实例，包含了工作流一次运行所需的各种信息。一般情况下由训练数据集提供，数据集中的每个样本都会被转化为一个 `Task` 实例。`Task` 的内容根据任务类型而异： - **数学问题**：包含问题和答案。 - **编程场景**：包含题目的描述、测试用例、运行环境等复杂信息。 -- **模型（Model）** ({class}`trinity.common.models.model.ModelWrapper`)：被训练的模型，工作流内需要使用该模型来执行推理。该实例由 Trinity-RFT 自动提供，支持同步以及异步的 `generate` 以及 `chat` 等方法，同时也提供了 OpenAI API 接口，能够兼容大部分 Agent 框架。 +- **模型（Rollout Model）** ({class}`trinity.common.models.model.ModelWrapper`)：被训练的模型。工作流通过模型暴露的 `base_url` 和 `api_key` 自行创建 OpenAI 客户端来调用模型推理接口；模型在响应的同时会**自动记录**生成过程并转化为可用于训练的 `Experience`，工作流无需手动构造。 -- **工作流（Workflow）** ({class}`trinity.common.workflows.Workflow`)：定义了 Agent 与 Environment 的交互流程。`Workflow` 通过 `Task` 中提供的信息初始化自身，并借助 `Model` 来执行其中定义好的交互流程。与常规 Agent 应用不同的是，工作流内部还需要计算奖励信号（reward）以指导训练过程。Trinity-RFT 包含多个内置工作流： - - `MathWorkflow` ({class}`trinity.common.workflows.MathWorkflow`)：用于数学场景，将问题提交给 LLM，解析 LLM 响应，并计算分数（奖励）。 - - `WebShopWorkflow` ({class}`trinity.common.workflows.WebShopWorkflow`)：用于 webshop 场景，包含与环境的多轮交互。 - - `AgentScopeReActWorkflow` ({class}`trinity.common.workflows.AgentScopeReActWorkflow`)：直接使用现有的 ReActAgent（基于 AgentScope）来解决问题。 +- **工作流（Workflow）** ({class}`trinity.common.workflows.WorkflowBase`)：定义了 Agent 与 Environment 的交互流程。`Workflow` 通过 `Task` 中提供的信息初始化自身，并借助 Rollout Model 执行其中定义好的交互流程。与常规 Agent 应用不同的是，工作流内部还需要计算奖励信号（reward）以指导训练过程，并通过 `update_reward` 方法将奖励回填到模型自动记录的 `Experience` 上。 -- **经验（Experience）** ({class}`trinity.common.experience.Experience`)：`Workflow` 的运行产出。产出的数量以及内部数据格式取决于所使用的训练算法。例如，对于常见的 PPO/GRPO 算法，`Experience` 包含 token ID 列表、动作掩码（标识哪些 token 是由 LLM 生成的）、每个 token 的对数概率（logprobs）、奖励信号（reward）等。 +- **经验（Experience）** ({class}`trinity.common.experience.Experience`)：训练所需的数据单元。`Experience` 会由 Rollout Model 在推理过程中自动记录产生，其数量与内部数据格式取决于所使用的训练算法。例如，对于常见的 PPO/GRPO 算法，`Experience` 包含 token ID 列表、动作掩码（标识哪些 token 是由 LLM 生成的）、每个 token 的对数概率（logprobs）、奖励信号（reward）等。工作流不需要、也不应该手动构造 `Experience` 对象。 --- @@ -37,8 +36,10 @@ flowchart LR 为处理 `Task` 内容的差异，Trinity-RFT 提供了一个统一的 `Task` 接口，包含以下字段： - **`workflow`** (`str`)：你的工作流类的注册名称。你可以在 YAML 配置文件的 `buffer.explorer_input.taskset.default_workflow_type` 中指定。 +- **`raw_task`** (`Dict`)：原始数据的记录，以 `Dict` 格式存储。对于高度定制化的工作流，你可以直接使用 `raw_task` 初始化 `Workflow` 实例，而不依赖后续的字段。 + +下面的字段都是可选字段，一般情况下无需设置： - **`reward_fn`** (`Optional[str]`)：你的奖励函数的注册名称。你可以在 `buffer.explorer_input.taskset.default_reward_fn_type` 中指定。注意某些工作流已内置奖励计算；此时可省略该字段。 -- **`raw_task`** (`Dict`)：原始数据的记录，以 `Dict` 格式存储。对于高度定制化的工作流，你可以直接使用 `raw_task` 初始化 `Workflow` 实例，而不依赖以下字段。 - **`format_args`** ({class}`trinity.common.config.FormatConfig`)：便于构造 `Workflow` 实例的参数。例如，`prompt_key` 和 `response_key` 可用于从 `raw_task` 中提取 prompt 和 response。这些设置来自 YAML 配置文件，可在 `buffer.explorer_input.task_set.format` 中设置。 - **`rollout_args`** ({class}`trinity.common.config.GenerationConfig`)：控制 rollout 过程的参数，如 `temperature`。该字段也来自 YAML 配置文件，可在 `buffer.explorer_input.task_set.rollout_args` 中设置。 - **`workflow_args`** (`Dict`)：用于构造 `Workflow` 实例的参数字典。相比 `format_args` 和 `rollout_args` 更灵活。该字段也来自 YAML 配置文件，可在 `buffer.explorer_input.task_set.workflow_args` 中设置。通常无需设置此字段。 @@ -66,7 +67,7 @@ flowchart LR buffer: explorer_input: taskset: - default_workflow: "math_workflow" + default_workflow_type: "math_workflow" path: ${oc.env:TRINITY_TASKSET_PATH} format: prompt_key: "question" @@ -82,62 +83,80 @@ buffer: ### 步骤 2：实现工作流 -`Workflow` 基类接口如下： +要实现一个新的工作流你需要继承 `WorkflowWithRecording` 基类： ```python -class Workflow(ABC): +class WorkflowWithRecording(WorkflowBase): def __init__( self, *, task: Task, model: ModelWrapper, - auxiliary_models: Optional[List[ModelWrapper]] = None, # 主要用于 LLM-as-a-judge 场景, 也可以用作distillation的techer + auxiliary_models: Optional[List[ModelWrapper]] = None, ): - self.task = task - self.model = model - self.auxiliary_model_wrappers = auxiliary_models - self.auxiliary_models = ... # 从 ModelWrapper 自动派生的 OpenAI client - self.logger = get_logger(__name__) # 用于运行时监控的内置 logger + """初始化工作流""" + + async def run_async(self) -> Metrics: + """运行工作流并返回一个 Metric 字典。""" + # 你需要实现该方法 + + @property + def base_url(self) -> str: + """返回 rollout 模型的 base_url。""" + + @property + def api_key(self) -> str: + """返回 rollout 模型的 api_key。""" + + @property + def model_name(self) -> str: + """返回 rollout 模型的 model_name。""" + + async def update_reward( + self, + reward: float, + info: Optional[Dict] = None, + ): + """将 reward 回填到模型自动记录的 Experience 上，同时可选附带额外信息 info。""" - @abstractmethod - def run(self) -> List[Experience]: - """Run the workflow and return a list of Experiences.""" ``` #### 初始化你的工作流 -`Workflow` 接受以下初始化参数： +`WorkflowWithRecording` 接受以下初始化参数： - `task`({class}`trinity.common.workflows.Task`)：数据集中的单个任务。 -- `model`({class}`trinity.common.models.model.ModelWrapper`)：正在训练的模型，提供类似于 OpenAI 的接口，能够接收对话消息列表并返回 LLM 生成的内容（包括回复文本 `response_text`、完整序列 token id `tokens`、prompt 部分 token 长度 `prompt_length`，以及输出 token 对数概率列表 `logprobs`）。 -- `auxiliary_models`(`List[ModelWrapper]`)：辅助模型的 ModelWrapper 列表。可通过 `self.auxiliary_models` 访问 OpenAI client（根据 workflow 的 `is_async` 自动派生）。 +- `model`({class}`trinity.common.models.model.ModelWrapper`)：正在训练的 rollout 模型，你可以直接通过 `WorkflowWithRecording` 的 `base_url`，`api_key` 以及 `model_name` 属性来创建 OpenAI 客户端从而调用模型推理接口。 +- `auxiliary_models`(`List[ModelWrapper]`)：辅助模型的 `ModelWrapper` 列表。每个元素同样暴露 `base_url`、`api_key`、`model_name`，可直接用于创建 OpenAI 客户端（详见 [LLM-as-a-judge 支持](#llm-as-a-judge-支持)）。 -以下是一个仅使用 `raw_task` 和 `rollout_args` 初始化简单工作流的示例。在更复杂的情况下，你可以使用 `format_args` 进行进一步自定义。 +以下是一个简单工作流的初始化示例。我们在 `__init__` 中使用 `base_url` 和 `api_key` 创建异步 OpenAI 客户端，并取出模型名称： ```python -class ExampleWorkflow(Workflow): +import openai +from trinity.common.workflows import WorkflowWithRecording - def __init__(self, task: Task, model: ModelWrapper, auxiliary_models: List): +class ExampleWorkflow(WorkflowWithRecording): + + def __init__(self, *, task: Task, model: ModelWrapper, auxiliary_models: List = None): super().__init__(task=task, model=model, auxiliary_models=auxiliary_models) self.question = task.raw_task.get("question") self.answer = task.raw_task.get("answer") self.rollout_args = task.rollout_args - # Optional: If you want to use OpenAI API in your workflow - # self.openai_client = self.model.get_openai_client() + # 通过 base_url 和 api_key 创建 OpenAI 客户端 + self.client = openai.AsyncOpenAI(base_url=self.base_url, api_key=self.api_key) ``` -#### 实现 `run` 方法 +#### 实现 `run_async` 方法 + +`run_async` 是工作流的核心方法。它没有输入参数，返回一个 `Metrics` 字典。 -`run` 方法是工作流的核心方法。该方法没有输入参数，返回一个 `Experience` 列表。 -以下是一个数学工作流的简单实现。 +工作流的职责是：调用模型完成 agent 任务、计算 reward、通过 `update_reward` 将 reward 回填到模型自动记录的 `Experience` 上，最后返回用于监控的 metric。 -我们首先调用模型，使用给定的问题和 rollout 参数生成答案。 -然后使用 `calculate_reward` 函数计算答案的奖励。 -最后，我们将生成的答案和奖励封装为`Experience` 实例并返回。 +以下是一个数学工作流的简单实现。我们先用 OpenAI 客户端生成答案，再计算奖励并回填： ```python -class ExampleWorkflow(Workflow): +class ExampleWorkflow(WorkflowWithRecording): # the __init__ function @@ -147,10 +166,11 @@ class ExampleWorkflow(Workflow): else: return 0.0 - def run(self) -> List[Experience]: - # call the model to generate multiple responses - responses = self.model.chat( - [ + async def run_async(self) -> Metrics: + # 调用模型生成回复 + resp = await self.client.chat.completions.create( + model=self.model_name, + messages=[ { "role": "user", "content": f"Question:\n{self.question}", @@ -158,115 +178,75 @@ class ExampleWorkflow(Workflow): ], temperature=self.rollout_args.temperature, ) - response = responses[0] # there is only one response - reward: float = self.calculate_reward(response.response_text, self.answer) - return [ - Experience( - tokens=response.tokens, - prompt_length=response.prompt_length, - reward=reward, - logprobs=response.logprobs, - ) - ] + response_text = resp.choices[0].message.content + # 计算 reward 并回填到模型自动记录的 Experience 上 + reward: float = self.calculate_reward(response_text, self.answer) + await self.update_reward(reward) + # 返回需要监控的 metric + return {"example/reward": reward} +``` + +```{note} +1. rollout 模型会自动记录每次 `chat.completions.create` 调用产生的训练数据并转化为 `Experience`。`update_reward` 会将 reward 精确回填到本次运行产生的 `Experience` 上。 +2. 对于包含多轮交互的工作流，`update_reward` 会将 reward 回填到本次运行产生的所有 `Experience` 上。 +3. `run_async` 返回的 `Metrics` 字典仅用于运行时监控与日志展示。 ``` #### 注册你的工作流 -为了让 Trinity-RFT 能够通过配置文件中的名称自动找到你的工作流，你需要在 `trinity/common/workflows/__init__.py` 中的 `default_mapping` 中注册。 +为了让 Trinity-RFT 能够通过配置文件中的名称自动找到你的工作流，你需要将其注册到 `WORKFLOWS` 注册表中。推荐使用装饰器方式注册： ```python -WORKFLOWS = Registry( - "workflows", - default_mapping={ - "example_workflow": "trinity.common.workflows.workflow.ExampleWorkflow", - }, -) +from trinity.common.workflows import WORKFLOWS, WorkflowWithRecording + +@WORKFLOWS.register_module(name="example_workflow") +class ExampleWorkflow(WorkflowWithRecording): + ... ``` -#### 性能调优 +也可以直接注册，或在 `trinity/common/workflows/__init__.py` 的 `default_mapping` 中添加一条 `"example_workflow": "path.to.module.ExampleWorkflow"` 映射。 -以下是一些可选的性能调优方法，能够提升工作流的运行效率。当然，这些方法并非所有工作流都需要实现，具体取决于你的工作流设计。 +#### 性能调优 -##### 避免重复初始化 +对于较为复杂的工作流，每次重新初始化会带来额外计算开销。此时，你可以设置 `can_reset` 类属性并实现 `reset` 方法以避免重复初始化。 -对于较为复杂的工作流，每次重新初始化会带来额外计算开销。 -此时，你可以设置 `can_reset` 属性并实现 `reset` 方法以避免重复初始化。 +注意在 `reset` 方法中必须使用输入的 `task` 覆盖工作流的 `task` 属性，并使用 `task.api_key` 更新模型和客户端的 API Key。 -`can_reset` 是一个类属性，表示工作流是否支持轻量化重置。 +> Trinity-RFT 内部借助 `api_key` 来区分不同任务产生的 Experience，如果不更新 API Key，可能会导致不同任务的 Experience 被错误地归类，导致 reward 回填错误。 -`reset` 方法接受一个新的 `Task` 实例，并使用该实例更新工作流的状态。 +以下是一个简单示例： ```python -class ExampleWorkflow(Workflow): +class ExampleWorkflow(WorkflowWithRecording): can_reset: bool = True # some code # ... def reset(self, task: Task): + self.task = task + self.model.set_api_key(task.api_key) + self.client.api_key = task.api_key self.question = task.raw_task.get("question") self.answer = task.raw_task.get("answer") ``` -##### 批量运行推理任务 - -当前流行的很多 RL 算法需要多次运行同一个任务(例如 GRPO)。该场景下一些简单任务可以直接通过模型批量推理来获得一个问题的多个回复以提升效率。 -针对该情况，你可以设置 `can_repeat` 属性并实现 `set_repeat_times` 方法。 - -`can_repeat` 是一个类属性，指示工作流是否支持在 `run` 方法内多次执行。 - -`set_repeat_times` 方法接受两个参数：`repeat_times` 指定了在 `run` 方法内需要执行的次数，`run_id_base` 是一个整数，用于标识多次运行中第一次的运行 ID，之后各次的 ID 基于此递增（该参数用于多轮交互场景，单次模型调用即可完成的任务可以忽略该项）。 - -```python -class ExampleWorkflow(Workflow): - can_repeat: bool = True - # some code - - def set_repeat_times(self, repeat_times, run_id_base): - self.repeat_times = repeat_times - self.run_id_base = run_id_base - - def run(self) -> List[Experience]: - # call the model to generate multiple responses - responses = self.model.chat( - [ - { - "role": "user", - "content": f"Question:\n{self.question}", - } - ], - n=self.repeat_times, # run multiple times in one call - temperature=self.rollout_args.temperature, - ) - experiences = [] - for response in responses: - # calculate reward - reward: float = self.calculate_reward(response.response_text, self.answer) - # construct Experience - experiences.append( - Experience( - tokens=response.tokens, - prompt_length=response.prompt_length, - reward=reward, - logprobs=response.logprobs, - ) - ) - return experiences -``` - - #### 完整代码示例 ```python -class ExampleWorkflow(Workflow): +import openai +from trinity.common.workflows import WORKFLOWS, WorkflowWithRecording + +@WORKFLOWS.register_module(name="example_workflow") +class ExampleWorkflow(WorkflowWithRecording): can_reset: bool = True - can_repeat: bool = True - def __init__(self, task: Task, model: ModelWrapper, auxiliary_models: List): + def __init__(self, *, task: Task, model: ModelWrapper, auxiliary_models: List = None): super().__init__(task=task, model=model, auxiliary_models=auxiliary_models) self.question = task.raw_task.get("question") self.answer = task.raw_task.get("answer") self.rollout_args = task.rollout_args + self.client = openai.AsyncOpenAI(base_url=self.base_url, api_key=self.api_key) def calculate_reward(self, response: str, truth: str) -> float: if response == truth: @@ -274,47 +254,35 @@ class ExampleWorkflow(Workflow): else: return 0.0 - def run(self) -> List[Experience]: - # call the model to generate multiple responses - responses = self.model.chat( - [ + async def run_async(self) -> Metrics: + resp = await self.client.chat.completions.create( + model=self.model_name, + messages=[ { "role": "user", "content": f"Question:\n{self.question}", } ], - n=self.rollout_args.n, temperature=self.rollout_args.temperature, ) - experiences = [] - for response in responses: - # calulcate reward - reward: float = self.calculate_reward(response.response_text, self.answer) - # construct Experience - experiences.append( - Experience( - tokens=response.tokens, - prompt_length=response.prompt_length, - reward=reward, - logprobs=response.logprobs, - ) - ) - return experiences + response_text = resp.choices[0].message.content + reward: float = self.calculate_reward(response_text, self.answer) + await self.update_reward(reward) + return {"example/reward": reward} def reset(self, task: Task): + self.task = task + self.model.set_api_key(task.api_key) + self.client.api_key = task.api_key self.question = task.raw_task.get("question") self.answer = task.raw_task.get("answer") - - def set_repeat_times(self, repeat_times, run_id_base): - self.repeat_times = repeat_times - self.run_id_base = run_id_base ``` --- ### 步骤 3：使用你的工作流 -实现并注册工作流后，就可以通过将配置文件中 `buffer.explorer_input.taskset` 的 `default_workflow_type` 域设置为你的工作流名称来使用它。例如： +实现并注册工作流后，就可以通过将配置文件中 `buffer.explorer_input.taskset` 的 `default_workflow_type` 设置为你的工作流名称来使用它。例如： ```yaml buffer: @@ -334,74 +302,7 @@ trinity run --config --- -### 其他进阶特性 - -#### async 支持 - -本节样例主要针对同步模式，如果你的工作流需要使用异步方法（例如异步 API）,你可以将 `is_async` 属性设置为 `True`，然后实现 `run_async` 方法，在这种情况下不再需要实现 `run` 方法，并且初始化参数 `auxiliary_models` 也会自动变为 `List[openai.AsyncOpenAI]` 类型，其余方法和属性保持不变。 - -```python -class ExampleWorkflowAsync(Workflow): - - is_async: bool = True - - async def run_async(self) -> List[Experience]: - # your async code here - - # no need to implement run() method -``` - -#### 使用 OpenAI API - -Trinity-RFT 的 Model 提供了 OpenAI API 接口，能够降低模型推理部分的学习成本并简化工作流的实现。 - -为了激活 OpenAI API 服务，你需要将配置文件中 `explorer.rollout_model.enable_openai_api` 设置为 `true` 。这样就可以通过 `Model` 实例的 `get_openai_client` 方法获取 `openai.OpenAI` 实例。 - -另外，由于 OpenAI API 无法提供训练所需的各项数据，你还需要将 `explorer.rollout_model.enable_history` 设置为 `true`，让框架自动记录可用于训练的数据并转化为 `Experience` 列表。你可以通过 `extract_experience_from_history` 方法来提取这些可用于训练的数据。 - - -```yaml -# example config snippet -explorer: - rollout_model: - enable_openai_api: true - enable_history: true - # Other fields -``` - -```python -class ExampleWorkflow(Workflow): - - def __init__(self, task: Task, model: ModelWrapper, auxiliary_models: List): - super().__init__(task=task, model=model, auxiliary_models=auxiliary_models) - self.model = model - self.client: openai.OpenAI = self.model.get_openai_client() - # or async client - # self.client: openai.AsyncOpenAI = self.model.get_openai_async_client() - self.agent = MyAgent(openai_client=self.client) - - def calculate_reward(self, response: str) -> float: - # your reward calculation logic - - def run(self) -> List[Experience]: - # run your agent - response = self.agent.run() - # calculate reward - reward = self.calculate_reward(response) - # extract experiences from history recorded in self.model - experiences = self.model.extract_experience_from_history() - for exp in experiences: - exp.reward = reward - return experiences -``` - -```{tip} -1. 当前的 OpenAI API 仅会自动记录 `openai.OpenAI.chat.completions.create` 以及 `openai.AsyncOpenAI.chat.completions.create` 方法的调用历史并转化为 `Experience` 结构，且不支持流式输出。 -2. 调用 `chat.completions.create` 时，其中的 `model` 字段可通过 `openai_client.models.list().data[0].id` 或 `openai_client.model_path` 获取。 -3. 更复杂的使用 OpenAI API 的工作流实例可参考 [ReAct Agent 训练](./example_react.md)。 -``` - -#### LLM-as-a-judge 支持 +### LLM-as-a-judge 支持 LLM-as-a-judge 是一种常见的奖励计算方法，尤其适用于开放式任务（如编程、写作等）。在这类场景下，Workflow 需要借助额外的 LLM 来评估答案质量并计算奖励信号（reward）。 @@ -430,24 +331,21 @@ explorer: 请注意，每个辅助模型会独立占用 `tensor_parallel_size * engine_num` 个 GPU，请根据硬件资源合理配置。在启用辅助模型后，Trainer 可用的 GPU 数量为总 GPU 数量减去所有辅助模型及被训练的推理模型（`rollout_model`）所占用的 GPU 数量。 -配置文件中指定的辅助模型会自动激活 OpenAI API，并将对应的 `openai.OpenAI` 或 `openai.AsyncOpenAI` 实例 (取决于 `is_async`) 传递给 `Workflow` 初始化方法的 `auxiliary_models` 参数。例如： +配置文件中指定的辅助模型会以 `ModelWrapper` 实例列表的形式传递给 `Workflow` 初始化方法的 `auxiliary_models` 参数。每个 `ModelWrapper` 同样暴露 `base_url`、`api_key`、`model_name`，推荐直接用它们创建 OpenAI 客户端来访问辅助模型： ```python -class MyWorkflow(Workflow): - def __init__( - self, - *, - task: Task, - model: ModelWrapper, - auxiliary_models: Optional[List[ModelWrapper]] = None, - ): +class MyWorkflow(WorkflowWithRecording): + def __init__(self, *, task, model, auxiliary_models=None): super().__init__(task=task, model=model, auxiliary_models=auxiliary_models) - self.judge_model = self.auxiliary_models[0] # 从 ModelWrapper 自动派生的 OpenAI client + self.judge = self.auxiliary_models[0] # ModelWrapper + self.judge_client = openai.AsyncOpenAI( + base_url=self.judge.base_url, api_key=self.judge.api_key + ) - def run(self) -> List[Experience]: - response = self.do_something() - reward_response = self.judge_model.chat.completions.create( - model=self.judge_model.model_path, + async def run_async(self) -> Metrics: + response = await self.do_something() + reward_response = await self.judge_client.chat.completions.create( + model=self.judge.model_name, messages=[ { "role": "system", @@ -455,7 +353,7 @@ class MyWorkflow(Workflow): }, { "role": "user", - "content": f"Question:\n{self.task.raw_task['question']}\nAnswer:\n{response.response_text}\nPlease give a score from 0 to 1.", + "content": f"Question:\n{self.task.raw_task['question']}\nAnswer:\n{response}\nPlease give a score from 0 to 1.", }, ], temperature=0.0, @@ -463,14 +361,8 @@ class MyWorkflow(Workflow): ) # 解析奖励分数 reward = float(reward_response.choices[0].message.content.strip()) - return [ - Experience( - tokens=response.tokens, - prompt_length=response.prompt_length, - reward=reward, - logprobs=response.logprobs, - ) - ] + await self.update_reward(reward, info={"source": "llm_as_a_judge"}) + return {"my_workflow/judge_reward": reward} ``` #### 调试模式（Debug Mode） @@ -504,7 +396,7 @@ trinity debug --config --module workflow --output-dir --module viewer --output-dir --port 8502 @@ -517,10 +409,10 @@ trinity debug --config --module viewer --output-dir ////log` 目录下的文件中。这样就可以方便地在训练过程中监控工作流的运行状态。由于所有 Workflow 子类均继承该 logger，因此你可以直接在自定义工作流中使用它来记录关键信息。 ```python -class ExampleWorkflow(Workflow): - def run(self) -> List[Experience]: +class ExampleWorkflow(WorkflowWithRecording): + async def run_async(self) -> Metrics: self.logger.info(f"Starting workflow for task: {self.task}") # your workflow logic if some_error_condition: self.logger.error("An error occurred during workflow execution.") self.logger.info(f"Completed workflow for task: {self.task}") - return experiences + return {"example/reward": reward} ``` 由于 Trinity-RFT 会自动创建一组 Workflow Runners 来并行执行 Workflow。每个运行器会将其日志输出到一个单独的日志文件中。日志文件的命名规则为 `explorer_runner_.log`，其中 `` 是工作流运行器的唯一标识符。通过这种设计，你可以独立地追踪正在并行执行的每个工作流实例的运行情况。日志文件的具体组织结构如下： @@ -557,3 +449,169 @@ class ExampleWorkflow(Workflow): ``` Trinity-RFT 还提供了一个方便的 `log` 命令来实时查看这些日志。你可以使用 `trinity log --log-dir /path/to/log/dir -k explorer_runner` 命令来过滤并查看所有 workflow runner 的日志，或者使用 `trinity log --log-dir /path/to/log/dir -k explorer_runner_0` 来查看特定 workflow runner 的日志。 + +--- + +### 附录：旧版 Workflow 接口（兼容） + +对于简单的单轮任务，Trinity-RFT 仍保留旧版 `Workflow` 接口。与 `WorkflowWithRecording` 不同，旧版接口要求工作流**手动构造并返回 `Experience` 列表**，模型也不会自动 recording。所有内置工作流（`MathWorkflow` 等）目前仍基于此接口。如果你不需要复杂的 agent 循环，可以继续使用它。 + +旧版 `Workflow` 基类接口如下： + +```python +class Workflow(WorkflowBase): + + def __init__( + self, + *, + task: Task, + model: ModelWrapper, + auxiliary_models: Optional[List[ModelWrapper]] = None, # 主要用于 LLM-as-a-judge 场景, 也可以用作distillation的techer + ): + self.task = task + self.model = model + self.auxiliary_model_wrappers = auxiliary_models + self.auxiliary_models = ... # 从 ModelWrapper 自动派生的 OpenAI client + self.logger = get_logger(__name__) # 用于运行时监控的内置 logger + + @abstractmethod + def run(self) -> List[Experience]: + """Run the workflow and return a list of Experiences.""" +``` + +##### 初始化与 `run` 方法 + +`Workflow` 接受与新版相同的初始化参数（`task`、`model`、`auxiliary_models`），但 `model` 提供的是同步/异步的 `generate` 以及 `chat` 方法，返回结构包含 `response_text`、`tokens`、`prompt_length`、`logprobs`。`auxiliary_models` 则是框架自动派生的 `openai.OpenAI` / `openai.AsyncOpenAI` 客户端列表。 + +以下是一个手动构造 `Experience` 的简单实现： + +```python +class ExampleWorkflow(Workflow): + + def __init__(self, task: Task, model: ModelWrapper, auxiliary_models: List): + super().__init__(task=task, model=model, auxiliary_models=auxiliary_models) + self.question = task.raw_task.get("question") + self.answer = task.raw_task.get("answer") + self.rollout_args = task.rollout_args + + def calculate_reward(self, response: str, truth: str) -> float: + if response == truth: + return 1.0 + else: + return 0.0 + + def run(self) -> List[Experience]: + responses = self.model.chat( + [ + { + "role": "user", + "content": f"Question:\n{self.question}", + } + ], + temperature=self.rollout_args.temperature, + ) + response = responses[0] + reward: float = self.calculate_reward(response.response_text, self.answer) + return [ + Experience( + tokens=response.tokens, + prompt_length=response.prompt_length, + reward=reward, + logprobs=response.logprobs, + ) + ] +``` + +##### 批量重复运行 + +旧版 `Workflow` 支持 `can_repeat` 与 `set_repeat_times`，可在一次 `run` 内通过模型批量推理获得同一问题的多个回复（适用于 GRPO 等算法）。`set_repeat_times` 接受 `repeat_times`（执行次数）和 `run_id_base`（首次运行 ID，多轮交互场景使用）： + +```python +class ExampleWorkflow(Workflow): + can_repeat: bool = True + + def set_repeat_times(self, repeat_times, run_id_base): + self.repeat_times = repeat_times + self.run_id_base = run_id_base + + def run(self) -> List[Experience]: + responses = self.model.chat( + [ + { + "role": "user", + "content": f"Question:\n{self.question}", + } + ], + n=self.repeat_times, + temperature=self.rollout_args.temperature, + ) + experiences = [] + for response in responses: + reward: float = self.calculate_reward(response.response_text, self.answer) + experiences.append( + Experience( + tokens=response.tokens, + prompt_length=response.prompt_length, + reward=reward, + logprobs=response.logprobs, + ) + ) + return experiences +``` + +##### 使用 OpenAI API 与 `extract_experience_from_history` + +旧版接口下若要使用 OpenAI API 风格调用模型，可通过 `self.model.get_openai_client()`（或 `get_openai_async_client()`）获取客户端。recording 与 OpenAI API 服务由框架自动开启（无需手动配置 `enable_history` / `enable_openai_api`），框架会自动记录可训练数据，你可通过 `extract_experience_from_history` 将其提取为 `Experience` 列表： + +```python +class ExampleWorkflow(Workflow): + + def __init__(self, task: Task, model: ModelWrapper, auxiliary_models: List): + super().__init__(task=task, model=model, auxiliary_models=auxiliary_models) + self.client: openai.OpenAI = self.model.get_openai_client() + self.agent = MyAgent(openai_client=self.client) + + def calculate_reward(self, response: str) -> float: + # your reward calculation logic + + def run(self) -> List[Experience]: + response = self.agent.run() + reward = self.calculate_reward(response) + experiences = self.model.extract_experience_from_history() + for exp in experiences: + exp.reward = reward + return experiences +``` + +```{tip} +1. 旧版 OpenAI API 仅自动记录 `openai.OpenAI.chat.completions.create` 及 `openai.AsyncOpenAI.chat.completions.create` 的调用历史，且不支持流式输出。 +2. 调用 `chat.completions.create` 时，`model` 字段可通过 `openai_client.models.list().data[0].id` 或 `openai_client.model_path` 获取。 +3. 更复杂的使用 OpenAI API 的工作流实例可参考 [ReAct Agent 训练](./example_react.md)。 +``` + +对于旧版接口下的 LLM-as-a-judge，`auxiliary_models` 是框架自动派生的 OpenAI client 列表，可直接调用： + +```python +class MyWorkflow(Workflow): + def __init__(self, *, task, model, auxiliary_models=None): + super().__init__(task=task, model=model, auxiliary_models=auxiliary_models) + self.judge_model = self.auxiliary_models[0] # 自动派生的 OpenAI client + + def run(self) -> List[Experience]: + response = self.do_something() + reward_response = self.judge_model.chat.completions.create( + model=self.judge_model.model_path, + messages=[...], + temperature=0.0, + max_tokens=10, + ) + reward = float(reward_response.choices[0].message.content.strip()) + return [ + Experience( + tokens=response.tokens, + prompt_length=response.prompt_length, + reward=reward, + logprobs=response.logprobs, + ) + ] +``` diff --git a/docs/sphinx_doc/source_zh/tutorial/trinity_configs.md b/docs/sphinx_doc/source_zh/tutorial/trinity_configs.md index 681f9e4f03a..3a5f335e331 100644 --- a/docs/sphinx_doc/source_zh/tutorial/trinity_configs.md +++ b/docs/sphinx_doc/source_zh/tutorial/trinity_configs.md @@ -418,8 +418,6 @@ explorer: engine_type: vllm engine_num: 1 tensor_parallel_size: 1 - enable_history: false - enable_openai_api: false nnodes: 1 auxiliary_models: - model_path: Qwen/Qwen2.5-7B-Instruct @@ -456,8 +454,6 @@ explorer: - `external`: 使用外部 API 引擎。 - `rollout_model.engine_num`: 推理引擎实例的数量。 - `rollout_model.tensor_parallel_size`: 每个实例的张量并行度。 -- `rollout_model.enable_history`: 是否启用模型调用历史记录功能。若设为 `True`，模型会自动记录调用返回的 experience。请定期通过 `extract_experience_from_history` 提取历史，以避免内存溢出。默认为 `False`。 -- `rollout_model.enable_openai_api`: 是否启用 OpenAI API 推理服务。默认为 `False`。 - `rollout_model.nnodes`: 部署每个推理引擎实例所需的节点数。默认为 `1`。仅在 `rollout_model.engine_type` 为 `vllm` 或 `sglang` 时生效。当 `nnodes` 大于 `1` 时，每个引擎实例将会占用完整的 `nnodes` 个节点的 GPU 资源 (`nnodes * cluster.gpu_per_node`)，不支持与其他实例共享节点。 - `auxiliary_models`: 用于自定义工作流的辅助模型，配置与 `rollout_model` 相同。 - `eval_interval`: 模型评估的间隔（以 step 为单位）。 diff --git a/tests/common/vllm_test.py b/tests/common/vllm_test.py index 4989ad026f2..f674fdddf49 100644 --- a/tests/common/vllm_test.py +++ b/tests/common/vllm_test.py @@ -330,10 +330,10 @@ async def test_generate(self): # noqa: C901 ) self.assertTrue(exp.logprobs.shape[0] == exp.tokens.shape[0] - prompt_length) self.assertTrue(torch.equal(result_dict["input_ids"][0], exp.tokens)) - if self.model_wrapper.config.enable_openai_api: - self.assertIsNotNone(self.model_wrapper.get_openai_client()) - else: - self.assertRaises(ValueError, self.model_wrapper.get_openai_client) + # The OpenAI API server is now always enabled for the rollout model + # (``enable_openai_api`` is a deprecated no-op), so the client is always + # available regardless of the requested value. + self.assertIsNotNone(self.model_wrapper.get_openai_client()) if self.enable_return_routed_experts: openai_messages = [ diff --git a/tests/explorer/workflow_test.py b/tests/explorer/workflow_test.py index 135cd749964..a3508572a15 100644 --- a/tests/explorer/workflow_test.py +++ b/tests/explorer/workflow_test.py @@ -905,7 +905,6 @@ class ConcurrentRecordingWorkflow(WorkflowWithRecording): def __init__(self, model: ModelWrapper, task: Task, auxiliary_models=None): super().__init__(task=task, model=model, auxiliary_models=auxiliary_models) self.client = openai.AsyncOpenAI(base_url=f"{self.base_url}/v1", api_key=self.api_key) - self.model_name = self.model.model_name def reset(self, task: Task): self.task = task diff --git a/trinity/common/config.py b/trinity/common/config.py index 751e9532d89..5e1632d12ce 100644 --- a/trinity/common/config.py +++ b/trinity/common/config.py @@ -1,5 +1,6 @@ # -*- coding: utf-8 -*- """Configs for RFT.""" + from __future__ import annotations import os @@ -567,25 +568,26 @@ class InferenceModelConfig: # For Qwen3 enable_thinking: Optional[bool] = None - # For experience recording. When enabled on vLLM rollout models, the engine - # wraps ``engine_client.generate`` and writes each finished turn as a - # Trinity ``Experience`` to the in-process ``MemoryStore``, keyed by the - # recording identity (``record_key``). This is the single switch for the - # recording flow — when on, the explorer also consumes from the store: the - # Workflow updates rewards by ``record_key`` and the Scheduler drains - # completed task experiences through rollout model actor methods. When off - # (default), recording APIs are disabled. When True, the Allocator forces - # ``enable_openai_api`` for engine-side request capture. - # VLLMModel mirrors the recording config onto the engine instance for the - # recorder to read. The capture width (top-k logprobs) reuses ``logprobs`` - # below (default 1). Requires ``enable_openai_api=True`` (the recording - # runner is the API server). Routed-experts capture is opt-in via - # ``enable_router_replay`` (mirrored to ``enable_return_routed_experts`` in - # ``config_validator``); it is not implied by ``enable_history``, so dense - # models can record history too. + # [Deprecated, not user-settable] Controls engine-side experience recording. + # When enabled, the engine wraps ``engine_client.generate`` / the API server + # middleware and writes each finished turn as a Trinity ``Experience`` to the + # in-process ``MemoryStore``, keyed by the recording identity (``record_key``). + # The ``ConfigValidator`` forces this to ``True`` for the rollout model of + # every engine type (the ``Workflow.execute`` overwrite path and the Scheduler + # drain both rely on experiences being captured) and to ``False`` for + # auxiliary models (which must never record). Any user-supplied value is + # overridden. The capture width (top-k logprobs) reuses ``logprobs`` below + # (default 1). Routed-experts capture is opt-in via ``enable_router_replay`` + # (mirrored to ``enable_return_routed_experts`` in ``config_validator``); it is + # not implied by ``enable_history``, so dense models can record history too. enable_history: bool = False - # For OpenAI API + # [Deprecated, not user-settable] Whether to start the OpenAI API server for + # this model. The API server is now always enabled: it hosts the recording + # runner (vLLM/SGLang) and backs the OpenAI client used by workflows. + # ``ConfigValidator`` forces this to ``True`` for both the rollout model and + # auxiliary models regardless of any user-supplied value. The field is kept + # only for backward compatibility with existing YAML configs. enable_openai_api: bool = False enable_log_requests: bool = False # whether to enable request logging in vLLM API server base_port: Optional[int] = None diff --git a/trinity/common/config_validator.py b/trinity/common/config_validator.py index 9a5558b45bc..d3150c441dd 100644 --- a/trinity/common/config_validator.py +++ b/trinity/common/config_validator.py @@ -705,21 +705,26 @@ def validate(self, config: Config) -> None: # noqa: C901 if config.mode == "serve": # in 'serve' mode, we always enable openai api for rollout model config.explorer.rollout_model.enable_openai_api = True - if config.explorer.rollout_model.engine_type.startswith("vllm") or ( - config.explorer.rollout_model.engine_type == "sglang" - ): - # enable_history is the single switch for engine-side recording and is - # mandatory for the rollout model: the Workflow.execute overwrite path - # and the Scheduler drain both rely on experiences being captured into - # the in-process store keyed by the recording identity. - if not config.explorer.rollout_model.enable_history: - config.explorer.rollout_model.enable_history = True - self.logger.warning( - "`explorer.rollout_model.enable_history` is required for the rollout " - "model's recording flow; force-set to True." - ) - # enable_history implies the OpenAI API server (the recording runner). - config.explorer.rollout_model.enable_openai_api = True + # ``enable_history`` is the single switch for engine-side recording and is + # mandatory for the rollout model of every engine type: the + # ``Workflow.execute`` overwrite path and the Scheduler drain both rely on + # experiences being captured into the in-process store keyed by the + # recording identity. vLLM/SGLang host the recorder in the API server; + # Tinker builds its own in-process recorder; external models are + # bench-only and never run the recording path, but the flag is set for + # consistency. This field is not user-settable; any user value is + # overridden here. Auxiliary models are forced to ``False`` below. + if not config.explorer.rollout_model.enable_history: + config.explorer.rollout_model.enable_history = True + self.logger.warning( + "`explorer.rollout_model.enable_history` is required for the rollout " + "model's recording flow; force-set to True." + ) + # The OpenAI API server is always enabled for the rollout model: it hosts + # the recording runner (vLLM/SGLang) and backs the OpenAI client used by + # workflows. ``enable_openai_api`` is a deprecated no-op kept only for + # backward compatibility; it is forced on here regardless of user setting. + config.explorer.rollout_model.enable_openai_api = True self._validate_inference_parallel_config(config.explorer.rollout_model, "rollout_model") # auxiliary models for aux_model in config.explorer.auxiliary_models: diff --git a/trinity/common/models/allocator.py b/trinity/common/models/allocator.py index 3d53380698a..991eee22cbb 100644 --- a/trinity/common/models/allocator.py +++ b/trinity/common/models/allocator.py @@ -83,16 +83,16 @@ async def create_engine( config.engine_id = engine_id if config.engine_type.startswith("vllm") or config.engine_type == "sglang": - # enable_history is the single switch for engine-side recording. The - # recorder runs inside the OpenAI API server, so it must be on. - # ``enable_return_routed_experts`` is NOT forced here: it is driven - # by the user's ``enable_router_replay`` (see ``config_validator``), - # so dense models can record history without vLLM's routed-experts - # capturer (which raises on configs lacking ``num_experts_per_tok``). - # The recorder simply leaves ``Experience.routed_experts`` as None - # when the engine did not capture any. - if config.enable_history: - config.enable_openai_api = True + # ``enable_history`` and ``enable_openai_api`` are both forced on for + # the rollout model by ``ConfigValidator`` (the recorder runs inside + # the OpenAI API server). Nothing to do here. Note: + # ``enable_return_routed_experts`` is NOT forced — it is driven by the + # user's ``enable_router_replay`` (see ``config_validator``), so dense + # models can record history without vLLM's routed-experts capturer + # (which raises on configs lacking ``num_experts_per_tok``). The + # recorder simply leaves ``Experience.routed_experts`` as None when + # the engine did not capture any. + pass actor_bundle_lists = [] for node_id in range(config.nnodes): diff --git a/trinity/common/models/model.py b/trinity/common/models/model.py index 4a1cc304184..767a90e530d 100644 --- a/trinity/common/models/model.py +++ b/trinity/common/models/model.py @@ -522,7 +522,12 @@ def __init__( async def prepare(self) -> None: """Prepare some necessary information for the model before inference.""" - if not self.config.enable_openai_api: + # The OpenAI API server is always enabled for vLLM/SGLang models; only the + # Tinker and external backends skip the HTTP probe — Tinker has no real + # API server (its OpenAI client is a Ray-remote shim), and external's + # address comes from the environment. This short-circuit is intentionally + # based on engine type, not on the deprecated ``enable_openai_api`` flag. + if self.config.engine_type in {"tinker", "external"}: return if self.api_address is None: if self.model is None: @@ -636,7 +641,7 @@ def base_url(self) -> str: """Get the base URL of the API server.""" if not self.api_address: raise ValueError("API address is not set. Cannot get base URL.") - return self.api_address + return f"{self.api_address}/v1" @property def api_key(self) -> str: @@ -706,11 +711,6 @@ def get_openai_client(self) -> "openai.OpenAI": """ import openai - if not self.config.enable_openai_api: - raise ValueError( - "OpenAI API is not enabled for this model. OpenAI client is unavailable." - ) - if self.openai_client is not None: setattr(self.openai_client, "model_path", self.config.model_path) return self.openai_client diff --git a/trinity/common/models/sglang_model.py b/trinity/common/models/sglang_model.py index 97da3c5d5be..b894734929c 100644 --- a/trinity/common/models/sglang_model.py +++ b/trinity/common/models/sglang_model.py @@ -277,9 +277,9 @@ def __init__( super().__init__(config) if config.cuda_visible_devices: os.environ["CUDA_VISIBLE_DEVICES"] = config.cuda_visible_devices - if not self.config.enable_openai_api: - self.logger.warning("SGLangRolloutModel requires OpenAI API to be enabled.") - self.config.enable_openai_api = True + # The OpenAI API server is always enabled (forced by ``ConfigValidator``); + # ``enable_openai_api`` is a deprecated no-op kept only for backward + # compatibility. os.environ["SGLANG_GRPC_PORT"] = "12345" # a dummy port not actually used os.environ["SGLANG_ENABLE_GRPC"] = "0" self.api_server_host: Optional[str] = None diff --git a/trinity/common/models/vllm_model.py b/trinity/common/models/vllm_model.py index b3935344aa7..cc4b7ba18b7 100644 --- a/trinity/common/models/vllm_model.py +++ b/trinity/common/models/vllm_model.py @@ -695,10 +695,6 @@ async def run_api_server(self) -> bool: Returns: success (bool): Whether the API server is started successfully. """ - if not self.config.enable_openai_api: - self.logger.info("OpenAI API server is not enabled. Skipping...") - return False # Not enabled - if self.api_server_host is not None and self.api_server_port is not None: self.logger.info("OpenAI API server is already running. Skipping...") return True # already running diff --git a/trinity/common/workflows/workflow.py b/trinity/common/workflows/workflow.py index 26ca8ba8470..6032007683c 100644 --- a/trinity/common/workflows/workflow.py +++ b/trinity/common/workflows/workflow.py @@ -442,7 +442,7 @@ class WorkflowWithRecording(WorkflowBase): the recorded experience data with the reward and optional info. """ - can_reset: bool = True + can_reset: bool = False is_async: bool = True def __init__( @@ -454,9 +454,7 @@ def __init__( ): super().__init__(task=task, model=model) # Store ModelWrapper instances - self.auxiliary_model_wrappers = auxiliary_models - # Get OpenAI clients from ModelWrapper - self.auxiliary_models = [m.get_openai_async_client() for m in (auxiliary_models or [])] + self.auxiliary_models = auxiliary_models @property def base_url(self) -> str: @@ -468,9 +466,10 @@ def api_key(self) -> str: """API_KEY of the OpenAI API of the rollout model.""" return self.task.api_key - def reset(self, task: Task): - """Reset the workflow.""" - raise NotImplementedError + @property + def model_name(self) -> str: + """Model name of the rollout model.""" + return self.model.model_name async def run_async(self) -> Metrics: """Run workflow asynchronously and return metrics for the completed run.""" diff --git a/trinity/explorer/proxy/service.py b/trinity/explorer/proxy/service.py index 4385a7adfb1..8e96c867365 100644 --- a/trinity/explorer/proxy/service.py +++ b/trinity/explorer/proxy/service.py @@ -117,7 +117,8 @@ async def allocate_model(self, increase_count: bool = True) -> Tuple[str, int]: self.running_model_ids.rotate(-1) if model.api_address is None: raise ValueError( - "Model does not have a valid API address, please set `enable_openai_api` to `True`." + "Model does not have a valid API address; the OpenAI API server " + "should have been started automatically during model preparation." ) return model.api_address, self.model_version_map[model_id] diff --git a/trinity/explorer/workflow_runner.py b/trinity/explorer/workflow_runner.py index a013c8ffef9..f1abfca9064 100644 --- a/trinity/explorer/workflow_runner.py +++ b/trinity/explorer/workflow_runner.py @@ -90,7 +90,7 @@ def _create_workflow_instance(self, task: Task) -> Workflow: if ( self.workflow_instance is None or not self.workflow_instance.__class__ == task.workflow - or not getattr(self.workflow_instance.__class__, "can_reset", True) + or not getattr(self.workflow_instance.__class__, "can_reset", False) ): # Pass ModelWrapper directly; Workflow.__init__ will get OpenAI clients automatically self.workflow_instance = task.to_workflow( From 815bd467e4bc81e1e0fe6dcdbb79d44abf73b9ca Mon Sep 17 00:00:00 2001 From: "panxuchen.pxc" Date: Wed, 1 Jul 2026 20:36:59 +0800 Subject: [PATCH 73/84] clean build experience --- trinity/explorer/rollout_coordinator.py | 20 ++++++++------------ 1 file changed, 8 insertions(+), 12 deletions(-) diff --git a/trinity/explorer/rollout_coordinator.py b/trinity/explorer/rollout_coordinator.py index a0044108202..247fc87997f 100644 --- a/trinity/explorer/rollout_coordinator.py +++ b/trinity/explorer/rollout_coordinator.py @@ -319,14 +319,6 @@ async def _finalize_train_batch( return self._finish_batch(batch_state, pipeline_metrics=pipeline_metrics) async def _discard_recorded_experiences(self, prefix: str) -> None: -<<<<<<< HEAD - """Delete recorded experiences matching a prefix from all rollout ranks.""" - results = await asyncio.gather( - *[ - actor.delete_experience_records.remote(prefix=prefix) - for actor in self._rollout_actors.values() - ], -======= """Block future writes and delete recorded experiences for a prefix. Blocking happens before deleting across all rollout ranks so that any @@ -335,14 +327,18 @@ async def _discard_recorded_experiences(self, prefix: str) -> None: rollout actor (batch_id is never reused), so the prefix stays unwritable for the lifetime of the process. """ - actors = self._resolve_rollout_actors() block_results = await asyncio.gather( - *[actor.block_experience_records.remote(prefix=prefix) for actor in actors.values()], + *[ + actor.block_experience_records.remote(prefix=prefix) + for actor in self._rollout_actors.values() + ], return_exceptions=True, ) delete_results = await asyncio.gather( - *[actor.delete_experience_records.remote(prefix=prefix) for actor in actors.values()], ->>>>>>> 1146972c33ab4417d0c6b15abe50d75d3136ee79 + *[ + actor.delete_experience_records.remote(prefix=prefix) + for actor in self._rollout_actors.values() + ], return_exceptions=True, ) for result in [*block_results, *delete_results]: From 0571edb93b4390d5844d74ceda1ccffe7314fcfa Mon Sep 17 00:00:00 2001 From: pxc Date: Wed, 1 Jul 2026 20:53:16 +0800 Subject: [PATCH 74/84] fix model version --- trinity/buffer/storage/queue.py | 14 ++++++++++---- trinity/common/models/sglang_model.py | 3 +++ trinity/common/models/vllm_model.py | 1 + 3 files changed, 14 insertions(+), 4 deletions(-) diff --git a/trinity/buffer/storage/queue.py b/trinity/buffer/storage/queue.py index 957ac1e6ab0..d10fd1cd095 100644 --- a/trinity/buffer/storage/queue.py +++ b/trinity/buffer/storage/queue.py @@ -59,7 +59,9 @@ def __init__(self, decay: float = 2.0): self.decay = decay def __call__(self, item: List[Experience]) -> Tuple[float, bool]: - priority = float(item[0].info["model_version"] - self.decay * item[0].info["use_count"]) + priority = float( + item[0].info["model_version"] - self.decay * item[0].info.get("use_count", 0) + ) put_into_queue = True return priority, put_into_queue @@ -82,11 +84,15 @@ def __init__(self, decay: float = 2.0, use_count_limit: int = 3, sigma: float = self.sigma = sigma def __call__(self, item: List[Experience]) -> Tuple[float, bool]: - priority = float(item[0].info["model_version"] - self.decay * item[0].info["use_count"]) + priority = float( + item[0].info["model_version"] - self.decay * item[0].info.get("use_count", 0) + ) if self.sigma > 0.0: priority += float(np.random.randn() * self.sigma) put_into_queue = ( - item[0].info["use_count"] < self.use_count_limit if self.use_count_limit > 0 else True + item[0].info.get("use_count", 0) < self.use_count_limit + if self.use_count_limit > 0 + else True ) return priority, put_into_queue @@ -293,7 +299,7 @@ async def get(self) -> List[Experience]: break for exp in item: - exp.info["use_count"] += 1 + exp.info["use_count"] = exp.info.get("use_count", 0) + 1 # Optionally resubmit the item after a cooldown if self.reuse_cooldown_time is not None: asyncio.create_task(self._put(item, delay=self.reuse_cooldown_time)) diff --git a/trinity/common/models/sglang_model.py b/trinity/common/models/sglang_model.py index b894734929c..7b9e6801ef5 100644 --- a/trinity/common/models/sglang_model.py +++ b/trinity/common/models/sglang_model.py @@ -482,6 +482,9 @@ async def generate( prompt_text=prompt_text, response_text=response_text, routed_experts=routed_experts, + info={ + "model_version": self.model_version, + }, ) ) return experiences diff --git a/trinity/common/models/vllm_model.py b/trinity/common/models/vllm_model.py index 95f7f497fe1..72835ab8474 100644 --- a/trinity/common/models/vllm_model.py +++ b/trinity/common/models/vllm_model.py @@ -334,6 +334,7 @@ async def generate( record_key=None, timestamp="", multi_modal_inputs=multi_modal_inputs, + model_version=self.model_version, prompt_text=self.tokenizer.decode(output.prompt_token_ids), include_routed_experts=self.config.enable_return_routed_experts, include_prompt_routed_experts=True, From 14d5bdc86f1ac37c4ec4c366b127b3ca4b6cd9ad Mon Sep 17 00:00:00 2001 From: "panxuchen.pxc" Date: Wed, 1 Jul 2026 21:01:36 +0800 Subject: [PATCH 75/84] fix model version drift --- .../models/sglang_patch/recording/models.py | 17 +++++++++++++++++ .../models/sglang_patch/recording/recorder.py | 5 +++++ .../models/vllm_patch/recording/models.py | 14 ++++++++++++++ .../models/vllm_patch/recording/recorder.py | 5 +++++ 4 files changed, 41 insertions(+) diff --git a/trinity/common/models/sglang_patch/recording/models.py b/trinity/common/models/sglang_patch/recording/models.py index 6caebc77a59..218f01e5582 100644 --- a/trinity/common/models/sglang_patch/recording/models.py +++ b/trinity/common/models/sglang_patch/recording/models.py @@ -49,6 +49,15 @@ def _sample_suffix(request_id: str, sample_index: int, num_samples: int) -> str: return f"{request_id}:{sample_index}" +def _model_version_drift(start: Optional[Any], end: Optional[Any]) -> int: + if start is None or end is None: + return 0 + try: + return int(end) - int(start) + except (TypeError, ValueError): + return 0 + + def _extract_routed_experts( routed_experts_value: Any, total_tokens: int, @@ -73,6 +82,7 @@ def build_sglang_experience( *, timestamp: str, model_version: Optional[Any] = None, + model_version_start: Optional[Any] = None, include_routed_experts: bool = True, routed_experts_layout: Optional[Tuple[int, int]] = None, ) -> List[Experience]: @@ -93,6 +103,9 @@ def build_sglang_experience( timestamp: UTC ISO-8601 string (caller-stamped to keep this pure). model_version: Checkpoint version fallback; overridden by ``meta_info.weight_version`` when present. + model_version_start: Checkpoint version captured when this generation + entered the rollout engine. Used to compute + ``info["model_version_drift"]``. include_routed_experts: Whether routed experts should be copied. routed_experts_layout: ``(num_layers, topk)`` for decoding base64-str routed experts (from ``BaseInferenceModel._get_routed_experts_layout``). @@ -160,6 +173,10 @@ def build_sglang_experience( "sample_index": sample_index, "timestamp": timestamp, "model_version": resolved_model_version, + "model_version_drift": _model_version_drift( + model_version_start, + resolved_model_version, + ), } experiences.append( diff --git a/trinity/common/models/sglang_patch/recording/recorder.py b/trinity/common/models/sglang_patch/recording/recorder.py index c602a707757..f82eb592d55 100644 --- a/trinity/common/models/sglang_patch/recording/recorder.py +++ b/trinity/common/models/sglang_patch/recording/recorder.py @@ -33,6 +33,7 @@ get_recording_record_key_from_context, ) from trinity.common.models.recording.recorder import ( + MODEL_VERSION_ATTR, TRINITY_RECORD_STORE_ATTR, TRINITY_RECORDER_ATTR, Recorder, @@ -267,6 +268,9 @@ async def _patched_generate_request(*args, **kwargs): state: dict = {} order: list = [] + model_version_start = ( + getattr(tokenizer_manager, MODEL_VERSION_ATTR, None) if recorder.enabled else None + ) # ``current`` is the original *bound* method captured pre-wrap, so it # still resolves ``self`` correctly. Yields each ret unchanged. async for out in current(*args, **kwargs): @@ -281,6 +285,7 @@ async def _patched_generate_request(*args, **kwargs): recorder.schedule_record( reconstructed, record_key, + model_version_start=model_version_start, include_routed_experts=True, routed_experts_layout=routed_experts_layout, ) diff --git a/trinity/common/models/vllm_patch/recording/models.py b/trinity/common/models/vllm_patch/recording/models.py index f9ff129c96a..f578018d747 100644 --- a/trinity/common/models/vllm_patch/recording/models.py +++ b/trinity/common/models/vllm_patch/recording/models.py @@ -88,6 +88,15 @@ def _sample_suffix(request_id: str, sample_index: int, num_samples: int) -> str: return f"{request_id}:{sample_index}" +def _model_version_drift(start: Optional[Any], end: Optional[Any]) -> int: + if start is None or end is None: + return 0 + try: + return int(end) - int(start) + except (TypeError, ValueError): + return 0 + + def _extract_routed_experts( output: Any, completion: Any, @@ -121,6 +130,7 @@ def build_experience( *, timestamp: str, model_version: Optional[int] = None, + model_version_start: Optional[Any] = None, multi_modal_inputs: Optional[dict] = None, prompt_text: Optional[str] = None, include_routed_experts: bool = True, @@ -140,6 +150,9 @@ def build_experience( timestamp: UTC ISO-8601 string (caller-stamped to keep this pure). model_version: Checkpoint version the serving policy was at; stamped into ``info`` for RL attribution (read in-actor by the recorder). + model_version_start: Checkpoint version captured when this generation + entered the rollout engine. Used to compute + ``info["model_version_drift"]``. multi_modal_inputs: Optional training-time multimodal tensors aligned with the prompt tokens. Response token type ids are appended per completion before storing on the ``Experience``. @@ -196,6 +209,7 @@ def build_experience( "sample_index": sample_index, "timestamp": timestamp, "model_version": model_version, + "model_version_drift": _model_version_drift(model_version_start, model_version), } experiences.append( diff --git a/trinity/common/models/vllm_patch/recording/recorder.py b/trinity/common/models/vllm_patch/recording/recorder.py index 8ae841bff5b..0bcc3711e70 100644 --- a/trinity/common/models/vllm_patch/recording/recorder.py +++ b/trinity/common/models/vllm_patch/recording/recorder.py @@ -29,6 +29,7 @@ get_recording_request_from_context, ) from trinity.common.models.recording.recorder import ( + MODEL_VERSION_ATTR, TRINITY_RECORD_STORE_ATTR, TRINITY_RECORDER_ATTR, Recorder, @@ -288,6 +289,9 @@ async def _patched_generate(*args, **kwargs): sampling_params.logprobs = ( max(cur, _RECORDER_LOGPROB_WIDTH) if cur is not None else _RECORDER_LOGPROB_WIDTH ) + model_version_start = ( + getattr(engine_client, MODEL_VERSION_ATTR, None) if recorder.enabled else None + ) is_delta_output = _is_delta_output(sampling_params) last = None @@ -321,6 +325,7 @@ async def _patched_generate(*args, **kwargs): recorder.schedule_record( record_output, record_key, + model_version_start=model_version_start, multi_modal_inputs=multi_modal_inputs, ) From 740805af517cb66a3d0ce718deba1d2d6bd1d524 Mon Sep 17 00:00:00 2001 From: pxc Date: Wed, 1 Jul 2026 21:37:24 +0800 Subject: [PATCH 76/84] fix workflow reset --- trinity/common/workflows/workflow.py | 10 +++------- trinity/explorer/workflow_runner.py | 2 ++ 2 files changed, 5 insertions(+), 7 deletions(-) diff --git a/trinity/common/workflows/workflow.py b/trinity/common/workflows/workflow.py index 6032007683c..e0099138217 100644 --- a/trinity/common/workflows/workflow.py +++ b/trinity/common/workflows/workflow.py @@ -107,8 +107,7 @@ class WorkflowBase: def __init__(self, task: Task, model: ModelWrapper) -> None: self.task = task self.model = model - if self.model is not None: - self.model.set_api_key(task.api_key) # set the API key for the rollout model + self.model.set_api_key(task.api_key) # set the API key for the rollout model self.logger = get_logger(__name__) @abstractmethod @@ -117,11 +116,8 @@ async def execute(self) -> Status: def reset(self, task: Task): """Reset the workflow with a new task.""" - self.task.batch_id = task.batch_id - self.task.task_id = task.task_id - self.task.run_id = task.run_id - if self.model is not None: - self.model.set_api_key(task.api_key) # set the API key for the rollout model + self.task = task + self.model.set_api_key(task.api_key) # set the API key for the rollout model class Workflow(WorkflowBase): diff --git a/trinity/explorer/workflow_runner.py b/trinity/explorer/workflow_runner.py index f1abfca9064..051f6501397 100644 --- a/trinity/explorer/workflow_runner.py +++ b/trinity/explorer/workflow_runner.py @@ -99,6 +99,8 @@ def _create_workflow_instance(self, task: Task) -> Workflow: ) else: self.workflow_instance.reset(task) + self.workflow_instance.task = task + self.workflow_instance.model.set_api_key(task.api_key) self.workflow_instance.set_single_run_context(task.run_id) return self.workflow_instance From 40a1dd69ca17515af52b4ec8191c696ff598918b Mon Sep 17 00:00:00 2001 From: pxc Date: Thu, 2 Jul 2026 10:28:26 +0800 Subject: [PATCH 77/84] fix comments --- tests/explorer/step_wise_workflow_test.py | 14 +++++++------- tests/explorer/workflow_test.py | 8 ++++---- tests/utils/plugins/my_workflow.py | 6 ++---- trinity/common/config.py | 2 +- trinity/common/workflows/workflow.py | 17 +++++------------ 5 files changed, 19 insertions(+), 28 deletions(-) diff --git a/tests/explorer/step_wise_workflow_test.py b/tests/explorer/step_wise_workflow_test.py index 389b8ec01cb..8a80be53e18 100644 --- a/tests/explorer/step_wise_workflow_test.py +++ b/tests/explorer/step_wise_workflow_test.py @@ -133,7 +133,7 @@ def test_step_wise_reward_workflow(self, workflow_cls) -> None: workflow_args={"max_env_steps": 10, "actual_steps": 5}, ) workflow = task.to_workflow(model=self.model) - if workflow.asynchronous: + if workflow.is_async: experiences = asyncio.run(workflow.run_async()) else: experiences = workflow.run() @@ -156,7 +156,7 @@ def test_reward_propagation_workflow(self, workflow_cls) -> None: workflow_args={"max_env_steps": 10, "actual_steps": 5}, ) workflow = task.to_workflow(model=self.model) - if workflow.asynchronous: + if workflow.is_async: experiences = asyncio.run(workflow.run_async()) else: experiences = workflow.run() @@ -169,14 +169,14 @@ def test_reward_propagation_workflow(self, workflow_cls) -> None: self.assertAlmostEqual(exp.reward, expected_reward) # type: ignore def test_workflows_stop_at_max_env_steps(self) -> None: - for workflow in _dummy_workflows: + for workflow_cls in _dummy_workflows: task = Task( - workflow=workflow, + workflow=workflow_cls, repeat_times=self.taskset_config.repeat_times, workflow_args={"max_env_steps": 3, "actual_steps": 100}, # actual > max ) workflow = task.to_workflow(model=self.model) - if workflow.asynchronous: + if workflow.is_async: experiences = asyncio.run(workflow.run_async()) # type: ignore else: experiences = workflow.run() @@ -184,9 +184,9 @@ def test_workflows_stop_at_max_env_steps(self) -> None: def test_workflows_raise_error(self) -> None: self.model.enable_history = False - for workflow in _dummy_workflows: + for workflow_cls in _dummy_workflows: task = Task( - workflow=workflow, + workflow=workflow_cls, repeat_times=self.taskset_config.repeat_times, workflow_args={"max_env_steps": 10, "actual_steps": 5}, ) diff --git a/tests/explorer/workflow_test.py b/tests/explorer/workflow_test.py index a3508572a15..0ff70fd4c88 100644 --- a/tests/explorer/workflow_test.py +++ b/tests/explorer/workflow_test.py @@ -426,13 +426,13 @@ def test_workflow_resettable(self, workflow_cls) -> None: workflow_args={"output_format": "yaml"}, ) workflow = json_task.to_workflow(model) - if workflow.asynchronous: + if workflow.is_async: answer = asyncio.run(workflow.run_async()) else: answer = workflow.run() self.assertEqual(answer[0].response_text, '{"a": 1}') workflow.reset(yaml_task) - if workflow.asynchronous: + if workflow.is_async: answer = asyncio.run(workflow.run_async()) else: answer = workflow.run() @@ -450,7 +450,7 @@ def test_workflow_repeatable(self, workflow_cls) -> None: workflow = task.to_workflow(model) workflow.set_repeat_times(2, run_id_base=0) self.assertEqual(workflow.repeat_times, 2) - if workflow.asynchronous: + if workflow.is_async: answer = asyncio.run(workflow.run_async()) else: answer = workflow.run() @@ -490,7 +490,7 @@ async def test_multi_turn_workflow(self): ) workflow = task.to_workflow(self.model_wrapper) workflow.set_repeat_times(2, run_id_base=0) - if workflow.asynchronous: + if workflow.is_async: answer = await workflow.run_async() else: answer = workflow.run() diff --git a/tests/utils/plugins/my_workflow.py b/tests/utils/plugins/my_workflow.py index 471b2371cff..1a321523cda 100644 --- a/tests/utils/plugins/my_workflow.py +++ b/tests/utils/plugins/my_workflow.py @@ -5,13 +5,11 @@ @WORKFLOWS.register_module("my_workflow") class MyWorkflow(Workflow): + can_repeat = True + def __init__(self, *, task, model, auxiliary_models=None): super().__init__(task=task, model=model, auxiliary_models=auxiliary_models) - @property - def repeatable(self): - return True - def set_repeat_times(self, repeat_times, run_id_base): pass diff --git a/trinity/common/config.py b/trinity/common/config.py index 5e1632d12ce..dd32b445b1a 100644 --- a/trinity/common/config.py +++ b/trinity/common/config.py @@ -588,7 +588,7 @@ class InferenceModelConfig: # ``ConfigValidator`` forces this to ``True`` for both the rollout model and # auxiliary models regardless of any user-supplied value. The field is kept # only for backward compatibility with existing YAML configs. - enable_openai_api: bool = False + enable_openai_api: bool = True enable_log_requests: bool = False # whether to enable request logging in vLLM API server base_port: Optional[int] = None api_key: str = "EMPTY" diff --git a/trinity/common/workflows/workflow.py b/trinity/common/workflows/workflow.py index e0099138217..68764309197 100644 --- a/trinity/common/workflows/workflow.py +++ b/trinity/common/workflows/workflow.py @@ -147,19 +147,13 @@ def __init__( # Get OpenAI clients from ModelWrapper (async or sync based on workflow type) self.auxiliary_models: Optional[Union[List[openai.OpenAI], List[openai.AsyncOpenAI]]] = None if auxiliary_models: - if self.__class__.is_async: + if self.is_async: self.auxiliary_models = [m.get_openai_async_client() for m in auxiliary_models] else: self.auxiliary_models = [m.get_openai_client() for m in auxiliary_models] self.run_id_base = 0 self.repeat_times = 1 - @property - def asynchronous(self): - """Deprecated, use cls.is_async instead. - Whether the workflow runs in async mode.""" - return self.__class__.is_async - def set_repeat_times(self, repeat_times: int, run_id_base: int) -> None: """ Set the number of times to repeat the workflow. @@ -189,7 +183,7 @@ async def run_async(self) -> List[Experience]: raise NotImplementedError async def execute(self) -> Status: - if self.asynchronous: + if self.is_async: exps = await self.run_async() else: exps = self.run() @@ -197,8 +191,8 @@ async def execute(self) -> Status: experiences=exps, key=self.task.api_key ) return Status( - completed_runs=self.__class__.can_repeat and self.repeat_times or 1, - total_runs=self.__class__.can_repeat and self.repeat_times or 1, + completed_runs=self.can_repeat and self.repeat_times or 1, + total_runs=self.can_repeat and self.repeat_times or 1, metrics=[exp.metrics for exp in exps if exp.metrics is not None], successful_ids=[self.task.api_key], ) @@ -309,8 +303,7 @@ def reset(self, task: Task): raise ValueError("`reward_fn` must be a subclass of `RewardFn`") def set_repeat_times(self, repeat_times, run_id_base): - self.repeat_times = repeat_times - self.run_id_base = run_id_base + super().set_repeat_times(repeat_times, run_id_base) self.task.rollout_args.n = repeat_times @property From a66afc81d9faa418af677e68733f882d45d78a18 Mon Sep 17 00:00:00 2001 From: pxc Date: Thu, 2 Jul 2026 13:48:40 +0800 Subject: [PATCH 78/84] fix workflow --- tests/explorer/workflow_test.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/explorer/workflow_test.py b/tests/explorer/workflow_test.py index 0ff70fd4c88..60c47a32e93 100644 --- a/tests/explorer/workflow_test.py +++ b/tests/explorer/workflow_test.py @@ -904,7 +904,7 @@ async def run_async(self): class ConcurrentRecordingWorkflow(WorkflowWithRecording): def __init__(self, model: ModelWrapper, task: Task, auxiliary_models=None): super().__init__(task=task, model=model, auxiliary_models=auxiliary_models) - self.client = openai.AsyncOpenAI(base_url=f"{self.base_url}/v1", api_key=self.api_key) + self.client = openai.AsyncOpenAI(base_url=self.base_url, api_key=self.api_key) def reset(self, task: Task): self.task = task From 2167e5315245ac0fd71cf30a3105ff49c952b6ce Mon Sep 17 00:00:00 2001 From: pxc Date: Thu, 2 Jul 2026 13:55:42 +0800 Subject: [PATCH 79/84] fix plugin tests --- tests/utils/plugin_test.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/tests/utils/plugin_test.py b/tests/utils/plugin_test.py index 3bda6754aef..09a3fea805c 100644 --- a/tests/utils/plugin_test.py +++ b/tests/utils/plugin_test.py @@ -3,6 +3,7 @@ import unittest from pathlib import Path from typing import Type +from unittest.mock import MagicMock import ray from parameterized import parameterized @@ -10,7 +11,7 @@ from tests.tools import TensorBoardParser, get_checkpoint_path, get_template_config from trinity.common.config import Config from trinity.common.constants import PLUGIN_DIRS_ENV_VAR -from trinity.common.workflows import WORKFLOWS, Workflow +from trinity.common.workflows import WORKFLOWS, Task, Workflow from trinity.utils.monitor import MONITOR from trinity.utils.plugin_loader import load_plugins @@ -43,7 +44,7 @@ def __init__( def run(self, workflow_cls=Type[Workflow]): if self.monitor: self.monitor.log({"rollout": 2}, step=1, commit=True) - return workflow_cls(task=None, model=None).run() + return workflow_cls(task=Task(), model=MagicMock()).run() class TestPluginLoader(unittest.TestCase): @@ -65,7 +66,7 @@ def test_load_plugins_local(self, plugin_dir): pass my_workflow_cls = WORKFLOWS.get("my_workflow") self.assertIsNotNone(my_workflow_cls) - my_plugin = my_workflow_cls(task=None, model=None, auxiliary_models=None) + my_plugin = my_workflow_cls(task=Task(), model=MagicMock(), auxiliary_models=None) self.assertTrue(my_plugin.__module__.startswith("trinity.plugins")) res = my_plugin.run() self.assertEqual(res[0], "Hello world") From fa7b2698b20d2bdb1bde52fcbfb3d7a97ccd76d5 Mon Sep 17 00:00:00 2001 From: pxc Date: Thu, 2 Jul 2026 14:00:52 +0800 Subject: [PATCH 80/84] fix registry --- tests/utils/registry_test.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/tests/utils/registry_test.py b/tests/utils/registry_test.py index e028cc2a0f3..3937f37dc82 100644 --- a/tests/utils/registry_test.py +++ b/tests/utils/registry_test.py @@ -1,6 +1,7 @@ # -*- coding: utf-8 -*- """Test cases for workflows registry mapping.""" import unittest +from unittest.mock import MagicMock import ray import torch @@ -27,7 +28,7 @@ from trinity.buffer.storage import PRIORITY_FUNC from trinity.buffer.storage.queue import PriorityFunction from trinity.common.rewards import REWARD_FUNCTIONS, RewardFn -from trinity.common.workflows import WORKFLOWS, Workflow +from trinity.common.workflows import WORKFLOWS, Task, Workflow from trinity.utils.monitor import MONITOR, Monitor @@ -51,7 +52,7 @@ def run(self): workflow_cls = WORKFLOWS.get("tests.utils.plugins.main.MainDummyWorkflow") assert issubclass(workflow_cls, Workflow) - workflow = workflow_cls(task=None, model=None) + workflow = workflow_cls(task=Task(), model=MagicMock()) res = workflow.run() assert res[0] == 0 assert res[1] == "0" From 8347dd8967a85d7128085c3fd96370cc19a01c33 Mon Sep 17 00:00:00 2001 From: "panxuchen.pxc" Date: Thu, 2 Jul 2026 14:13:19 +0800 Subject: [PATCH 81/84] finish harbor workflow --- trinity/common/workflows/__init__.py | 2 + trinity/common/workflows/harbor_workflow.py | 286 ++++++++++++++++++++ 2 files changed, 288 insertions(+) create mode 100644 trinity/common/workflows/harbor_workflow.py diff --git a/trinity/common/workflows/__init__.py b/trinity/common/workflows/__init__.py index 140fbcd2217..d8530001dfc 100644 --- a/trinity/common/workflows/__init__.py +++ b/trinity/common/workflows/__init__.py @@ -51,6 +51,8 @@ "math_trainable_ruler_workflow": "trinity.common.workflows.math_trainable_ruler_workflow.MathTrainableRULERWorkflow", "rubric_judge_workflow": "trinity.common.workflows.rubric_judge_workflow.RubricJudgeWorkflow", "rubric_judge_openai_workflow": "trinity.common.workflows.rubric_judge_workflow.RubricJudgeWorkflowWithAPI", + # harbor + "harbor_workflow": "trinity.common.workflows.harbor_workflow.HarborWorkflow", # others "simple_mm_workflow": "trinity.common.workflows.simple_mm_workflow.SimpleMMWorkflow", "async_simple_mm_workflow": "trinity.common.workflows.simple_mm_workflow.AsyncSimpleMMWorkflow", diff --git a/trinity/common/workflows/harbor_workflow.py b/trinity/common/workflows/harbor_workflow.py new file mode 100644 index 00000000000..dca80b567ac --- /dev/null +++ b/trinity/common/workflows/harbor_workflow.py @@ -0,0 +1,286 @@ +# -*- coding: utf-8 -*- +"""Base workflow for Harbor directory tasks.""" + +from __future__ import annotations + +from pathlib import Path +from typing import TYPE_CHECKING, Any, List, Optional + +from trinity.common.workflows.workflow import Metrics, Task, WorkflowWithRecording + +if TYPE_CHECKING: + from harbor.models.task.config import TaskConfig + from harbor.models.task.paths import TaskPaths + from harbor.models.trial.config import TrialConfig + from harbor.models.trial.result import TrialResult + from harbor.viewer.task_scanner import TaskDefinitionScanner + + from trinity.common.models.model import ModelWrapper + + +class HarborWorkflow(WorkflowWithRecording): + """Workflow that runs a Harbor task and writes the verifier reward back. + + ``task.raw_task["task_dir"]`` points at one Harbor task directory. The + workflow builds one Harbor trial, injects Trinity's rollout OpenAI endpoint + into the Harbor agent, runs the trial, and then calls ``update_reward`` so + the rollout model's recorded experiences are labeled with Harbor's verifier + reward. + """ + + DEFAULT_AGENT_NAME = "openclaw" + DEFAULT_TRIALS_DIR = "harbor_trials" + DEFAULT_REWARD_KEY = "reward" + + def __init__( + self, + *, + task: Task, + model: ModelWrapper, + auxiliary_models: Optional[List[ModelWrapper]] = None, + ): + super().__init__( + task=task, + model=model, + auxiliary_models=auxiliary_models, + ) + self.harbor_task_dir = self._get_harbor_task_dir(task) + self.harbor_task_name = self.harbor_task_dir.name + ( + self.harbor_scanner, + self.harbor_task_paths, + self.harbor_task_config, + self.harbor_instruction, + self.harbor_task_paths_info, + ) = self._load_harbor_task(self.harbor_task_dir) + + def _get_harbor_task_dir(self, task: Task) -> Path: + if task.raw_task is None: + raise ValueError("HarborWorkflow requires `task.raw_task` to be configured.") + + task_dir = task.raw_task.get("task_dir") + if task_dir is None: + raise ValueError("HarborWorkflow requires `task.raw_task['task_dir']`.") + + task_dir_path = Path(task_dir).expanduser().resolve() + if not task_dir_path.exists(): + raise FileNotFoundError(f"Harbor task directory does not exist: {task_dir_path}") + if not task_dir_path.is_dir(): + raise ValueError(f"Harbor task path must be a directory: {task_dir_path}") + return task_dir_path + + def _load_harbor_task( + self, + task_dir: Path, + ) -> tuple["TaskDefinitionScanner", "TaskPaths", "TaskConfig", str | None, dict[str, bool]]: + try: + from harbor.models.task.paths import TaskPaths + from harbor.viewer.task_scanner import TaskDefinitionScanner + except ImportError as exc: + raise ImportError( + "HarborWorkflow requires the `harbor` package to be installed." + ) from exc + + scanner = TaskDefinitionScanner(task_dir.parent) + task_name = task_dir.name + config = scanner.get_task_config(task_name) + if config is None: + raise ValueError( + f"Failed to load Harbor task config from: {task_dir / TaskPaths.CONFIG_FILENAME}" + ) + + paths = TaskPaths(task_dir) + instruction = scanner.get_instruction(task_name) + paths_info = scanner.get_task_paths_info(task_name) + return scanner, paths, config, instruction, paths_info + + async def run_async(self) -> Metrics: + trial_config = self._build_trial_config() + trial_result = await self._run_harbor_trial(trial_config) + reward, reward_info = self._extract_reward(trial_result) + info = self._build_reward_info(trial_result, reward_info) + await self.update_reward(reward, info=info) + return self._build_metrics(trial_result, reward) + + def _build_trial_config(self) -> "TrialConfig": + try: + from harbor.models.trial.config import AgentConfig, EnvironmentConfig + from harbor.models.trial.config import TaskConfig as TrialTaskConfig + from harbor.models.trial.config import TrialConfig, VerifierConfig + except ImportError as exc: + raise ImportError( + "HarborWorkflow requires the `harbor` package to be installed." + ) from exc + + workflow_args = self.task.workflow_args or {} + return TrialConfig( + task=TrialTaskConfig(path=self.harbor_task_dir), + trial_name=self._get_optional_str_arg("trial_name"), + trials_dir=self._get_trials_dir(), + agent=self._build_agent_config(AgentConfig), + environment=self._build_model_config( + EnvironmentConfig, workflow_args.get("environment") + ), + verifier=self._build_model_config(VerifierConfig, workflow_args.get("verifier")), + artifacts=list(workflow_args.get("artifacts", [])), + extra_instruction_paths=[ + Path(path).expanduser() for path in workflow_args.get("extra_instruction_paths", []) + ], + timeout_multiplier=workflow_args.get("timeout_multiplier", 1.0), + agent_timeout_multiplier=workflow_args.get("agent_timeout_multiplier"), + verifier_timeout_multiplier=workflow_args.get("verifier_timeout_multiplier"), + agent_setup_timeout_multiplier=workflow_args.get("agent_setup_timeout_multiplier"), + environment_build_timeout_multiplier=workflow_args.get( + "environment_build_timeout_multiplier" + ), + ) + + def _build_agent_config(self, agent_config_cls): + workflow_args = self.task.workflow_args or {} + agent_env = dict(workflow_args.get("agent_env", {})) + agent_env.update(self._rollout_env()) + agent_name = workflow_args.get("agent_name") + agent_import_path = workflow_args.get("agent_import_path") + if agent_name is None and agent_import_path is None: + agent_name = self.DEFAULT_AGENT_NAME + + return agent_config_cls( + name=agent_name, + import_path=agent_import_path, + model_name=workflow_args.get("harbor_model_name", self.model_name), + kwargs=dict(workflow_args.get("agent_kwargs", {})), + env=agent_env, + ) + + def _build_model_config(self, config_cls, raw_config: Any): + if raw_config is None: + return config_cls() + if isinstance(raw_config, config_cls): + return raw_config + if isinstance(raw_config, dict): + return config_cls(**raw_config) + raise TypeError( + f"Expected {config_cls.__name__} config as a dict or {config_cls.__name__}, " + f"got {type(raw_config).__name__}." + ) + + def _rollout_env(self) -> dict[str, str]: + workflow_args = self.task.workflow_args or {} + base_url = workflow_args.get("rollout_base_url", self.base_url) + api_key = workflow_args.get("rollout_api_key", self.api_key) + model_name = workflow_args.get("harbor_model_name", self.model_name) + env = { + "OPENAI_BASE_URL": str(base_url), + "OPENAI_API_KEY": str(api_key), + "OPENAI_MODEL": str(model_name), + } + + provider_env = workflow_args.get("provider_env", {}) + if provider_env: + env.update({str(key): str(value) for key, value in provider_env.items()}) + return env + + def _get_trials_dir(self) -> Path: + workflow_args = self.task.workflow_args or {} + trials_dir = workflow_args.get("trials_dir", self.DEFAULT_TRIALS_DIR) + return Path(trials_dir).expanduser() + + def _get_optional_str_arg(self, key: str) -> str: + value = (self.task.workflow_args or {}).get(key, "") + return "" if value is None else str(value) + + async def _run_harbor_trial(self, trial_config: "TrialConfig") -> "TrialResult": + try: + from harbor import Trial + except ImportError as exc: + raise ImportError( + "HarborWorkflow requires the `harbor` package to be installed." + ) from exc + + trial = await Trial.create(trial_config) + return await trial.run() + + def _extract_reward(self, trial_result: "TrialResult") -> tuple[float, dict[str, Any]]: + rewards = self._get_harbor_rewards(trial_result) + if not rewards: + return 0.0, { + "harbor_reward_missing": True, + "harbor_reward_reason": "verifier_result.rewards is empty", + } + + reward_key = self._get_reward_key(rewards) + reward = rewards[reward_key] + return float(reward), {"harbor_reward_key": reward_key} + + def _get_harbor_rewards(self, trial_result: "TrialResult") -> dict[str, float | int]: + verifier_result = trial_result.verifier_result + if verifier_result is None or verifier_result.rewards is None: + return {} + return dict(verifier_result.rewards) + + def _get_reward_key(self, rewards: dict[str, float | int]) -> str: + workflow_args = self.task.workflow_args or {} + configured_key = workflow_args.get("reward_key") + if configured_key: + reward_key = str(configured_key) + if reward_key not in rewards: + raise KeyError( + f"Harbor reward key {reward_key!r} not found. " + f"Available reward keys: {sorted(rewards)}" + ) + return reward_key + + if self.DEFAULT_REWARD_KEY in rewards: + return self.DEFAULT_REWARD_KEY + if len(rewards) == 1: + return next(iter(rewards)) + raise ValueError( + "Harbor verifier returned multiple rewards. Set " + "`workflow_args.reward_key` to choose one. " + f"Available reward keys: {sorted(rewards)}" + ) + + def _build_reward_info( + self, + trial_result: "TrialResult", + reward_info: dict[str, Any], + ) -> dict[str, Any]: + info: dict[str, Any] = { + "harbor_task_name": trial_result.task_name, + "harbor_trial_name": trial_result.trial_name, + "harbor_trial_uri": trial_result.trial_uri, + "harbor_rewards": self._get_harbor_rewards(trial_result), + **reward_info, + } + if trial_result.exception_info is not None: + info.update( + { + "harbor_exception_type": trial_result.exception_info.exception_type, + "harbor_exception_message": (trial_result.exception_info.exception_message), + } + ) + return info + + def _build_metrics(self, trial_result: "TrialResult", reward: float) -> Metrics: + rewards = self._get_harbor_rewards(trial_result) + metrics: Metrics = { + "harbor/reward": float(reward), + "harbor/has_exception": (1.0 if trial_result.exception_info is not None else 0.0), + "harbor/num_rewards": float(len(rewards)), + } + + ( + input_tokens, + cache_tokens, + output_tokens, + cost_usd, + ) = trial_result.compute_token_cost_totals() + if input_tokens is not None: + metrics["harbor/input_tokens"] = float(input_tokens) + if cache_tokens is not None: + metrics["harbor/cache_tokens"] = float(cache_tokens) + if output_tokens is not None: + metrics["harbor/output_tokens"] = float(output_tokens) + if cost_usd is not None: + metrics["harbor/cost_usd"] = float(cost_usd) + return metrics From f010e6462c2dbdf372b991c3f7049b0103202e41 Mon Sep 17 00:00:00 2001 From: pxc Date: Thu, 2 Jul 2026 16:55:52 +0800 Subject: [PATCH 82/84] move recorder to a separate thread --- tests/common/recording_recorder_test.py | 38 ++-- trinity/buffer/store/memory_store.py | 184 +++++++++++--------- trinity/common/models/recording/recorder.py | 55 ++++-- 3 files changed, 162 insertions(+), 115 deletions(-) diff --git a/tests/common/recording_recorder_test.py b/tests/common/recording_recorder_test.py index 05820aee53e..51ae6cafa75 100644 --- a/tests/common/recording_recorder_test.py +++ b/tests/common/recording_recorder_test.py @@ -51,8 +51,8 @@ async def test_prefix_experiences_merge_and_keep_final_sample_id(self): logprobs=[-0.4, -0.5, -0.6], ) - await recorder._safe_append(first) - await recorder._safe_append(second) + recorder._safe_append(first) + recorder._safe_append(second) recorded = store.get(record_key) self.assertEqual(len(recorded), 1) @@ -89,7 +89,7 @@ async def test_non_prefix_experiences_do_not_merge(self): ) record_key = "0/task_a/1" - await recorder._safe_append( + recorder._safe_append( make_turn( request_id="req-1", record_key=record_key, @@ -98,7 +98,7 @@ async def test_non_prefix_experiences_do_not_merge(self): logprobs=[-0.2], ) ) - await recorder._safe_append( + recorder._safe_append( make_turn( request_id="req-2", record_key=record_key, @@ -143,9 +143,9 @@ async def test_merge_head_replaces_only_matching_sample_stream(self): sample_index=1, ) - await recorder._safe_append(sample_zero) - await recorder._safe_append(sample_one_first) - await recorder._safe_append(sample_one_final) + recorder._safe_append(sample_zero) + recorder._safe_append(sample_one_first) + recorder._safe_append(sample_one_final) recorded = store.get(record_key) self.assertEqual(len(recorded), 2) @@ -199,10 +199,10 @@ async def test_interleaved_branches_with_shared_sample_index_merge_independently sample_index=0, ) - await recorder._safe_append(branch_a_first) - await recorder._safe_append(branch_b_first) - await recorder._safe_append(branch_a_final) - await recorder._safe_append(branch_b_final) + recorder._safe_append(branch_a_first) + recorder._safe_append(branch_b_first) + recorder._safe_append(branch_a_final) + recorder._safe_append(branch_b_final) recorded = store.get(record_key) self.assertEqual(len(recorded), 2) @@ -252,10 +252,10 @@ async def test_multi_head_merge_uses_longest_matching_prefix(self): logprobs=[-0.5], ) - await recorder._safe_append(short_prefix) - await recorder._safe_append(long_prefix) - await recorder._safe_append(unrelated) - await recorder._safe_append(final) + recorder._safe_append(short_prefix) + recorder._safe_append(long_prefix) + recorder._safe_append(unrelated) + recorder._safe_append(final) recorded = store.get(record_key) self.assertEqual(len(recorded), 2) @@ -286,8 +286,8 @@ async def test_same_prompt_independent_completions_do_not_merge(self): logprobs=[-0.4, -0.5, -0.6, -0.7], ) - await recorder._safe_append(short_completion) - await recorder._safe_append(long_completion) + recorder._safe_append(short_completion) + recorder._safe_append(long_completion) recorded = store.get(record_key) self.assertEqual(len(recorded), 2) @@ -317,9 +317,9 @@ async def test_stale_merge_head_falls_back_to_append(self): logprobs=[-0.3], ) - await recorder._safe_append(first) + recorder._safe_append(first) store.remove(record_key) - await recorder._safe_append(second) + recorder._safe_append(second) recorded = store.get(record_key) self.assertEqual(len(recorded), 1) diff --git a/trinity/buffer/store/memory_store.py b/trinity/buffer/store/memory_store.py index e4b875e7b22..ec94d64ba76 100644 --- a/trinity/buffer/store/memory_store.py +++ b/trinity/buffer/store/memory_store.py @@ -1,6 +1,7 @@ """In-memory implementation of the experience store interface.""" import logging +import threading from collections import OrderedDict from typing import Iterable, List @@ -50,6 +51,12 @@ class MemoryStore(RecordStore): ``add``, ``overwrite`` and ``update`` require complete keys in the form ``//``. ``get`` and ``remove`` also accept prefixes so callers can drain a batch or task at once. + + A coarse reentrant lock guards every public method: the store may be written + to from the recorder's background flusher thread while the engine event loop + reads/updates it (e.g. ``VLLMModel.generate`` short-circuit adds, actor + ``block_prefix``/``update`` calls). The lock is uncontended on the fast path + (single consumer) and only serializes the rare cross-thread overlap. """ def __init__(self) -> None: @@ -63,66 +70,73 @@ def __init__(self) -> None: # finalized batches); see ``block_prefix``. Only grows since batch_id # is never reused. self._blocked_batches: set[str] = set() + # Serializes all dict mutations/reads across the flusher thread and the + # engine event loop. Reentrant so ``overwrite`` -> ``add`` works. + self._lock = threading.RLock() def __len__(self) -> int: - return sum(len(exps) for exps in self._records.values()) + with self._lock: + return sum(len(exps) for exps in self._records.values()) def add(self, key: str, exps: List[Experience]) -> None: - batch, task, _ = self._parse_complete_key(key) # validate key format - if batch in self._blocked_batches: - _logger.debug( - "Dropping write to blocked batch '%s' (key=%s, %d exps).", - batch, - key, - len(exps), - ) - return - if not exps: - return - - records = self._records.setdefault(key, OrderedDict()) - self._index_key(batch, task, key) - for exp in exps: - sample_id = get_sample_id(exp) - owner_key = self._sample_to_key.get(sample_id) - if owner_key is not None: - raise ValueError( - f"Duplicate sample_id '{sample_id}' already exists under key '{owner_key}'." + with self._lock: + batch, task, _ = self._parse_complete_key(key) # validate key format + if batch in self._blocked_batches: + _logger.debug( + "Dropping write to blocked batch '%s' (key=%s, %d exps).", + batch, + key, + len(exps), ) - records[sample_id] = exp - self._sample_to_key[sample_id] = key + return + if not exps: + return + + records = self._records.setdefault(key, OrderedDict()) + self._index_key(batch, task, key) + for exp in exps: + sample_id = get_sample_id(exp) + owner_key = self._sample_to_key.get(sample_id) + if owner_key is not None: + raise ValueError( + f"Duplicate sample_id '{sample_id}' already exists under key '{owner_key}'." + ) + records[sample_id] = exp + self._sample_to_key[sample_id] = key def overwrite(self, key: str, exps: List[Experience]) -> None: - self._parse_complete_key(key) # validate key format - self._drop_key(key) - self.add(key, exps) + with self._lock: + self._parse_complete_key(key) # validate key format + self._drop_key(key) + self.add(key, exps) def replace(self, key: str, old_sample_id: str, exp: Experience) -> None: - self._parse_complete_key(key) # validate key format - records = self._records.get(key) - if records is None: - raise KeyError(f"Key '{key}' does not exist.") - if old_sample_id not in records: - raise KeyError(f"sample_id '{old_sample_id}' does not exist under key '{key}'.") - - new_sample_id = get_sample_id(exp) - owner_key = self._sample_to_key.get(new_sample_id) - if owner_key is not None and (owner_key != key or new_sample_id != old_sample_id): - raise ValueError( - f"Duplicate sample_id '{new_sample_id}' already exists under key '{owner_key}'." - ) - - items = [] - for sample_id, record in records.items(): - if sample_id == old_sample_id: - items.append((new_sample_id, exp)) - else: - items.append((sample_id, record)) - - records.clear() - records.update(items) - self._sample_to_key.pop(old_sample_id, None) - self._sample_to_key[new_sample_id] = key + with self._lock: + self._parse_complete_key(key) # validate key format + records = self._records.get(key) + if records is None: + raise KeyError(f"Key '{key}' does not exist.") + if old_sample_id not in records: + raise KeyError(f"sample_id '{old_sample_id}' does not exist under key '{key}'.") + + new_sample_id = get_sample_id(exp) + owner_key = self._sample_to_key.get(new_sample_id) + if owner_key is not None and (owner_key != key or new_sample_id != old_sample_id): + raise ValueError( + f"Duplicate sample_id '{new_sample_id}' already exists under key '{owner_key}'." + ) + + items = [] + for sample_id, record in records.items(): + if sample_id == old_sample_id: + items.append((new_sample_id, exp)) + else: + items.append((sample_id, record)) + + records.clear() + records.update(items) + self._sample_to_key.pop(old_sample_id, None) + self._sample_to_key[new_sample_id] = key def update( self, @@ -130,49 +144,55 @@ def update( update: ExperienceUpdate, sample_ids: List[str] | None, ) -> None: - batch, task, run = self._parse_complete_key(key) # validate key format - records = self._records.get(key) - if records is None: - raise KeyError(f"Key '{key}' does not exist.") - target_ids: Iterable[str] = list(records.keys()) if sample_ids is None else sample_ids - for sample_id in target_ids: - if sample_id not in records: - raise KeyError(f"sample_id '{sample_id}' does not exist under key '{key}'.") - exp = records[sample_id] - exp.eid.batch = batch - exp.eid.task = task - exp.eid.run = run - if update.reward is not None: - exp.reward = update.reward - if update.info: - if exp.info is None: - exp.info = {} - exp.info.update(update.info) - if update.teacher_logprobs is not None: - exp.teacher_logprobs = update.teacher_logprobs + with self._lock: + batch, task, run = self._parse_complete_key(key) # validate key format + records = self._records.get(key) + if records is None: + raise KeyError(f"Key '{key}' does not exist.") + target_ids: Iterable[str] = list(records.keys()) if sample_ids is None else sample_ids + for sample_id in target_ids: + if sample_id not in records: + raise KeyError(f"sample_id '{sample_id}' does not exist under key '{key}'.") + exp = records[sample_id] + exp.eid.batch = batch + exp.eid.task = task + exp.eid.run = run + if update.reward is not None: + exp.reward = update.reward + if update.info: + if exp.info is None: + exp.info = {} + exp.info.update(update.info) + if update.teacher_logprobs is not None: + exp.teacher_logprobs = update.teacher_logprobs def get(self, key: str) -> List[Experience]: - result: List[Experience] = [] - for matched_key in self._matching_keys(key): - result.extend(self._records[matched_key].values()) - return result + with self._lock: + result: List[Experience] = [] + for matched_key in self._matching_keys(key): + result.extend(self._records[matched_key].values()) + return result def remove(self, key: str) -> List[Experience]: - result: List[Experience] = [] - for matched_key in self._matching_keys(key): - result.extend(self._drop_key(matched_key)) - return result + with self._lock: + result: List[Experience] = [] + for matched_key in self._matching_keys(key): + result.extend(self._drop_key(matched_key)) + return result def keys(self) -> list[str]: - return list(self._records.keys()) + with self._lock: + return list(self._records.keys()) def block_prefix(self, prefix: str) -> None: """Mark a batch prefix as blocked; future ``add``/``overwrite`` are dropped.""" - self._blocked_batches.add(prefix) + with self._lock: + self._blocked_batches.add(prefix) def is_prefix_blocked(self, prefix: str) -> bool: """Return whether the given batch prefix is blocked.""" - return prefix in self._blocked_batches + with self._lock: + return prefix in self._blocked_batches @staticmethod def _parse_complete_key(key: str) -> tuple[str, str, int]: diff --git a/trinity/common/models/recording/recorder.py b/trinity/common/models/recording/recorder.py index aee6366c195..be64f6d9792 100644 --- a/trinity/common/models/recording/recorder.py +++ b/trinity/common/models/recording/recorder.py @@ -2,6 +2,8 @@ import asyncio import logging +import queue +import threading from collections.abc import Callable, Sequence from datetime import datetime, timezone from typing import Any, Optional @@ -17,12 +19,22 @@ BuildExperiencesFn = Callable[..., Sequence[Experience]] +# Sentinel pushed onto the queue to request the flusher thread to exit. +_STOP = object() + class Recorder: - """Drains finished turns into a ``RecordStore`` from a background task. + """Drains finished turns into a ``RecordStore`` from a background thread. Engine-specific code supplies ``build_experiences``, which converts a finished engine output object into Trinity ``Experience`` instances. + + The CPU-bound store writes (and prefix merging) run on a dedicated daemon + thread so they never stall the engine's event loop. Producers (engine-loop + coroutines) push onto an unbounded, thread-safe ``queue.Queue`` and never + block. ``MemoryStore`` is expected to be thread-safe (it guards its + dictionaries with a lock) since engine-loop code may touch the store + concurrently with the flusher. """ def __init__( @@ -41,23 +53,31 @@ def __init__( self.engine_client = engine_client self.merge_prefix_experiences = merge_prefix_experiences self._build_experiences = build_experiences - self._queue: "asyncio.Queue[Optional[Experience]]" = asyncio.Queue() - self._flusher: Optional[asyncio.Task] = None + # Thread-safe unbounded queue; producers never block on put. + self._queue: "queue.Queue[Optional[Experience]]" = queue.Queue() + self._flusher: Optional[threading.Thread] = None self._pending: "set[asyncio.Task]" = set() self._prefix_merger = PrefixExperienceMerger(store) def start(self) -> None: - """Start the background flusher. Idempotent.""" + """Start the background flusher thread. Idempotent.""" if self._flusher is not None or not self.enabled: return - self._flusher = asyncio.create_task(self._flush_loop()) + self._flusher = threading.Thread( + target=self._flush_loop, + name=f"trinity-recorder-flusher-rank{self.rank}", + daemon=True, + ) + self._flusher.start() async def stop(self) -> None: - """Drain in-flight + queued turns, then stop the flusher.""" + """Drain in-flight + queued turns, then stop the flusher thread.""" if self._flusher is None: return await self.flush() - self._flusher.cancel() + self._queue.put(_STOP) # type: ignore [arg-type] + # join() blocks; run it off the event loop. + await asyncio.to_thread(self._flusher.join) self._flusher = None def schedule_record(self, output: Any, record_key: Optional[str], **builder_kwargs) -> None: @@ -71,7 +91,9 @@ async def flush(self) -> None: if self._pending: await asyncio.gather(*self._pending, return_exceptions=True) if self._flusher is not None: - await self._queue.join() + # queue.Queue.join() is a blocking call; run it off the event loop + # so the engine keeps scheduling while the flusher drains. + await asyncio.to_thread(self._queue.join) async def _record(self, output: Any, record_key: Optional[str], **builder_kwargs) -> None: if skip_recording_ctx.get(): @@ -86,19 +108,24 @@ async def _record(self, output: Any, record_key: Optional[str], **builder_kwargs **builder_kwargs, ) for exp in exps: - await self._queue.put(exp) + self._queue.put(exp) # non-blocking on an unbounded queue - async def _flush_loop(self) -> None: + def _flush_loop(self) -> None: while True: - exp = await self._queue.get() + exp = self._queue.get() try: - if exp is None: + if exp is _STOP: return - await self._safe_append(exp) + self._safe_append(exp) # type: ignore [arg-type] + except Exception: + logging.getLogger(__name__).exception( + "recording flusher failed for request %s", + exp.eid.suffix if exp is not _STOP else None, + ) finally: self._queue.task_done() - async def _safe_append(self, exp: Experience) -> None: + def _safe_append(self, exp: Experience) -> None: try: record_key = get_record_key(exp) if self.merge_prefix_experiences and self._prefix_merger.try_merge(record_key, exp): From 722bf72bfc7f79cdace8b6838a3c07e81925899f Mon Sep 17 00:00:00 2001 From: "panxuchen.pxc" Date: Thu, 2 Jul 2026 17:22:17 +0800 Subject: [PATCH 83/84] add harbor tests --- tests/explorer/harbor_workflow_test.py | 257 ++++++++++++++++++ .../arithmetic-task/environment/.gitkeep | 1 + .../harbor/arithmetic-task/instruction.md | 3 + .../data/harbor/arithmetic-task/task.toml | 28 ++ .../data/harbor/arithmetic-task/tests/test.sh | 2 + 5 files changed, 291 insertions(+) create mode 100644 tests/explorer/harbor_workflow_test.py create mode 100644 tests/template/data/harbor/arithmetic-task/environment/.gitkeep create mode 100644 tests/template/data/harbor/arithmetic-task/instruction.md create mode 100644 tests/template/data/harbor/arithmetic-task/task.toml create mode 100644 tests/template/data/harbor/arithmetic-task/tests/test.sh diff --git a/tests/explorer/harbor_workflow_test.py b/tests/explorer/harbor_workflow_test.py new file mode 100644 index 00000000000..6ca212e6356 --- /dev/null +++ b/tests/explorer/harbor_workflow_test.py @@ -0,0 +1,257 @@ +# -*- coding: utf-8 -*- +"""Tests for HarborWorkflow without starting Docker or another sandbox.""" + +from __future__ import annotations + +import os +import re +import sys +import tempfile +import unittest +from pathlib import Path +from typing import Any + +import openai +import ray + +REPO_ROOT = Path(__file__).resolve().parents[2] +HARBOR_SRC = REPO_ROOT / "thirdparties" / "harbor" / "src" +if str(HARBOR_SRC) not in sys.path: + sys.path.insert(0, str(HARBOR_SRC)) + +from harbor.agents.base import BaseAgent +from harbor.environments.base import BaseEnvironment, ExecResult +from harbor.environments.capabilities import EnvironmentCapabilities +from harbor.models.agent.context import AgentContext +from harbor.models.verifier.result import VerifierResult +from harbor.verifier.base import BaseVerifier + +from tests.tools import get_model_path, get_template_config +from trinity.common.constants import MODEL_PATH_ENV_VAR +from trinity.common.models.allocator import Allocator +from trinity.common.workflows.harbor_workflow import HarborWorkflow +from trinity.common.workflows.workflow import Task + + +class InProcessHarborEnvironment(BaseEnvironment): + """A minimal Harbor environment for unit tests. + + It satisfies Harbor's trial lifecycle but never starts Docker, a sandbox, + or a subprocess. The arithmetic agent either calls a real OpenAI-compatible + vLLM service or uses a deterministic fallback response, so the environment + only needs to accept lifecycle calls. + """ + + @staticmethod + def type() -> str: + return "in-process" + + @property + def capabilities(self) -> EnvironmentCapabilities: + return EnvironmentCapabilities(mounted=True) + + def _validate_definition(self) -> None: + pass + + async def start(self, force_build: bool) -> None: + pass + + async def stop(self, delete: bool) -> None: + pass + + async def upload_file(self, source_path: Path | str, target_path: str) -> None: + pass + + async def upload_dir(self, source_dir: Path | str, target_dir: str) -> None: + pass + + async def download_file(self, source_path: str, target_path: Path | str) -> None: + Path(target_path).parent.mkdir(parents=True, exist_ok=True) + + async def download_dir(self, source_dir: str, target_dir: Path | str) -> None: + Path(target_dir).mkdir(parents=True, exist_ok=True) + + async def exec( + self, + command: str, + cwd: str | None = None, + env: dict[str, str] | None = None, + timeout_sec: int | None = None, + user: str | int | None = None, + ) -> ExecResult: + return ExecResult(stdout="", stderr="", return_code=0) + + +class OpenAIArithmeticAgent(BaseAgent): + """A Harbor agent that solves the task via the injected OpenAI endpoint.""" + + @staticmethod + def name() -> str: + return "openai-arithmetic-agent" + + def version(self) -> str: + return "1.0.0" + + async def setup(self, environment: BaseEnvironment) -> None: + pass + + async def run( + self, + instruction: str, + environment: BaseEnvironment, + context: AgentContext, + ) -> None: + if self.extra_env.get("TRINITY_FAKE_OPENAI_RESPONSE") is not None: + answer = self.extra_env["TRINITY_FAKE_OPENAI_RESPONSE"] + else: + client = openai.AsyncOpenAI( + base_url=self.extra_env["OPENAI_BASE_URL"], + api_key=self.extra_env["OPENAI_API_KEY"], + ) + response = await client.chat.completions.create( + model=self.extra_env["OPENAI_MODEL"], + messages=[ + { + "role": "system", + "content": "Solve the problem. Return only the final answer.", + }, + {"role": "user", "content": instruction}, + ], + temperature=0.0, + max_tokens=16, + ) + answer = response.choices[0].message.content or "" + self.logs_dir.mkdir(parents=True, exist_ok=True) + (self.logs_dir / "answer.txt").write_text(answer.strip()) + + +class ExactAnswerHarborVerifier(BaseVerifier): + """Rule-based verifier for answer-only arithmetic tasks.""" + + def __init__(self, *args: Any, expected_answer: str, **kwargs: Any) -> None: + super().__init__(*args, **kwargs) + self.expected_answer = expected_answer + + async def verify(self) -> VerifierResult: + answer_path = self.trial_paths.agent_dir / "answer.txt" + answer = answer_path.read_text().strip() if answer_path.exists() else "" + integers = re.findall(r"-?\d+", answer) + final_answer = integers[-1] if integers else "" + reward = 1.0 if final_answer == self.expected_answer else 0.0 + self.trial_paths.reward_json_path.parent.mkdir(parents=True, exist_ok=True) + self.trial_paths.reward_json_path.write_text( + f'{{"reward": {reward}, "exact_match": {reward}}}' + ) + return VerifierResult(rewards={"reward": reward, "exact_match": reward}) + + +class DummyRecordingModelWrapper: + def __init__( + self, + *, + base_url: str = "http://127.0.0.1:8000/v1", + model_name: str = "dummy-rollout-model", + ) -> None: + self.base_url = base_url + self.model_name = model_name + self.api_key = "EMPTY" + self.reward_updates: list[dict[str, Any]] = [] + + def set_api_key(self, api_key: str) -> None: + self.api_key = api_key + + async def update_experience_reward_async( + self, + *, + key: str, + reward: float, + info: dict | None = None, + sample_ids: list[str] | None = None, + ) -> None: + self.reward_updates.append( + { + "key": key, + "reward": reward, + "info": info, + "sample_ids": sample_ids, + } + ) + + +class HarborWorkflowTest(unittest.IsolatedAsyncioTestCase): + async def _create_model_wrapper(self): + if not os.environ.get(MODEL_PATH_ENV_VAR): + return DummyRecordingModelWrapper(), False, [] + + ray.init(ignore_reinit_error=True, namespace="trinity_unittest") + config = get_template_config() + config.mode = "explore" + config.model.model_path = get_model_path() + config.explorer.rollout_model.engine_num = 1 + config.explorer.rollout_model.enable_history = True + config.explorer.rollout_model.enable_openai_api = True + config.check_and_update() + + allocator = Allocator(config.explorer) + rollout_models, _ = await allocator.create_all_models() + return rollout_models[0], True, rollout_models + + async def test_harbor_workflow_runs_arithmetic_task(self) -> None: + task_dir = REPO_ROOT / "tests" / "template" / "data" / "harbor" / "arithmetic-task" + model, uses_real_vllm, rollout_models = await self._create_model_wrapper() + expected_reward = 1.0 if uses_real_vllm else 0.0 + + with tempfile.TemporaryDirectory() as tmp_dir: + task = Task( + workflow=HarborWorkflow, + raw_task={"task_dir": str(task_dir)}, + workflow_args={ + "agent_import_path": ( + "tests.explorer.harbor_workflow_test:" + "OpenAIArithmeticAgent" + ), + "environment": { + "import_path": ( + "tests.explorer.harbor_workflow_test:" + "InProcessHarborEnvironment" + ), + }, + "verifier": { + "import_path": ( + "tests.explorer.harbor_workflow_test:" + "ExactAnswerHarborVerifier" + ), + "kwargs": {"expected_answer": "42"}, + }, + "provider_env": ( + {} + if uses_real_vllm + else {"TRINITY_FAKE_OPENAI_RESPONSE": "not-a-number"} + ), + "trial_name": "harbor-workflow-vllm-arithmetic", + "trials_dir": tmp_dir, + }, + batch_id="harbor", + task_id=0, + run_id=0, + ) + + try: + metrics = await HarborWorkflow(task=task, model=model).run_async() + finally: + for wrapper in rollout_models: + await wrapper.shutdown() + if uses_real_vllm: + ray.shutdown() + + self.assertEqual(metrics["harbor/reward"], expected_reward) + self.assertEqual(metrics["harbor/has_exception"], 0.0) + self.assertEqual(metrics["harbor/num_rewards"], 2.0) + if isinstance(model, DummyRecordingModelWrapper): + self.assertEqual(len(model.reward_updates), 1) + self.assertEqual(model.reward_updates[0]["key"], "harbor/0/0") + self.assertEqual(model.reward_updates[0]["reward"], expected_reward) + self.assertEqual( + model.reward_updates[0]["info"]["harbor_rewards"], + {"reward": expected_reward, "exact_match": expected_reward}, + ) diff --git a/tests/template/data/harbor/arithmetic-task/environment/.gitkeep b/tests/template/data/harbor/arithmetic-task/environment/.gitkeep new file mode 100644 index 00000000000..8b137891791 --- /dev/null +++ b/tests/template/data/harbor/arithmetic-task/environment/.gitkeep @@ -0,0 +1 @@ + diff --git a/tests/template/data/harbor/arithmetic-task/instruction.md b/tests/template/data/harbor/arithmetic-task/instruction.md new file mode 100644 index 00000000000..47977b5319c --- /dev/null +++ b/tests/template/data/harbor/arithmetic-task/instruction.md @@ -0,0 +1,3 @@ +What is 19 + 23? + +Return only the decimal integer. diff --git a/tests/template/data/harbor/arithmetic-task/task.toml b/tests/template/data/harbor/arithmetic-task/task.toml new file mode 100644 index 00000000000..c816f05cb2e --- /dev/null +++ b/tests/template/data/harbor/arithmetic-task/task.toml @@ -0,0 +1,28 @@ +version = "1.0" + +[task] +name = "trinity/arithmetic-harbor-workflow" +authors = [] +keywords = ["test", "arithmetic"] + +[metadata] +difficulty = "easy" +category = "math" +tags = ["unit", "vllm"] + +[agent] +timeout_sec = 30.0 + +[verifier] +timeout_sec = 10.0 + +[environment] +cpus = 1 +memory_mb = 128 +storage_mb = 128 +gpus = 0 +mcp_servers = [] + +[verifier.env] + +[solution.env] diff --git a/tests/template/data/harbor/arithmetic-task/tests/test.sh b/tests/template/data/harbor/arithmetic-task/tests/test.sh new file mode 100644 index 00000000000..321c13d4241 --- /dev/null +++ b/tests/template/data/harbor/arithmetic-task/tests/test.sh @@ -0,0 +1,2 @@ +#!/bin/sh +echo "This script is present only so Harbor validates the task." From 64b001d10259cec6d2782ef76082b9c02d4746de Mon Sep 17 00:00:00 2001 From: "panxuchen.pxc" Date: Thu, 2 Jul 2026 17:30:32 +0800 Subject: [PATCH 84/84] fix pre-commit --- tests/explorer/harbor_workflow_test.py | 22 +++++-------------- .../arithmetic-task/environment/.gitkeep | 1 - trinity/common/workflows/workflow.py | 6 ++--- trinity/explorer/workflow_runner.py | 12 +++++----- 4 files changed, 15 insertions(+), 26 deletions(-) diff --git a/tests/explorer/harbor_workflow_test.py b/tests/explorer/harbor_workflow_test.py index 6ca212e6356..04613b6002c 100644 --- a/tests/explorer/harbor_workflow_test.py +++ b/tests/explorer/harbor_workflow_test.py @@ -5,7 +5,6 @@ import os import re -import sys import tempfile import unittest from pathlib import Path @@ -13,12 +12,6 @@ import openai import ray - -REPO_ROOT = Path(__file__).resolve().parents[2] -HARBOR_SRC = REPO_ROOT / "thirdparties" / "harbor" / "src" -if str(HARBOR_SRC) not in sys.path: - sys.path.insert(0, str(HARBOR_SRC)) - from harbor.agents.base import BaseAgent from harbor.environments.base import BaseEnvironment, ExecResult from harbor.environments.capabilities import EnvironmentCapabilities @@ -32,6 +25,8 @@ from trinity.common.workflows.harbor_workflow import HarborWorkflow from trinity.common.workflows.workflow import Task +REPO_ROOT = Path(__file__).resolve().parents[2] + class InProcessHarborEnvironment(BaseEnvironment): """A minimal Harbor environment for unit tests. @@ -207,26 +202,21 @@ async def test_harbor_workflow_runs_arithmetic_task(self) -> None: raw_task={"task_dir": str(task_dir)}, workflow_args={ "agent_import_path": ( - "tests.explorer.harbor_workflow_test:" - "OpenAIArithmeticAgent" + "tests.explorer.harbor_workflow_test:" "OpenAIArithmeticAgent" ), "environment": { "import_path": ( - "tests.explorer.harbor_workflow_test:" - "InProcessHarborEnvironment" + "tests.explorer.harbor_workflow_test:" "InProcessHarborEnvironment" ), }, "verifier": { "import_path": ( - "tests.explorer.harbor_workflow_test:" - "ExactAnswerHarborVerifier" + "tests.explorer.harbor_workflow_test:" "ExactAnswerHarborVerifier" ), "kwargs": {"expected_answer": "42"}, }, "provider_env": ( - {} - if uses_real_vllm - else {"TRINITY_FAKE_OPENAI_RESPONSE": "not-a-number"} + {} if uses_real_vllm else {"TRINITY_FAKE_OPENAI_RESPONSE": "not-a-number"} ), "trial_name": "harbor-workflow-vllm-arithmetic", "trials_dir": tmp_dir, diff --git a/tests/template/data/harbor/arithmetic-task/environment/.gitkeep b/tests/template/data/harbor/arithmetic-task/environment/.gitkeep index 8b137891791..e69de29bb2d 100644 --- a/tests/template/data/harbor/arithmetic-task/environment/.gitkeep +++ b/tests/template/data/harbor/arithmetic-task/environment/.gitkeep @@ -1 +0,0 @@ - diff --git a/trinity/common/workflows/workflow.py b/trinity/common/workflows/workflow.py index 68764309197..75a60d454f7 100644 --- a/trinity/common/workflows/workflow.py +++ b/trinity/common/workflows/workflow.py @@ -40,7 +40,7 @@ def ok(self) -> bool: class Task(dict): """A Task class that defines a task and its associated reward function / workflow.""" - workflow: Type[Workflow] = None + workflow: Type[WorkflowBase] = None repeat_times: Optional[int] = None format_args: FormatConfig = field(default_factory=FormatConfig) rollout_args: GenerationConfig = field(default_factory=GenerationConfig) @@ -61,7 +61,7 @@ def to_workflow( self, model: ModelWrapper, auxiliary_models: Optional[List[ModelWrapper]] = None, - ) -> Workflow: + ) -> "WorkflowBase": """Convert the task to a workflow. Args: @@ -104,7 +104,7 @@ def to_dict(self) -> dict: class WorkflowBase: """The base workflow interface.""" - def __init__(self, task: Task, model: ModelWrapper) -> None: + def __init__(self, task: Task, model: ModelWrapper, **kwargs) -> None: self.task = task self.model = model self.model.set_api_key(task.api_key) # set the API key for the rollout model diff --git a/trinity/explorer/workflow_runner.py b/trinity/explorer/workflow_runner.py index 051f6501397..59b1be88eae 100644 --- a/trinity/explorer/workflow_runner.py +++ b/trinity/explorer/workflow_runner.py @@ -14,7 +14,7 @@ from trinity.common.experience import Experience from trinity.common.models.allocator import Allocator from trinity.common.models.model import ModelWrapper -from trinity.common.workflows import Status, Task, Workflow +from trinity.common.workflows import Status, Task, WorkflowBase from trinity.utils.log import get_logger @@ -46,7 +46,7 @@ def __init__( ) for index, auxiliary_model_id in enumerate(auxiliary_model_ids or []) ] - self.workflow_instance: Workflow = None + self.workflow_instance: WorkflowBase = None self.rollout_model_id = rollout_model_id self.runner_id = runner_id self.runner_state = { @@ -84,7 +84,7 @@ async def prepare(self) -> None: def is_alive(self): return True - def _create_workflow_instance(self, task: Task) -> Workflow: + def _create_workflow_instance(self, task: Task) -> WorkflowBase: if task.workflow is None: raise ValueError("Workflow is not set in the task.") if ( @@ -104,7 +104,7 @@ def _create_workflow_instance(self, task: Task) -> Workflow: self.workflow_instance.set_single_run_context(task.run_id) return self.workflow_instance - async def _run_workflow(self, workflow_instance: Workflow) -> Status: + async def _run_workflow(self, workflow_instance: WorkflowBase) -> Status: status = await workflow_instance.execute() if not isinstance(status, Status): raise TypeError( @@ -113,7 +113,7 @@ async def _run_workflow(self, workflow_instance: Workflow) -> Status: ) return status - def _create_isolated_workflow_instance(self, task: Task, run_id: int) -> Workflow: + def _create_isolated_workflow_instance(self, task: Task, run_id: int) -> WorkflowBase: model_wrapper = self.model_wrapper.clone_with_isolated_state() # only a shallow copy is enough; use copy.copy so the result stays a Task # (Task inherits dict, so task.copy() would return a plain dict) @@ -238,7 +238,7 @@ async def run_single(i: int) -> Status: async def _execute_single_run( self, - workflow: Workflow, + workflow: WorkflowBase, ) -> Status: st = time.time() self.runner_state["terminate_time"] = None