From 16f507bb845eb74eca9d30036e2f839850e8f6fa Mon Sep 17 00:00:00 2001 From: Fred Wei <20514172+WeiHaocheng@users.noreply.github.com> Date: Sat, 7 Feb 2026 03:49:57 +0000 Subject: [PATCH 01/18] draft init scaffolding init scaffolding run gsm8k fix workflow no skip tokenizer fix reward multi-thread and oom --- areal/experimental/scaffolding/__init__.py | 49 ++ areal/experimental/scaffolding/_compat.py | 192 +++++ areal/experimental/scaffolding/controllers.py | 680 ++++++++++++++++++ areal/experimental/scaffolding/task.py | 232 ++++++ areal/experimental/scaffolding/worker.py | 231 ++++++ areal/experimental/scaffolding/workflow.py | 263 +++++++ areal/reward/__init__.py | 53 +- examples/scaffolding/README.md | 225 ++++++ .../scaffolding/gsm8k_rlvr_scaffolding.py | 61 ++ .../scaffolding/gsm8k_rlvr_scaffolding.yaml | 177 +++++ 10 files changed, 2146 insertions(+), 17 deletions(-) create mode 100644 areal/experimental/scaffolding/__init__.py create mode 100644 areal/experimental/scaffolding/_compat.py create mode 100644 areal/experimental/scaffolding/controllers.py create mode 100644 areal/experimental/scaffolding/task.py create mode 100644 areal/experimental/scaffolding/worker.py create mode 100644 areal/experimental/scaffolding/workflow.py create mode 100644 examples/scaffolding/README.md create mode 100644 examples/scaffolding/gsm8k_rlvr_scaffolding.py create mode 100644 examples/scaffolding/gsm8k_rlvr_scaffolding.yaml diff --git a/areal/experimental/scaffolding/__init__.py b/areal/experimental/scaffolding/__init__.py new file mode 100644 index 0000000000..24cda41f7e --- /dev/null +++ b/areal/experimental/scaffolding/__init__.py @@ -0,0 +1,49 @@ +""" +Scaffolding Framework Integration for AReaL. + +This module provides integration between TensorRT-LLM's Scaffolding framework +and AReaL's RL training pipeline. + +Key Components: +- ScaffoldingWorkflow: RolloutWorkflow implementation that wraps ScaffoldingLlm +- RLVRRewardTask: Task for computing verifiable rewards +- RLVRRewardController: Controller for computing verifiable rewards +- PipelineTrajectoryMaker: Controller for composing generation and reward pipelines +- ChatTracer: TaskCollection for tracing multi-turn chat conversations +- TraceTrajectoryMaker: Controller that traces ChatTask objects during rollout +- TraceGenerationTask: Task for tracing multi-turn generation +- ChatRewardTask: Task for computing rewards on traced interactions +- CreateWorkerFromEngine: Creates a scaffolding Worker from AReaL's InferenceEngine +- SGLangWorker: Worker implementation for SGLang engines + +Note: Requires tensorrt_llm to be installed for full functionality. +""" + +from areal.experimental.scaffolding._compat import HAS_TENSORRT_LLM +from areal.experimental.scaffolding.controllers import ( + ChatTracer, + PipelineTrajectoryMaker, + RLVRRewardController, + TraceTrajectoryMaker, +) +from areal.experimental.scaffolding.task import ( + ChatRewardTask, + RLVRRewardTask, + TraceGenerationTask, +) +from areal.experimental.scaffolding.worker import CreateWorkerFromEngine, SGLangWorker +from areal.experimental.scaffolding.workflow import ScaffoldingWorkflow + +__all__ = [ + "HAS_TENSORRT_LLM", + "ScaffoldingWorkflow", + "RLVRRewardTask", + "RLVRRewardController", + "PipelineTrajectoryMaker", + "ChatTracer", + "TraceTrajectoryMaker", + "TraceGenerationTask", + "ChatRewardTask", + "CreateWorkerFromEngine", + "SGLangWorker", +] diff --git a/areal/experimental/scaffolding/_compat.py b/areal/experimental/scaffolding/_compat.py new file mode 100644 index 0000000000..3b845056ee --- /dev/null +++ b/areal/experimental/scaffolding/_compat.py @@ -0,0 +1,192 @@ +"""Compatibility layer for optional tensorrt_llm.scaffolding dependency. + +Provides imports from tensorrt_llm.scaffolding when available, or lightweight +standalone implementations when not installed. +""" + +from __future__ import annotations + +import enum +from dataclasses import dataclass, field +from typing import Any + +try: + from tensorrt_llm.scaffolding import ( + NativeGenerationController, + ScaffoldingLlm, + ) + from tensorrt_llm.scaffolding.controller import Controller + from tensorrt_llm.scaffolding.result import ScaffoldingOutput + from tensorrt_llm.scaffolding.task import ( + AssistantMessage, + ChatTask, + GenerationTask, + Task, + TaskStatus, + ) + from tensorrt_llm.scaffolding.task_collection import ( + TaskCollection, + with_task_collection, + ) + from tensorrt_llm.scaffolding.worker import OpenaiWorker, Worker + + HAS_TENSORRT_LLM = True + +except ImportError: + HAS_TENSORRT_LLM = False + + # ---- Standalone lightweight implementations ---- + # These provide the scaffolding interfaces so the framework works + # without tensorrt_llm installed. + + class Controller: + """Lightweight Controller base class.""" + + def process(self, tasks: list, **kwargs) -> Any: + yield tasks + + @dataclass + class Task: + """Lightweight Task base class.""" + + worker_tag: Any = None + + @dataclass + class GenerationTask(Task): + """Lightweight GenerationTask.""" + + input_str: str | None = None + output_str: str | None = None + input_tokens: list | None = None + output_tokens: list | None = None + logprobs: Any = None + finish_reason: str | None = None + perf_metrics: Any = None + customized_result_fields: dict = field(default_factory=dict) + + @dataclass + class ChatTask(Task): + """Lightweight ChatTask.""" + + messages: list = field(default_factory=list) + completion: Any = None + tools: list | None = None + finish_reason: str | None = None + input_tokens: list | None = None + output_tokens: list | None = None + enable_token_counting: bool = False + prompt_tokens_num: int = 0 + completion_tokens_num: int = 0 + reasoning_tokens_num: int = 0 + perf_metrics: Any = None + + @staticmethod + def create_from_prompt(prompt: str) -> ChatTask: + return ChatTask(messages=[{"role": "user", "content": prompt}]) + + def messages_to_dict_content(self) -> list: + return self.messages + + class TaskStatus(enum.Enum): + """Lightweight TaskStatus.""" + + SUCCESS = "success" + WORKER_EXECEPTION = "worker_exception" # noqa: S105 (matches upstream typo) + + class AssistantMessage: + """Lightweight AssistantMessage.""" + + def __init__( + self, + content: str | None = None, + reasoning: str | None = None, + reasoning_content: str | None = None, + tool_calls: list | None = None, + ): + self.content = content + self.reasoning = reasoning + self.reasoning_content = reasoning_content + self.tool_calls = tool_calls + + class TaskCollection: + """Lightweight TaskCollection base class.""" + + def before_yield(self, tasks: list) -> None: + pass + + def after_yield(self, tasks: list) -> None: + pass + + def with_task_collection(name: str, collection_cls: type): + """Decorator that attaches a TaskCollection to a Controller class.""" + + def decorator(cls): + if not hasattr(cls, "task_collections"): + cls.task_collections = {} + cls.task_collections[name] = collection_cls() + return cls + + return decorator + + class Worker: + """Lightweight Worker base class.""" + + class OpenaiWorker(Worker): + """Lightweight OpenaiWorker base class.""" + + def __init__(self, async_client: Any = None, model: str = "", **kwargs): + self.async_client = async_client + self.model = model + + def convert_task_params(self, task: Any) -> dict: + return {} + + @dataclass + class ScaffoldingOutput: + """Lightweight ScaffoldingOutput.""" + + text: str = "" + token_ids: list = field(default_factory=list) + + class NativeGenerationController(Controller): + """Lightweight NativeGenerationController.""" + + class WorkerTag(enum.Enum): + GENERATION = "generation" + + def process(self, tasks: list, **kwargs) -> Any: + yield tasks + + class ScaffoldingLlm: + """Lightweight ScaffoldingLlm.""" + + def __init__(self, controller: Controller, workers: dict | None = None): + self.controller = controller + self.workers = workers or {} + + def generate(self, prompt: str, **kwargs) -> Any: + return None + + async def generate_async(self, prompt: str, **kwargs) -> Any: + return None + + def shutdown(self) -> None: + pass + + +__all__ = [ + "HAS_TENSORRT_LLM", + "AssistantMessage", + "ChatTask", + "Controller", + "GenerationTask", + "NativeGenerationController", + "OpenaiWorker", + "ScaffoldingLlm", + "ScaffoldingOutput", + "Task", + "TaskCollection", + "TaskStatus", + "Worker", + "with_task_collection", +] diff --git a/areal/experimental/scaffolding/controllers.py b/areal/experimental/scaffolding/controllers.py new file mode 100644 index 0000000000..d191317587 --- /dev/null +++ b/areal/experimental/scaffolding/controllers.py @@ -0,0 +1,680 @@ +""" +RLVR Controllers for Scaffolding Framework. + +This module provides controllers for RLVR (Reinforcement Learning with Verifiable +Rewards) that integrate with TensorRT-LLM's scaffolding framework. + +Key Components: +- RLVRRewardController: Controller that processes reward computation +- PipelineTrajectoryMaker: Controller that composes generation and reward pipelines +- ChatTracer: TaskCollection for tracing multi-turn chat conversations +- TraceTrajectoryMaker: Controller that traces ChatTask objects during rollout +""" + +from __future__ import annotations + +from collections.abc import Callable +from enum import Enum +from typing import TYPE_CHECKING, Any + +from areal.api.reward_api import AsyncRewardWrapper +from areal.experimental.openai.cache import InteractionCache +from areal.experimental.openai.types import InteractionWithTokenLogpReward +from areal.experimental.scaffolding._compat import ( + ChatTask, + Controller, + GenerationTask, + Task, + TaskCollection, + with_task_collection, +) +from areal.experimental.scaffolding.task import ( + ChatRewardTask, + RLVRRewardTask, + TraceGenerationTask, +) +from areal.utils import logging + +if TYPE_CHECKING: + pass + +logger = logging.getLogger("RLVRControllers") + + +class RLVRRewardController(Controller): + """Controller for computing RLVR (verifiable) rewards. + + This controller processes RLVRRewardTask objects and computes rewards + using a provided reward function. The reward function should verify + whether the generated answer is correct. + + The reward computation follows the pattern from RLVRWorkflow._compute_rewards: + 1. Decode output tokens to string (if needed) + 2. Call reward_fn(prompt_str, completion_str, input_tokens, output_tokens, **task_data) + 3. Store the reward in the task and update the interaction object + + Parameters + ---------- + reward_fn : Callable + The reward function that takes (prompt, completions, prompt_ids, completion_ids, **data) + and returns a reward value (typically 0 or 1 for verifiable rewards). + + Example + ------- + ```python + from areal.reward.gsm8k import gsm8k_reward_fn + + reward_controller = RLVRRewardController(gsm8k_reward_fn) + ``` + """ + + class WorkerTag(Enum): + """Worker tag for reward computation.""" + + REWARD = "rlvr_reward" + + def __init__(self, reward_fn: Callable[..., Any]): + """Initialize the RLVR reward controller. + + Parameters + ---------- + reward_fn : Callable + The reward function for verifying answers. + """ + super().__init__() + self.reward_fn = reward_fn + self.async_reward_fn = AsyncRewardWrapper(reward_fn) + self.scores: list[float] | None = None + + def process(self, tasks: list[Task], **kwargs) -> Any: + """Process reward tasks and compute rewards. + + This method computes rewards for each task using the reward function. + The rewards are stored in: + 1. task.reward - the computed reward value + 2. task.interaction.reward - if an interaction object is provided + 3. self.scores - list of all computed rewards + + Parameters + ---------- + tasks : list[Task] + List of RLVRRewardTask objects to process. + **kwargs + Additional keyword arguments. + + Yields + ------ + list[Task] + The processed tasks with rewards computed. + """ + # Mark tasks with worker tag (for potential worker-based execution) + for task in tasks: + task.worker_tag = self.WorkerTag.REWARD + + # Compute rewards synchronously + # Note: For async execution, this would be handled by a worker + self.scores = [] + for task in tasks: + if isinstance(task, RLVRRewardTask): + reward = self._compute_reward(task) + task.reward = reward + self.scores.append(reward) + + # Update the interaction object if provided + if task.interaction is not None: + task.interaction.reward = reward + elif isinstance(task, GenerationTask): + # For generation tasks, compute reward from customized fields + reward = self._compute_reward_from_generation_task(task, **kwargs) + task.customized_result_fields["reward"] = reward + self.scores.append(reward) + + yield tasks + + def _compute_reward(self, task: RLVRRewardTask) -> float: + """Compute reward for an RLVR reward task. + + Parameters + ---------- + task : RLVRRewardTask + The reward task containing prompt, completion, and task data. + + Returns + ------- + float + The computed reward value. + """ + reward = self.reward_fn( + task.prompt_str, + task.completion_str, + task.input_tokens, + task.output_tokens, + **task.task_data, + ) + return float(reward) + + def _compute_reward_from_generation_task( + self, task: GenerationTask, **kwargs + ) -> float: + """Compute reward from a generation task. + + Parameters + ---------- + task : GenerationTask + The completed generation task. + **kwargs + Should contain 'task_data' with ground truth. + + Returns + ------- + float + The computed reward value. + """ + task_data = kwargs.get("task_data", {}) + prompt_str = kwargs.get("prompt_str", task.input_str or "") + + reward = self.reward_fn( + prompt_str, + task.output_str or "", + list(task.input_tokens or []), + list(task.output_tokens or []), + **task_data, + ) + return float(reward) + + async def aprocess(self, tasks: list[Task], **kwargs) -> Any: + """Process reward tasks asynchronously. + + This method computes rewards asynchronously using AsyncRewardWrapper. + + Parameters + ---------- + tasks : list[Task] + List of RLVRRewardTask objects to process. + **kwargs + Additional keyword arguments. + + Returns + ------- + list[Task] + The processed tasks with rewards computed. + """ + # Mark tasks with worker tag + for task in tasks: + task.worker_tag = self.WorkerTag.REWARD + + # Compute rewards asynchronously + self.scores = [] + for task in tasks: + if isinstance(task, RLVRRewardTask): + reward = await self._acompute_reward(task) + task.reward = reward + self.scores.append(reward) + + # Update the interaction object if provided + if task.interaction is not None: + task.interaction.reward = reward + + return tasks + + async def _acompute_reward(self, task: RLVRRewardTask) -> float: + """Compute reward asynchronously for an RLVR reward task. + + Parameters + ---------- + task : RLVRRewardTask + The reward task containing prompt, completion, and task data. + + Returns + ------- + float + The computed reward value. + """ + reward = await self.async_reward_fn( + task.prompt_str, + task.completion_str, + task.input_tokens, + task.output_tokens, + **task.task_data, + ) + return float(reward) + + +class PipelineTrajectoryMaker(Controller): + """Controller that composes generation and reward controllers into a pipeline. + + This controller orchestrates the full RLVR pipeline: + 1. Run generation via the generation controller + 2. Compute rewards via the reward controller + 3. Assemble results into InteractionWithTokenLogpReward objects + + Parameters + ---------- + generation_controller : Controller + The controller for text generation (e.g., NativeGenerationController). + reward_controller : RLVRRewardController + The controller for reward computation. + task_data : dict[str, Any] + Task data containing ground truth (e.g., "answer" field) for reward computation. + prompt_str : str + The prompt string used for generation. + + Example + ------- + ```python + from tensorrt_llm.scaffolding import NativeGenerationController + + gen_controller = NativeGenerationController() + reward_controller = RLVRRewardController(gsm8k_reward_fn) + trajectory_maker = PipelineTrajectoryMaker( + gen_controller, + reward_controller, + task_data={"answer": "42"}, + prompt_str="What is the answer?", + ) + ``` + """ + + def __init__( + self, + generation_controller: Controller, + reward_controller: RLVRRewardController, + task_data: dict[str, Any] | None = None, + prompt_str: str = "", + ): + """Initialize the pipeline trajectory maker. + + Parameters + ---------- + generation_controller : Controller + The generation controller. + reward_controller : RLVRRewardController + The reward controller. + task_data : dict[str, Any], optional + Task data containing ground truth for reward computation. + prompt_str : str, optional + The prompt string used for generation. + """ + super().__init__() + self.generation_controller = generation_controller + self.reward_controller = reward_controller + self.task_data = task_data if task_data is not None else {} + self.prompt_str = prompt_str + + def process(self, tasks: list[Task], **kwargs) -> Any: + """Process tasks through the generation and reward pipeline. + + Parameters + ---------- + tasks : list[Task] + List of generation tasks to process. + **kwargs + Additional keyword arguments. + + Yields + ------ + dict[str, InteractionWithTokenLogpReward] + Dictionary mapping task IDs to their interaction results. + """ + # Step 1: Run generation + yield from self.generation_controller.process(tasks, **kwargs) + + reward_tasks = [] + interactions = {} + + for i, task in enumerate(tasks): + if isinstance(task, GenerationTask): + # Create interaction object + interaction = self._create_interaction_from_task(task) + task_id = f"task_{i}" + interactions[task_id] = interaction + + # Create reward task using constructor-provided task_data and prompt_str + reward_task = RLVRRewardTask.create_from_generation_task( + gen_task=task, + prompt_str=self.prompt_str or task.input_str or "", + task_data=self.task_data, + interaction=interaction, + ) + reward_tasks.append(reward_task) + + # Step 3: Process reward tasks + yield from self.reward_controller.process(reward_tasks, **kwargs) + + # The interactions now have rewards set + # Return as the final result + yield interactions + + def _create_interaction_from_task( + self, task: GenerationTask + ) -> InteractionWithTokenLogpReward: + """Create an InteractionWithTokenLogpReward from a generation task. + + Parameters + ---------- + task : GenerationTask + The completed generation task. + + Returns + ------- + InteractionWithTokenLogpReward + The interaction object with model response data. + """ + from areal.api.io_struct import ModelResponse + + # Build ModelResponse from task data + input_tokens = list(task.input_tokens or []) + output_tokens = list(task.output_tokens or []) + output_logprobs = task.customized_result_fields.get("output_logprobs", []) + output_versions = task.customized_result_fields.get("output_versions", []) + + # Create ModelResponse + model_response = ModelResponse( + input_tokens=input_tokens, + output_tokens=output_tokens, + output_logprobs=list(output_logprobs) + if output_logprobs + else [0.0] * len(output_tokens), + output_versions=list(output_versions) + if output_versions + else [-1] * len(output_tokens), + ) + + # Create interaction + interaction = InteractionWithTokenLogpReward( + model_response=model_response, + reward=None, # Will be set by reward controller + ) + + return interaction + + +class ChatTracer(TaskCollection): + """TaskCollection for tracing multi-turn chat conversations. + + This class traces ChatTask objects during the controller's process execution. + A multi-turn conversation uses the same ChatTask object across multiple yields, + allowing us to track the evolution of the conversation. + + The tracer: + 1. In `before_yield`: Records the state of ChatTask before worker execution + 2. In `after_yield`: Captures the new messages added by the worker and creates + InteractionWithTokenLogpReward objects + + The traced results can be exported via `get_trace_results()`, which returns + a dict[str, InteractionWithTokenLogpReward] similar to client.py's export_interactions. + + Parameters + ---------- + reward_discount : float + Discount factor for backward reward propagation across turns. + export_style : str + Export style for interactions: 'concat' (tree structure) or 'individual'. + + Example + ------- + ```python + tracer = ChatTracer(reward_discount=0.9, export_style="individual") + # Used via with_task_collection decorator or TraceTrajectoryMaker + ``` + """ + + def __init__( + self, + reward_discount: float = 1.0, + export_style: str = "individual", + ): + super().__init__() + self.reward_discount = reward_discount + self.export_style = export_style + + # Cache for storing interactions, similar to InteractionCache in client.py + self._cache = InteractionCache() + + def before_yield(self, tasks: list[Task]): + """Called before tasks are yielded to workers. + + Parameters + ---------- + tasks : list[Task] + List of tasks about to be yielded. + """ + pass + + def after_yield(self, tasks: list[Task]): + """Called after tasks return from workers. + + Creates InteractionWithTokenLogpReward objects for each ChatTask. + Uses task.completion.id as the interaction ID. + + Parameters + ---------- + tasks : list[Task] + List of tasks that have been processed by workers. + """ + for task in tasks: + if not isinstance(task, ChatTask): + continue + + interaction = self._create_interaction_from_chat_task(task) + # Use completion.id as the interaction key + completion_id = task.completion.id + self._cache[completion_id] = interaction + + def _create_interaction_from_chat_task( + self, + task: ChatTask, + ) -> InteractionWithTokenLogpReward: + """Create an InteractionWithTokenLogpReward from a ChatTask. + + Parameters + ---------- + task : ChatTask + The ChatTask. Must contain a `completion` attribute + with the ChatCompletion object. + + Returns + ------- + InteractionWithTokenLogpReward + The interaction object capturing this turn. + """ + from areal.api.io_struct import ModelResponse + + # Extract all messages + messages = [ + msg.to_dict() if hasattr(msg, "to_dict") else msg for msg in task.messages + ] + + # Create ModelResponse from task data + input_tokens = list(task.input_tokens or []) + output_tokens = list(task.output_tokens or []) + + model_response = ModelResponse( + input_tokens=input_tokens, + output_tokens=output_tokens, + output_logprobs=[0.0] * len(output_tokens), + output_versions=[-1] * len(output_tokens), + ) + + # Get completion from task (ChatTask will contain the ChatCompletion) + completion = task.completion + + interaction = InteractionWithTokenLogpReward( + model_response=model_response, + reward=None, + messages=messages, + output_message_list=[], + completion=completion, + chat_template_type=self.export_style, + ) + + return interaction + + def get_trace_results(self) -> dict[str, InteractionWithTokenLogpReward]: + """Export traced interactions. + + Returns the traced interactions in the specified export style. + Applies reward discount before export if configured. + + Returns + ------- + dict[str, InteractionWithTokenLogpReward] + Dictionary mapping interaction IDs to their data. + + See Also + -------- + client.py : export_interactions method for similar functionality + """ + if len(self._cache) == 0: + return {} + + return self._cache.export_interactions( + style=self.export_style, + reward_discount=self.reward_discount, + ) + + def clear(self) -> None: + """Clear all traced data.""" + self._cache.clear() + + +@with_task_collection("chat_tracer", ChatTracer) +class TraceTrajectoryMaker(Controller): + """Controller that traces ChatTask objects during rollout using ChatTracer. + + This controller uses the @with_task_collection decorator to automatically + apply ChatTracer's before_yield and after_yield hooks around each yield + in the rollout controller's process execution. + + A multi-turn conversation uses the same ChatTask object, which is traced + across all yields. The trace results are stored in the TraceGenerationTask + after processing. + + Parameters + ---------- + rollout_controller : Controller + The controller for rollout (e.g., a chat or agent controller). + reward_controller : Controller + The controller for computing rewards on traced interactions. + + Example + ------- + ```python + from tensorrt_llm.scaffolding import NativeGenerationController + + chat_controller = SomeChatController() + reward_controller = RLVRRewardController(gsm8k_reward_fn) + + trace_maker = TraceTrajectoryMaker( + rollout_controller=chat_controller, + reward_controller=reward_controller, + ) + + # Process tasks + result = trace_maker.generate(prompt) + + # Or use process directly + task = TraceGenerationTask.create_from_prompt(prompt) + for _ in trace_maker.process([task]): + pass + trace_results = task.trace_results + ``` + """ + + def __init__( + self, + rollout_controller: Controller, + reward_controller: Controller, + ): + """Initialize the trace trajectory maker. + + Parameters + ---------- + rollout_controller : Controller + The controller for rollout execution. + reward_controller : Controller + The controller for reward computation. + """ + super().__init__() + self.rollout_controller = rollout_controller + self.reward_controller = reward_controller + + def process(self, tasks: list[Task], **kwargs) -> Any: + """Process tasks through the rollout and reward pipeline with tracing. + + This method: + 1. Extracts the generation_task from the TraceGenerationTask + 2. Runs the rollout_controller.process() with ChatTracer tracing + 3. Gets trace results from the ChatTracer + 4. Creates ChatRewardTask objects for each traced interaction + 5. Runs the reward_controller.process() to compute rewards + 6. Stores the trace results in the original task + + Parameters + ---------- + tasks : list[Task] + List of TraceGenerationTask objects to process. + **kwargs + Additional keyword arguments. + + Yields + ------ + Any + Results from the controllers. + """ + # Get the generation task from the first TraceGenerationTask + task = tasks[0] + if isinstance(task, TraceGenerationTask): + generation_task = task.generation_task + else: + generation_task = task + + # Run rollout with tracing (ChatTracer hooks applied via decorator) + yield from self.rollout_controller.process([generation_task], **kwargs) + + # Get trace results from the ChatTracer (registered via decorator) + chat_tracer = self.task_collections["chat_tracer"] + trace_results = chat_tracer.get_trace_results() + + # Create reward tasks for each traced interaction + reward_tasks = [ + ChatRewardTask.create_from_trace_result(interaction_id, interaction) + for interaction_id, interaction in trace_results.items() + ] + + # Run reward computation + if reward_tasks: + yield from self.reward_controller.process(reward_tasks, **kwargs) + + # Update trace_results with computed rewards + for reward_task in reward_tasks: + if ( + reward_task.interaction is not None + and reward_task.reward is not None + ): + reward_task.interaction.reward = reward_task.reward + + # Store trace results in the original task + if isinstance(task, TraceGenerationTask): + task.trace_results = trace_results + + def generate(self, prompt: str, **kwargs) -> Any: + """Generate with tracing from a prompt string. + + Parameters + ---------- + prompt : str + The input prompt. + **kwargs + Additional keyword arguments. + + Returns + ------- + Any + The scaffolding output. + """ + task = TraceGenerationTask.create_from_prompt(prompt) + + yield from self.process([task], **kwargs) + + return task.create_scaffolding_output() diff --git a/areal/experimental/scaffolding/task.py b/areal/experimental/scaffolding/task.py new file mode 100644 index 0000000000..762d762311 --- /dev/null +++ b/areal/experimental/scaffolding/task.py @@ -0,0 +1,232 @@ +""" +RLVR Tasks for Scaffolding Framework. + +This module provides task definitions for RLVR (Reinforcement Learning with +Verifiable Rewards) that integrate with TensorRT-LLM's scaffolding framework. +""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import TYPE_CHECKING, Any + +from areal.experimental.scaffolding._compat import ( + ChatTask, + GenerationTask, + ScaffoldingOutput, + Task, +) + +if TYPE_CHECKING: + from areal.experimental.openai.types import InteractionWithTokenLogpReward + + +@dataclass +class RLVRRewardTask(Task): + """Task for computing RLVR (verifiable) rewards. + + This task contains the necessary information to verify whether a generated + response is correct and compute the corresponding reward. + + Attributes + ---------- + prompt_str : str + The prompt string that was used for generation. + completion_str : str + The generated completion string to verify. + input_tokens : list[int] + The input token IDs. + output_tokens : list[int] + The output token IDs. + output_logprobs : list[float] + The log probabilities of output tokens. + output_versions : list[int] + The weight versions for output tokens. + task_data : dict[str, Any] + Additional task data containing ground truth (e.g., "answer" field). + interaction : InteractionWithTokenLogpReward + The interaction object to store the computed reward. + reward : float + The computed reward value (output field, set after processing). + """ + + # Input fields + prompt_str: str = field(default="") + completion_str: str = field(default="") + input_tokens: list[int] = field(default_factory=list) + output_tokens: list[int] = field(default_factory=list) + output_logprobs: list[float] = field(default_factory=list) + output_versions: list[int] = field(default_factory=list) + task_data: dict[str, Any] = field(default_factory=dict) + + # The interaction object to update with the reward + interaction: InteractionWithTokenLogpReward | None = None + + # Output field + reward: float | None = None + + @staticmethod + def create_from_generation_task( + gen_task: GenerationTask, + prompt_str: str, + task_data: dict[str, Any], + interaction: InteractionWithTokenLogpReward | None = None, + ) -> RLVRRewardTask: + """Create a reward task from a completed generation task. + + Parameters + ---------- + gen_task : GenerationTask + The completed generation task with output. + prompt_str : str + The original prompt string. + task_data : dict[str, Any] + Task data containing ground truth answer. + interaction : InteractionWithTokenLogpReward, optional + The interaction object to update with reward. + + Returns + ------- + RLVRRewardTask + The reward task ready for processing. + """ + reward_task = RLVRRewardTask( + prompt_str=prompt_str, + completion_str=gen_task.output_str or "", + input_tokens=list(gen_task.input_tokens or []), + output_tokens=list(gen_task.output_tokens or []), + output_logprobs=list( + gen_task.customized_result_fields.get("output_logprobs", []) + ), + output_versions=list( + gen_task.customized_result_fields.get("output_versions", []) + ), + task_data=task_data, + interaction=interaction, + ) + return reward_task + + +@dataclass +class TraceGenerationTask(Task): + """Task for tracing multi-turn generation with ChatTracer. + + This task wraps a ChatTask (or GenerationTask) for tracing purposes. + The trace results are stored after processing. + + Attributes + ---------- + generation_task : ChatTask | GenerationTask + The underlying task to be processed and traced. + trace_results : dict[str, InteractionWithTokenLogpReward] + The traced interaction results (output field, set after processing). + """ + + # The underlying generation/chat task + generation_task: ChatTask | GenerationTask | None = None + + # Output field - trace results after processing + trace_results: dict[str, InteractionWithTokenLogpReward] | None = None + + @staticmethod + def create_from_prompt(prompt: str) -> TraceGenerationTask: + """Create a TraceGenerationTask from a prompt string. + + Parameters + ---------- + prompt : str + The input prompt string. + + Returns + ------- + TraceGenerationTask + The task ready for processing. + """ + # Create underlying ChatTask + chat_task = ChatTask.create_from_prompt(prompt) + return TraceGenerationTask(generation_task=chat_task) + + @staticmethod + def create_from_chat_task(chat_task: ChatTask) -> TraceGenerationTask: + """Create a TraceGenerationTask from an existing ChatTask. + + Parameters + ---------- + chat_task : ChatTask + The ChatTask to wrap. + + Returns + ------- + TraceGenerationTask + The task ready for processing. + """ + return TraceGenerationTask(generation_task=chat_task) + + def create_scaffolding_output(self) -> ScaffoldingOutput: + """Create a ScaffoldingOutput from the trace results. + + Returns + ------- + ScaffoldingOutput + The output containing traced results. + """ + # Return the trace results as the output + if self.generation_task is not None and hasattr( + self.generation_task, "output_str" + ): + return ScaffoldingOutput( + text=self.generation_task.output_str or "", + token_ids=list(self.generation_task.output_tokens or []), + ) + return ScaffoldingOutput(text="", token_ids=[]) + + +@dataclass +class ChatRewardTask(Task): + """Task for computing rewards on traced chat interactions. + + This task contains a traced InteractionWithTokenLogpReward and is used + by the reward controller to compute and set rewards. + + Attributes + ---------- + interaction : InteractionWithTokenLogpReward + The traced interaction to compute reward for. + interaction_id : str + The ID of the interaction. + reward : float + The computed reward value (output field, set after processing). + """ + + # The traced interaction + interaction: InteractionWithTokenLogpReward | None = None + + # Interaction ID for reference + interaction_id: str = field(default="") + + # Output field + reward: float | None = None + + @staticmethod + def create_from_trace_result( + interaction_id: str, + interaction: InteractionWithTokenLogpReward, + ) -> ChatRewardTask: + """Create a ChatRewardTask from a trace result. + + Parameters + ---------- + interaction_id : str + The ID of the interaction. + interaction : InteractionWithTokenLogpReward + The traced interaction. + + Returns + ------- + ChatRewardTask + The reward task ready for processing. + """ + return ChatRewardTask( + interaction=interaction, + interaction_id=interaction_id, + ) diff --git a/areal/experimental/scaffolding/worker.py b/areal/experimental/scaffolding/worker.py new file mode 100644 index 0000000000..479ae1dd88 --- /dev/null +++ b/areal/experimental/scaffolding/worker.py @@ -0,0 +1,231 @@ +""" +Worker implementations for Scaffolding Framework. + +This module provides Worker implementations that wrap AReaL inference engines +for use with TensorRT-LLM's scaffolding framework. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +import openai + +from areal.experimental.scaffolding._compat import ( + AssistantMessage, + ChatTask, + GenerationTask, + OpenaiWorker, + TaskStatus, +) + +if TYPE_CHECKING: + from areal.engine.sglang_remote import RemoteSGLangEngine + + +class SGLangWorker(OpenaiWorker): + """Worker that wraps an SGLang engine for scaffolding. + + This worker connects to an SGLang server via its OpenAI-compatible API + and handles generation and chat tasks. + + Parameters + ---------- + async_client : openai.AsyncOpenAI + The OpenAI async client configured to connect to SGLang server. + model : str + The model name to use for requests. + engine : RemoteSGLangEngine + The underlying SGLang engine (kept for reference and potential future use). + """ + + def __init__( + self, + async_client: openai.AsyncOpenAI, + model: str, + engine: RemoteSGLangEngine, + ): + super().__init__(async_client, model, kv_cache_hint_enabled=False) + self.engine = engine + + async def chat_handler(self, task: ChatTask) -> TaskStatus: + """Handle chat completion requests. + + This method extends the base OpenaiWorker's chat handler to also + store the ChatCompletion object in the task for tracing purposes. + + Parameters + ---------- + task : ChatTask + The chat task to process. + + Returns + ------- + TaskStatus + The status of the task execution. + """ + params = self.convert_task_params(task) + params["messages"] = task.messages_to_dict_content() + params["model"] = self.model + if task.tools is not None: + params["tools"] = [tool.to_dict() for tool in task.tools] + + try: + response = await self.async_client.chat.completions.create(**params) + + # Store the completion in the task for tracing + task.completion = response + + task.finish_reason = response.choices[0].finish_reason + if hasattr(response, "perf_metrics"): + task.perf_metrics = response.perf_metrics + + content = response.choices[0].message.content + reasoning = getattr(response.choices[0].message, "reasoning", None) + reasoning_content = getattr( + response.choices[0].message, "reasoning_content", None + ) + tool_calls = response.choices[0].message.tool_calls + + task.messages.append( + AssistantMessage(content, reasoning, reasoning_content, tool_calls) + ) + + if task.enable_token_counting and response.usage: + task.prompt_tokens_num = response.usage.prompt_tokens + task.completion_tokens_num = response.usage.completion_tokens + if ( + hasattr(response.usage, "completion_tokens_details") + and response.usage.completion_tokens_details is not None + ): + task.reasoning_tokens_num = ( + response.usage.completion_tokens_details.reasoning_tokens + ) + + return TaskStatus.SUCCESS + + except Exception as e: + print(f"SGLang chat client exception: {e}") + return TaskStatus.WORKER_EXECEPTION + + async def generation_handler(self, task: GenerationTask) -> TaskStatus: + """Handle text generation requests. + + Parameters + ---------- + task : GenerationTask + The generation task to process. + + Returns + ------- + TaskStatus + The status of the task execution. + """ + params = self.convert_task_params(task) + + try: + response = await self.async_client.completions.create(**params) + + task.output_str = response.choices[0].text + if hasattr(response.choices[0], "token_ids"): + task.output_tokens = response.choices[0].token_ids + task.finish_reason = response.choices[0].finish_reason + if hasattr(response.choices[0], "logprobs"): + task.logprobs = response.choices[0].logprobs + if hasattr(response, "perf_metrics"): + task.perf_metrics = response.perf_metrics + + return TaskStatus.SUCCESS + + except Exception as e: + print(f"SGLang completion client exception: {e}") + return TaskStatus.WORKER_EXECEPTION + + # Register task handlers + task_handlers = { + GenerationTask: generation_handler, + ChatTask: chat_handler, + } + + +def CreateWorkerFromEngine( + engine: RemoteSGLangEngine, + model: str = "default", +) -> SGLangWorker: + """Create a scaffolding Worker from an AReaL SGLang engine. + + This function creates a Worker that wraps the given SGLang engine, + allowing it to be used with TensorRT-LLM's scaffolding framework. + The worker uses the SGLang server's OpenAI-compatible API. + + Parameters + ---------- + engine : RemoteSGLangEngine + The AReaL SGLang inference engine (must be initialized). + model : str, optional + The model name to use for API requests. Defaults to "default". + + Returns + ------- + SGLangWorker + A Worker instance that can be used with ScaffoldingLlm. + + Example + ------- + ```python + from areal.engine.sglang_remote import RemoteSGLangEngine + from areal.experimental.scaffolding import CreateWorkerFromEngine + + # Initialize the engine + engine = RemoteSGLangEngine(config) + engine.initialize() + + # Create a worker + worker = CreateWorkerFromEngine(engine) + + # Use with ScaffoldingLlm + from tensorrt_llm.scaffolding import ScaffoldingLlm, NativeGenerationController + llm = ScaffoldingLlm( + controller, + {NativeGenerationController.WorkerTag.GENERATION: worker}, + ) + ``` + + Raises + ------ + RuntimeError + If the engine is not initialized. + """ + if not engine.initialized: + raise RuntimeError( + "Engine must be initialized before creating a worker. " + "Call engine.initialize() first." + ) + + # Get the server address from the engine + # The internal engine stores server info + internal_engine = engine._engine + server_addrs = internal_engine._server_addrs + + if not server_addrs: + raise RuntimeError("No server addresses found in engine.") + + # Use the first server address for the OpenAI client + # SGLang servers support OpenAI-compatible API at /v1/ + base_url = server_addrs[0] + if not base_url.startswith("http"): + base_url = f"http://{base_url}" + if not base_url.endswith("/v1"): + base_url = f"{base_url}/v1" + + # Create an async OpenAI client pointing to the SGLang server + async_client = openai.AsyncOpenAI( + base_url=base_url, + api_key="EMPTY", # SGLang doesn't require API key by default + ) + + return SGLangWorker( + async_client=async_client, + model=model, + engine=engine, + ) diff --git a/areal/experimental/scaffolding/workflow.py b/areal/experimental/scaffolding/workflow.py new file mode 100644 index 0000000000..1ab2c28937 --- /dev/null +++ b/areal/experimental/scaffolding/workflow.py @@ -0,0 +1,263 @@ +""" +ScaffoldingWorkflow - RolloutWorkflow with generation and reward via Scaffolding. + +Architecture +------------ +- Generation: via scaffolding Worker (SGLangWorker calls SGLang OpenAI API) +- Reward: via scaffolding RLVRRewardController +- Logprobs: placeholder (0.0) since recompute_logprob=true in training config + causes the actor to recompute exact logprobs during PPO update. +- Worker & ScaffoldingLlm: lazily created from engine server addresses, + exposed for subclasses (e.g., multi-turn workflows). +""" + +from __future__ import annotations + +from collections.abc import Callable +from typing import Any + +import torch +from transformers import PreTrainedTokenizerFast + +from areal.api.cli_args import GenerationHyperparameters +from areal.api.engine_api import InferenceEngine +from areal.api.workflow_api import RolloutWorkflow +from areal.core import workflow_context +from areal.experimental.scaffolding._compat import ( + GenerationTask, + NativeGenerationController, + ScaffoldingLlm, +) +from areal.experimental.scaffolding.controllers import ( + PipelineTrajectoryMaker, + RLVRRewardController, +) +from areal.experimental.scaffolding.task import RLVRRewardTask +from areal.experimental.scaffolding.worker import SGLangWorker +from areal.utils import logging, stats_tracker +from areal.utils.dynamic_import import import_from_string +from areal.utils.perf_tracer import session_context, trace_session + +logger = logging.getLogger("ScaffoldingWorkflow") + + +class ScaffoldingWorkflow(RolloutWorkflow): + """RolloutWorkflow with generation and reward via scaffolding components. + + Both generation and reward computation go through scaffolding: + - Generation: SGLangWorker calls SGLang's OpenAI-compatible completions API + - Reward: RLVRRewardController computes verifiable rewards + + Since the OpenAI API does not return per-token logprobs in AReaL's format, + placeholder logprobs are used. Set ``recompute_logprob: true`` in the actor + config so the training engine recomputes exact logprobs during PPO update. + + Parameters + ---------- + reward_fn : Callable | str + The reward function, or an importable string path. + gconfig : GenerationHyperparameters + Generation hyperparameters. + tokenizer : PreTrainedTokenizerFast | str + Tokenizer or path to load it. + enable_thinking : bool + Whether to enable thinking tokens. + """ + + def __init__( + self, + reward_fn: Callable[..., Any] | str, + gconfig: GenerationHyperparameters, + tokenizer: PreTrainedTokenizerFast | str, + enable_thinking: bool = False, + ): + if isinstance(reward_fn, str): + reward_fn = import_from_string(reward_fn) + self.reward_fn = reward_fn + + self.tokenizer = tokenizer + if isinstance(self.tokenizer, str): + from areal.utils.hf_utils import load_hf_tokenizer + + self.tokenizer = load_hf_tokenizer(self.tokenizer) + self.gconfig = gconfig.new_with_stop_and_pad_token_ids(self.tokenizer) + self.enable_thinking = enable_thinking + + # Scaffolding controllers + self.reward_controller = RLVRRewardController(self.reward_fn) + self.gen_controller = NativeGenerationController() + + # Lazily created from engine server addresses + self.worker: SGLangWorker | None = None + self.trajectory_maker: PipelineTrajectoryMaker | None = None + self.scaffolding_llm: ScaffoldingLlm | None = None + + def _lazy_init_scaffolding(self, engine: InferenceEngine) -> None: + """Create Worker, PipelineTrajectoryMaker, and ScaffoldingLlm.""" + import openai + + addr = engine.addresses[0] + base_url = f"http://{addr}" + if not base_url.endswith("/v1"): + base_url = f"{base_url}/v1" + + async_client = openai.AsyncOpenAI(base_url=base_url, api_key="EMPTY") + self.worker = SGLangWorker( + async_client=async_client, model="default", engine=engine + ) + + self.trajectory_maker = PipelineTrajectoryMaker( + self.gen_controller, self.reward_controller + ) + + self.scaffolding_llm = ScaffoldingLlm( + self.trajectory_maker, + {NativeGenerationController.WorkerTag.GENERATION: self.worker}, + ) + logger.info(f"Initialized scaffolding components with server at {addr}") + + async def _generate_via_worker( + self, prompt_str: str, input_ids: list[int] + ) -> GenerationTask: + """Run generation through scaffolding Worker (SGLang OpenAI API). + + Parameters + ---------- + prompt_str : str + The prompt string. + input_ids : list[int] + The tokenized input IDs. + + Returns + ------- + GenerationTask + Completed task with output_str and output_tokens. + """ + # Build generation params for SGLang completions API + stop_strings = [] + if self.gconfig.stop_token_ids: + for tid in self.gconfig.stop_token_ids: + decoded = self.tokenizer.decode([tid]) + if decoded: + stop_strings.append(decoded) + + response = await self.worker.async_client.completions.create( + model=self.worker.model, + prompt=prompt_str, + max_tokens=self.gconfig.max_new_tokens, + temperature=self.gconfig.temperature or 1.0, + stop=stop_strings or None, + ) + + output_str = response.choices[0].text + # Tokenize to get output token IDs + output_token_ids = self.tokenizer.encode( + output_str, add_special_tokens=False + ) + + # Package as a GenerationTask (scaffolding data structure) + gen_task = GenerationTask( + input_str=prompt_str, + input_tokens=input_ids, + output_str=output_str, + output_tokens=output_token_ids, + finish_reason=response.choices[0].finish_reason, + ) + return gen_task + + @trace_session("reward") + async def _compute_rewards_via_controller( + self, + gen_task: GenerationTask, + prompt_str: str, + task_data: dict[str, Any], + ) -> float: + """Compute reward via scaffolding RLVRRewardController.""" + reward_task = RLVRRewardTask( + prompt_str=prompt_str, + completion_str=gen_task.output_str or "", + input_tokens=list(gen_task.input_tokens or []), + output_tokens=list(gen_task.output_tokens or []), + task_data=task_data, + ) + for _ in self.reward_controller.process([reward_task]): + pass + return float(reward_task.reward) + + @session_context() + async def _collect_samples( + self, + prompt_str: str, + input_ids: list[int], + task_data: dict[str, Any], + ) -> tuple[GenerationTask, float]: + """Generate via Worker, compute reward via Controller.""" + gen_task = await self._generate_via_worker(prompt_str, input_ids) + + reward = await self._compute_rewards_via_controller( + gen_task, prompt_str, task_data + ) + stats_tracker.get(workflow_context.stat_scope()).scalar(reward=reward) + + return gen_task, reward + + async def arun_episode( + self, engine: InferenceEngine, data: dict[str, Any] + ) -> dict[str, torch.Tensor]: + """Run a single episode via scaffolding pipeline. + + 1. Generation: SGLangWorker -> SGLang completions API + 2. Reward: RLVRRewardController.process() + 3. Output: tensor dict for PPO training + + Note: logprobs are placeholders (0.0). Set ``recompute_logprob: true`` + in actor config so the training engine computes exact logprobs. + + Parameters + ---------- + engine : InferenceEngine + The inference engine (used for server addresses on first call). + data : dict[str, Any] + Input data containing messages and ground truth. + + Returns + ------- + dict[str, torch.Tensor] + Trajectory tensors for PPO training. + """ + if self.worker is None: + self._lazy_init_scaffolding(engine) + + # Tokenize prompt + input_ids = list( + self.tokenizer.apply_chat_template( + data["messages"], + tokenize=True, + add_generation_prompt=True, + enable_thinking=self.enable_thinking, + ) + ) + prompt_str = self.tokenizer.decode(input_ids) + + # Scaffolding pipeline: Worker (generate) + Controller (reward) + gen_task, reward = await self._collect_samples( + prompt_str, input_ids, data + ) + + # Build tensor dict for PPO training + output_tokens = list(gen_task.output_tokens or []) + seq = input_ids + output_tokens + # Placeholder logprobs — recompute_logprob=true will replace these + logprobs = [0.0] * len(seq) + loss_mask = [0] * len(input_ids) + [1] * len(output_tokens) + versions = [-1] * len(seq) + + res = { + "input_ids": torch.tensor(seq, dtype=torch.int32), + "loss_mask": torch.tensor(loss_mask, dtype=torch.int32), + "logprobs": torch.tensor(logprobs, dtype=torch.float32), + "versions": torch.tensor(versions, dtype=torch.int32), + "attention_mask": torch.ones(len(seq), dtype=torch.bool), + "rewards": torch.tensor(reward, dtype=torch.float32), + } + return {k: v.unsqueeze(0) for k, v in res.items()} diff --git a/areal/reward/__init__.py b/areal/reward/__init__.py index 2a00d44f43..f54bf05dd0 100644 --- a/areal/reward/__init__.py +++ b/areal/reward/__init__.py @@ -1,5 +1,5 @@ -from math_verify.metric import math_metric -from math_verify.parser import ExprExtractionConfig, LatexExtractionConfig +from math_verify.grader import verify as math_verify_verify +from math_verify.parser import ExprExtractionConfig, LatexExtractionConfig, parse from areal.utils import logging @@ -27,6 +27,11 @@ def get_custom_reward_fn(path: str, **kwargs): class MathVerifyWorker: """Thin wrapper over math_verify with configurable extraction/precision. + Uses ``parse()`` + ``verify()`` directly instead of ``math_metric()`` + so that signal-based timeouts can be disabled (``parsing_timeout=None``, + ``timeout_seconds=None``). This avoids ``signal.alarm()`` which only + works in the main thread. + Args: try_extract_without_anchor: When False, only answers with explicit anchors (e.g., "answer = 1", "final answer = 1") are matched. When True, @@ -38,27 +43,41 @@ class MathVerifyWorker: """ def __init__(self, try_extract_without_anchor=True, precision: int = 6): - self.verify_func = math_metric( - gold_extraction_target=( - ExprExtractionConfig( - try_extract_without_anchor=try_extract_without_anchor - ), - LatexExtractionConfig(), + self.gold_extraction_target = ( + ExprExtractionConfig( + try_extract_without_anchor=try_extract_without_anchor ), - pred_extraction_target=( - ExprExtractionConfig( - try_extract_without_anchor=try_extract_without_anchor - ), - LatexExtractionConfig(), + LatexExtractionConfig(), + ) + self.pred_extraction_target = ( + ExprExtractionConfig( + try_extract_without_anchor=try_extract_without_anchor ), - precision=precision, + LatexExtractionConfig(), ) + self.precision = precision def verify(self, response: str, ground_truth: str) -> float: - # ground_truth_parsable = "\\boxed{" + ground_truth + "}" try: - ret_score, _ = self.verify_func([ground_truth], [response]) - return float(ret_score) + gold_parsed = parse( + ground_truth, + extraction_config=self.gold_extraction_target, + parsing_timeout=None, + ) + pred_parsed = parse( + response, + extraction_config=self.pred_extraction_target, + parsing_timeout=None, + ) + if not gold_parsed or not pred_parsed: + return 0.0 + result = math_verify_verify( + gold_parsed, + pred_parsed, + float_rounding=self.precision, + timeout_seconds=None, + ) + return 1.0 if result else 0.0 except Exception: logger.warning( f"Exception in MathVerifyWorker.verify for response={response} and ground_truth={ground_truth}", diff --git a/examples/scaffolding/README.md b/examples/scaffolding/README.md new file mode 100644 index 0000000000..9840045275 --- /dev/null +++ b/examples/scaffolding/README.md @@ -0,0 +1,225 @@ +# Scaffolding Framework Examples for AReaL + +This directory contains examples demonstrating how to use the TensorRT-LLM Scaffolding +framework with AReaL for reinforcement learning training. + +## Overview + +The scaffolding framework provides a modular and extensible way to compose inference-time +compute methods with RL training. It decouples the inference logic (Controllers) from the +execution backend (Workers), enabling flexible composition of different methods. + +### Key Components + +1. **Controller**: Defines the inference-time compute logic (e.g., generation, reward + computation) +2. **Worker**: Handles the actual execution of tasks (e.g., TRT-LLM, OpenAI API) +3. **ScaffoldingLlm**: Orchestrates controllers and workers together +4. **ScaffoldingWorkflow**: Wraps ScaffoldingLlm as a RolloutWorkflow for AReaL training + +### AReaL-Specific Components + +The following components are implemented in `areal/experimental/scaffolding/`: + +- **`CreateWorkerFromEngine`**: Creates a scaffolding Worker from AReaL's InferenceEngine + (e.g., RemoteSGLangEngine). The returned Worker is similar to scaffolding's `OpenaiWorker` + but integrated with AReaL's engine. + +- **`RLVRRewardController`**: A Controller that computes rewards for generated samples + using verifiable reward functions (e.g., math answer verification). + +- **`PipelineTrajectoryMaker`**: A Controller that composes generation and reward + controllers into a pipeline that produces training trajectories. + +- **`ScaffoldingWorkflow`**: A `RolloutWorkflow` implementation that wraps ScaffoldingLlm + for integration with AReaL's training pipeline. + +## RLVR Example with GSM8K + +### Quick Start + +```bash +python examples/scaffolding/gsm8k_rlvr_scaffolding.py \ + --config examples/scaffolding/gsm8k_rlvr_scaffolding.yaml +``` + +### Architecture + +The scaffolding workflow follows this pattern from the RFC: + +```python +# Step 1: Create Worker from the SGLang engine +rollout_worker = CreateWorkerFromEngine(engine) + +# Step 2: Create controllers +rollout_controller = NativeGenerationController() +reward_controller = RLVRRewardController(gsm8k_reward_fn) + +# Step 3: Create trajectory maker (composes the controllers) +trajectory_maker = PipelineTrajectoryMaker(rollout_controller, reward_controller) + +# Step 4: Create ScaffoldingLlm (orchestrates controllers with workers) +scaffolding_llm = ScaffoldingLlm( + trajectory_maker, + {NativeGenerationController.WorkerTag.GENERATION: rollout_worker}, +) + +# Step 5: Create ScaffoldingWorkflow (wraps as RolloutWorkflow) +scaffolding_workflow = ScaffoldingWorkflow(scaffolding_llm) +``` + +### Data Flow Diagram + +``` + ┌─────────────────────────────────────────────────┐ + │ ScaffoldingWorkflow │ + │ │ + │ ┌───────────────────────────────────────────┐ │ + │ │ ScaffoldingLlm │ │ + │ │ │ │ + │ │ ┌─────────────────────────────────────┐ │ │ + │ │ │ PipelineTrajectoryMaker │ │ │ + │ │ │ │ │ │ + │ │ │ ┌───────────────────────────────┐ │ │ │ +Data ─────────────────────────┼──┼──┼──► NativeGenerationController │ │ │ │ + │ │ │ │ (from tensorrt_llm) │ │ │ │ + │ │ │ └───────────────┬───────────────┘ │ │ │ + │ │ │ │ │ │ │ + │ │ │ ▼ │ │ │ + │ │ │ ┌───────────────────────────────┐ │ │ │ + │ │ │ │ RLVRRewardController │ │ │ │ + │ │ │ │ (from areal.experimental) │ │ │ │ + │ │ │ └───────────────┬───────────────┘ │ │ │ + │ │ │ │ │ │ │ + │ │ └──────────────────┼──────────────────┘ │ │ + │ │ │ │ │ + │ └─────────────────────┼─────────────────────┘ │ + │ │ │ + └────────────────────────┼────────────────────────┘ + │ + ▼ Trajectories + ┌─────────────────────────────┐ + │ PPOTrainer │ + │ (GRPO/PPO Training) │ + └─────────────────────────────┘ + │ + via CreateWorkerFromEngine │ + ▼ + ┌─────────────────────────────────────────┐ + │ RemoteSGLangEngine │ + │ (AReaL Inference Backend) │ + └─────────────────────────────────────────┘ +``` + +### How It Works + +1. **Engine Initialization**: `RemoteSGLangEngine` is initialized with the rollout + configuration and connected to the model server. + +2. **Worker Creation**: `CreateWorkerFromEngine(engine)` wraps the engine into a + scaffolding-compatible Worker. This allows scaffolding controllers to use AReaL's + inference backends. + +3. **Controller Pipeline**: + - `NativeGenerationController()`: Handles text generation by yielding `GenerationTask` + objects to the Worker. + - `RLVRRewardController(reward_fn)`: Computes rewards for generated samples using + the provided reward function. + - `PipelineTrajectoryMaker(gen_ctrl, reward_ctrl)`: Composes these controllers into + a pipeline that produces training trajectories. + +4. **ScaffoldingLlm**: Orchestrates the trajectory maker with the worker, handling the + async execution of tasks. + +5. **ScaffoldingWorkflow**: Wraps the ScaffoldingLlm as a `RolloutWorkflow` that can be + used directly with AReaL's `PPOTrainer`. + +6. **Training**: The trainer calls the workflow to generate trajectories, which are then + used for GRPO/PPO training. + +### Configuration + +See `gsm8k_rlvr_scaffolding.yaml` for the full configuration. Key options: + +```yaml +# Model configuration +pretrain_path: Qwen/Qwen2.5-3B-Instruct +tokenizer_path: Qwen/Qwen2.5-3B-Instruct + +# Generation hyperparameters +gconfig: + max_new_tokens: 1024 + temperature: 1.0 + top_p: 1.0 + n_samples: 8 + +# Inference engine configuration +engine: + type: sglang + tp: 1 + max_model_len: 4096 +``` + +## Extending the Framework + +### Custom Reward Controllers + +You can create custom reward controllers by subclassing the base Controller: + +```python +from tensorrt_llm.scaffolding import Controller + +class CustomRewardController(Controller): + def __init__(self, reward_fn): + super().__init__() + self.reward_fn = reward_fn + + def process(self, tasks, **kwargs): + # Compute rewards for completed generation tasks + for task in tasks: + reward = self.reward_fn( + prompt=task.input_str, + completion=task.output_str, + **kwargs + ) + task.customized_result_fields["reward"] = reward + yield tasks +``` + +### Custom Trajectory Makers + +For different RL algorithms, you may need different trajectory formats: + +```python +from tensorrt_llm.scaffolding import Controller +import torch + +class CustomTrajectoryMaker(Controller): + def __init__(self, generation_controller, reward_controller): + super().__init__() + self.generation_controller = generation_controller + self.reward_controller = reward_controller + + def process(self, tasks, **kwargs): + # Run generation + yield from self.generation_controller.process(tasks, **kwargs) + + # Run reward computation + yield from self.reward_controller.process(tasks, **kwargs) + + # Build trajectories + trajectories = [] + for task in tasks: + trajectory = { + "input_ids": torch.tensor(task.output_tokens), + "rewards": torch.tensor(task.customized_result_fields["reward"]), + } + trajectories.append(trajectory) + yield trajectories +``` + +## References + +- [TensorRT-LLM Scaffolding README](https://github.com/NVIDIA/TensorRT-LLM/tree/main/tensorrt_llm/scaffolding) +- [AReaL Workflow Documentation](../../docs/customization/workflow.md) +- [RFC: Scaffolding Integration](https://github.com/inclusionAI/AReaL/issues/818) diff --git a/examples/scaffolding/gsm8k_rlvr_scaffolding.py b/examples/scaffolding/gsm8k_rlvr_scaffolding.py new file mode 100644 index 0000000000..bbc2ffdab2 --- /dev/null +++ b/examples/scaffolding/gsm8k_rlvr_scaffolding.py @@ -0,0 +1,61 @@ +""" +RLVR (Reinforcement Learning with Verifiable Rewards) Example using Scaffolding Framework. + +This example demonstrates how to use the scaffolding framework for RLVR training +on the GSM8K math dataset. The ScaffoldingWorkflow uses AReaL's engine for +generation and scaffolding controllers for reward computation. + +Usage: + python examples/scaffolding/gsm8k_rlvr_scaffolding.py \ + --config examples/scaffolding/gsm8k_rlvr_scaffolding.yaml \ + +scheduler.type=local experiment_name=areal trial_name=scaffolding +""" + +import sys + +from areal.api.cli_args import GRPOConfig, load_expr_config +from areal.dataset import get_custom_dataset +from areal.experimental.trainer import PPOTrainer +from areal.utils.hf_utils import load_hf_tokenizer + + +def main(args): + """Main entry point for RLVR training with scaffolding.""" + config, _ = load_expr_config(args, GRPOConfig) + tokenizer = load_hf_tokenizer(config.tokenizer_path) + + train_dataset = get_custom_dataset( + split="train", + dataset_config=config.train_dataset, + tokenizer=tokenizer, + ) + valid_dataset = get_custom_dataset( + split="test", + dataset_config=config.valid_dataset, + tokenizer=tokenizer, + ) + + workflow_kwargs = dict( + reward_fn="areal.reward.gsm8k.gsm8k_reward_fn", + gconfig=config.gconfig, + tokenizer=config.tokenizer_path, + enable_thinking=False, + ) + eval_workflow_kwargs = workflow_kwargs.copy() + eval_workflow_kwargs["gconfig"] = config.gconfig.new(temperature=0.6) + + with PPOTrainer( + config, + train_dataset=train_dataset, + valid_dataset=valid_dataset, + ) as trainer: + trainer.train( + workflow="areal.experimental.scaffolding.workflow.ScaffoldingWorkflow", + workflow_kwargs=workflow_kwargs, + eval_workflow="areal.experimental.scaffolding.workflow.ScaffoldingWorkflow", + eval_workflow_kwargs=eval_workflow_kwargs, + ) + + +if __name__ == "__main__": + main(sys.argv[1:]) diff --git a/examples/scaffolding/gsm8k_rlvr_scaffolding.yaml b/examples/scaffolding/gsm8k_rlvr_scaffolding.yaml new file mode 100644 index 0000000000..c2e55954fe --- /dev/null +++ b/examples/scaffolding/gsm8k_rlvr_scaffolding.yaml @@ -0,0 +1,177 @@ +# RLVR Scaffolding Example Configuration for GSM8K +# Compatible with GRPOConfig, for 8-GPU setup (4 inference + 4 training) + +experiment_name: gsm8k-rlvr-scaffolding +trial_name: trial0 + +seed: 1 +enable_offload: false +total_train_epochs: 10 +tokenizer_path: ${actor.path} + +cluster: + n_nodes: 1 + n_gpus_per_node: 8 + fileroot: /tmp/areal/experiments + name_resolve: + type: nfs + nfs_record_root: /tmp/areal/name_resolve + +allocation_mode: sglang:d4p1t1+d4p1t1 + +rollout: + experiment_name: ${experiment_name} + trial_name: ${trial_name} + max_concurrent_rollouts: 256 + queue_size: null + consumer_batch_size: ${train_dataset.batch_size} + max_head_offpolicyness: 2 + enable_rollout_tracing: false + scheduling_spec: ${actor.scheduling_spec} + fileroot: ${cluster.fileroot} + tokenizer_path: ${tokenizer_path} + dump_to_file: true + +gconfig: + n_samples: 8 + min_new_tokens: 0 + max_new_tokens: 1024 + greedy: false + temperature: 1.0 + +actor: + experiment_name: ${experiment_name} + trial_name: ${trial_name} + path: Qwen/Qwen2.5-3B-Instruct + init_from_scratch: false + disable_dropout: true + gradient_checkpointing: true + dtype: bfloat16 + mb_spec: + max_tokens_per_mb: 10240 + optimizer: + type: adam + lr: 1.0e-6 + weight_decay: 0.01 + beta1: 0.9 + beta2: 0.999 + eps: 1e-8 + lr_scheduler_type: constant + gradient_clipping: 1.0 + warmup_steps_proportion: 0.001 + eps_clip: 0.2 + temperature: ${gconfig.temperature} + reward_scaling: 10.0 + reward_bias: -0.5 + kl_ctl: 0.0 + ppo_n_minibatches: 1 + recompute_logprob: true + use_decoupled_loss: true + behav_imp_weight_cap: 5.0 + reward_norm: + mean_level: group + std_level: group + group_size: ${gconfig.n_samples} + adv_norm: + mean_level: batch + std_level: batch + weight_update_mode: disk + max_new_tokens: ${gconfig.max_new_tokens} + scheduling_spec: + - task_type: worker + port_count: 2 + gpu: 1 + mem: 32 + cmd: python3 -m areal.scheduler.rpc.rpc_server + env_vars: {} + +ref: + experiment_name: ${experiment_name} + trial_name: ${trial_name} + path: ${actor.path} + init_from_scratch: false + disable_dropout: true + dtype: ${actor.dtype} + mb_spec: + max_tokens_per_mb: 10240 + optimizer: null + scheduling_strategy: + type: colocation + target: actor + scheduling_spec: ${actor.scheduling_spec} + +# SGLang +sglang: + model_path: ${actor.path} + random_seed: ${seed} + skip_tokenizer_init: false + dtype: ${actor.dtype} + max_running_requests: null + context_length: 4096 + mem_fraction_static: 0.65 + +vllm: + model: ${actor.path} + seed: ${seed} + skip_tokenizer_init: false + dtype: ${actor.dtype} + max_model_len: 4096 + gpu_memory_utilization: 0.8 + +# Datasets +train_dataset: + batch_size: 256 + shuffle: true + pin_memory: true + num_workers: 4 + path: openai/gsm8k + type: rl + max_length: 2048 + +valid_dataset: + batch_size: 256 + pin_memory: true + num_workers: 4 + path: openai/gsm8k + type: rl + +# Utilities +saver: + experiment_name: ${experiment_name} + trial_name: ${trial_name} + fileroot: ${cluster.fileroot} + freq_epochs: 1 + freq_steps: null + freq_secs: null + +recover: + mode: disabled + experiment_name: ${experiment_name} + trial_name: ${trial_name} + fileroot: ${cluster.fileroot} + freq_epochs: 1 + freq_steps: null + freq_secs: 3600 + +evaluator: + experiment_name: ${experiment_name} + trial_name: ${trial_name} + fileroot: ${cluster.fileroot} + freq_epochs: 1 + freq_steps: null + freq_secs: null + +stats_logger: + experiment_name: ${experiment_name} + trial_name: ${trial_name} + fileroot: ${cluster.fileroot} + wandb: + mode: disabled + +perf_tracer: + experiment_name: ${experiment_name} + trial_name: ${trial_name} + fileroot: ${cluster.fileroot} + enabled: false + session_tracer: + enabled: false From ca3e3db01b7a660f5df9687fbb6e3e6b60631389 Mon Sep 17 00:00:00 2001 From: root Date: Thu, 19 Feb 2026 09:49:40 +0000 Subject: [PATCH 02/18] test: add unit tests for scaffolding TestTraceTrajectoryMaker Co-Authored-By: Claude Opus 4.6 --- areal/experimental/scaffolding/controllers.py | 298 +++--- areal/experimental/scaffolding/workflow.py | 2 +- .../experimental/scaffolding/__init__.py | 0 .../scaffolding/test_controllers.py | 931 ++++++++++++++++++ .../scaffolding/gsm8k_rlvr_scaffolding.py | 2 +- .../scaffolding/gsm8k_rlvr_scaffolding.yaml | 7 +- 6 files changed, 1086 insertions(+), 154 deletions(-) create mode 100644 areal/tests/experimental/scaffolding/__init__.py create mode 100644 areal/tests/experimental/scaffolding/test_controllers.py diff --git a/areal/experimental/scaffolding/controllers.py b/areal/experimental/scaffolding/controllers.py index d191317587..e8b68535fa 100644 --- a/areal/experimental/scaffolding/controllers.py +++ b/areal/experimental/scaffolding/controllers.py @@ -240,155 +240,6 @@ async def _acompute_reward(self, task: RLVRRewardTask) -> float: return float(reward) -class PipelineTrajectoryMaker(Controller): - """Controller that composes generation and reward controllers into a pipeline. - - This controller orchestrates the full RLVR pipeline: - 1. Run generation via the generation controller - 2. Compute rewards via the reward controller - 3. Assemble results into InteractionWithTokenLogpReward objects - - Parameters - ---------- - generation_controller : Controller - The controller for text generation (e.g., NativeGenerationController). - reward_controller : RLVRRewardController - The controller for reward computation. - task_data : dict[str, Any] - Task data containing ground truth (e.g., "answer" field) for reward computation. - prompt_str : str - The prompt string used for generation. - - Example - ------- - ```python - from tensorrt_llm.scaffolding import NativeGenerationController - - gen_controller = NativeGenerationController() - reward_controller = RLVRRewardController(gsm8k_reward_fn) - trajectory_maker = PipelineTrajectoryMaker( - gen_controller, - reward_controller, - task_data={"answer": "42"}, - prompt_str="What is the answer?", - ) - ``` - """ - - def __init__( - self, - generation_controller: Controller, - reward_controller: RLVRRewardController, - task_data: dict[str, Any] | None = None, - prompt_str: str = "", - ): - """Initialize the pipeline trajectory maker. - - Parameters - ---------- - generation_controller : Controller - The generation controller. - reward_controller : RLVRRewardController - The reward controller. - task_data : dict[str, Any], optional - Task data containing ground truth for reward computation. - prompt_str : str, optional - The prompt string used for generation. - """ - super().__init__() - self.generation_controller = generation_controller - self.reward_controller = reward_controller - self.task_data = task_data if task_data is not None else {} - self.prompt_str = prompt_str - - def process(self, tasks: list[Task], **kwargs) -> Any: - """Process tasks through the generation and reward pipeline. - - Parameters - ---------- - tasks : list[Task] - List of generation tasks to process. - **kwargs - Additional keyword arguments. - - Yields - ------ - dict[str, InteractionWithTokenLogpReward] - Dictionary mapping task IDs to their interaction results. - """ - # Step 1: Run generation - yield from self.generation_controller.process(tasks, **kwargs) - - reward_tasks = [] - interactions = {} - - for i, task in enumerate(tasks): - if isinstance(task, GenerationTask): - # Create interaction object - interaction = self._create_interaction_from_task(task) - task_id = f"task_{i}" - interactions[task_id] = interaction - - # Create reward task using constructor-provided task_data and prompt_str - reward_task = RLVRRewardTask.create_from_generation_task( - gen_task=task, - prompt_str=self.prompt_str or task.input_str or "", - task_data=self.task_data, - interaction=interaction, - ) - reward_tasks.append(reward_task) - - # Step 3: Process reward tasks - yield from self.reward_controller.process(reward_tasks, **kwargs) - - # The interactions now have rewards set - # Return as the final result - yield interactions - - def _create_interaction_from_task( - self, task: GenerationTask - ) -> InteractionWithTokenLogpReward: - """Create an InteractionWithTokenLogpReward from a generation task. - - Parameters - ---------- - task : GenerationTask - The completed generation task. - - Returns - ------- - InteractionWithTokenLogpReward - The interaction object with model response data. - """ - from areal.api.io_struct import ModelResponse - - # Build ModelResponse from task data - input_tokens = list(task.input_tokens or []) - output_tokens = list(task.output_tokens or []) - output_logprobs = task.customized_result_fields.get("output_logprobs", []) - output_versions = task.customized_result_fields.get("output_versions", []) - - # Create ModelResponse - model_response = ModelResponse( - input_tokens=input_tokens, - output_tokens=output_tokens, - output_logprobs=list(output_logprobs) - if output_logprobs - else [0.0] * len(output_tokens), - output_versions=list(output_versions) - if output_versions - else [-1] * len(output_tokens), - ) - - # Create interaction - interaction = InteractionWithTokenLogpReward( - model_response=model_response, - reward=None, # Will be set by reward controller - ) - - return interaction - - class ChatTracer(TaskCollection): """TaskCollection for tracing multi-turn chat conversations. @@ -538,6 +389,155 @@ def clear(self) -> None: self._cache.clear() +class PipelineTrajectoryMaker(Controller): + """Controller that composes generation and reward controllers into a pipeline. + + This controller orchestrates the full RLVR pipeline: + 1. Run generation via the generation controller + 2. Compute rewards via the reward controller + 3. Assemble results into InteractionWithTokenLogpReward objects + + Parameters + ---------- + generation_controller : Controller + The controller for text generation (e.g., NativeGenerationController). + reward_controller : RLVRRewardController + The controller for reward computation. + task_data : dict[str, Any] + Task data containing ground truth (e.g., "answer" field) for reward computation. + prompt_str : str + The prompt string used for generation. + + Example + ------- + ```python + from tensorrt_llm.scaffolding import NativeGenerationController + + gen_controller = NativeGenerationController() + reward_controller = RLVRRewardController(gsm8k_reward_fn) + trajectory_maker = PipelineTrajectoryMaker( + gen_controller, + reward_controller, + task_data={"answer": "42"}, + prompt_str="What is the answer?", + ) + ``` + """ + + def __init__( + self, + generation_controller: Controller, + reward_controller: RLVRRewardController, + task_data: dict[str, Any] | None = None, + prompt_str: str = "", + ): + """Initialize the pipeline trajectory maker. + + Parameters + ---------- + generation_controller : Controller + The generation controller. + reward_controller : RLVRRewardController + The reward controller. + task_data : dict[str, Any], optional + Task data containing ground truth for reward computation. + prompt_str : str, optional + The prompt string used for generation. + """ + super().__init__() + self.generation_controller = generation_controller + self.reward_controller = reward_controller + self.task_data = task_data if task_data is not None else {} + self.prompt_str = prompt_str + + def process(self, tasks: list[Task], **kwargs) -> Any: + """Process tasks through the generation and reward pipeline. + + Parameters + ---------- + tasks : list[Task] + List of generation tasks to process. + **kwargs + Additional keyword arguments. + + Yields + ------ + dict[str, InteractionWithTokenLogpReward] + Dictionary mapping task IDs to their interaction results. + """ + # Step 1: Run generation + yield from self.generation_controller.process(tasks, **kwargs) + + reward_tasks = [] + interactions = {} + + for i, task in enumerate(tasks): + if isinstance(task, GenerationTask): + # Create interaction object + interaction = self._create_interaction_from_task(task) + task_id = f"task_{i}" + interactions[task_id] = interaction + + # Create reward task using constructor-provided task_data and prompt_str + reward_task = RLVRRewardTask.create_from_generation_task( + gen_task=task, + prompt_str=self.prompt_str or task.input_str or "", + task_data=self.task_data, + interaction=interaction, + ) + reward_tasks.append(reward_task) + + # Step 3: Process reward tasks + yield from self.reward_controller.process(reward_tasks, **kwargs) + + # The interactions now have rewards set + # Return as the final result + yield interactions + + def _create_interaction_from_task( + self, task: GenerationTask + ) -> InteractionWithTokenLogpReward: + """Create an InteractionWithTokenLogpReward from a generation task. + + Parameters + ---------- + task : GenerationTask + The completed generation task. + + Returns + ------- + InteractionWithTokenLogpReward + The interaction object with model response data. + """ + from areal.api.io_struct import ModelResponse + + # Build ModelResponse from task data + input_tokens = list(task.input_tokens or []) + output_tokens = list(task.output_tokens or []) + output_logprobs = task.customized_result_fields.get("output_logprobs", []) + output_versions = task.customized_result_fields.get("output_versions", []) + + # Create ModelResponse + model_response = ModelResponse( + input_tokens=input_tokens, + output_tokens=output_tokens, + output_logprobs=list(output_logprobs) + if output_logprobs + else [0.0] * len(output_tokens), + output_versions=list(output_versions) + if output_versions + else [-1] * len(output_tokens), + ) + + # Create interaction + interaction = InteractionWithTokenLogpReward( + model_response=model_response, + reward=None, # Will be set by reward controller + ) + + return interaction + + @with_task_collection("chat_tracer", ChatTracer) class TraceTrajectoryMaker(Controller): """Controller that traces ChatTask objects during rollout using ChatTracer. diff --git a/areal/experimental/scaffolding/workflow.py b/areal/experimental/scaffolding/workflow.py index 1ab2c28937..35d0a9d736 100644 --- a/areal/experimental/scaffolding/workflow.py +++ b/areal/experimental/scaffolding/workflow.py @@ -22,7 +22,7 @@ from areal.api.cli_args import GenerationHyperparameters from areal.api.engine_api import InferenceEngine from areal.api.workflow_api import RolloutWorkflow -from areal.core import workflow_context +from areal import workflow_context from areal.experimental.scaffolding._compat import ( GenerationTask, NativeGenerationController, diff --git a/areal/tests/experimental/scaffolding/__init__.py b/areal/tests/experimental/scaffolding/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/areal/tests/experimental/scaffolding/test_controllers.py b/areal/tests/experimental/scaffolding/test_controllers.py new file mode 100644 index 0000000000..4fd2de3a86 --- /dev/null +++ b/areal/tests/experimental/scaffolding/test_controllers.py @@ -0,0 +1,931 @@ +"""Unit tests for scaffolding controllers (TraceTrajectoryMaker, PipelineTrajectoryMaker). + +Tests use fake workers/controllers that simulate LLM inference responses +without requiring an SGLang backend or GPU. + +Design Notes +------------ +- ``FakeGenerationController`` fills ``GenerationTask`` fields in-memory + (single-turn generation). +- ``FakeChatRolloutController`` appends assistant messages to ``ChatTask`` + across multiple yields (multi-turn chat). It manually calls + ``ChatTracer.before_yield / after_yield`` because the lightweight + ``with_task_collection`` decorator in ``_compat.py`` only attaches the + ``ChatTracer`` instance to the class; it does NOT wrap ``process`` to + invoke hooks automatically (that is the tensorrt_llm implementation's + responsibility). +- ``FakeChatRewardController`` assigns predetermined rewards to + ``ChatRewardTask`` objects. +- The lightweight ``with_task_collection`` creates the ``ChatTracer`` as a + **class-level** attribute, so all ``TraceTrajectoryMaker`` instances share + one tracer. Each test that uses it must call ``tracer.clear()`` (or use a + fresh instance) to avoid inter-test pollution. +""" + +from __future__ import annotations + +from typing import Any +from unittest.mock import MagicMock + +import pytest + +from areal.api.io_struct import ModelResponse +from areal.experimental.openai.types import InteractionWithTokenLogpReward +from areal.experimental.scaffolding._compat import ( + AssistantMessage, + ChatTask, + Controller, + GenerationTask, + Task, +) +from areal.experimental.scaffolding.controllers import ( + ChatTracer, + PipelineTrajectoryMaker, + RLVRRewardController, + TraceTrajectoryMaker, +) +from areal.experimental.scaffolding.task import ( + ChatRewardTask, + RLVRRewardTask, + TraceGenerationTask, +) + + +# --------------------------------------------------------------------------- +# Fake / stub helpers +# --------------------------------------------------------------------------- + +FAKE_INPUT_TOKENS = [101, 102, 103] +FAKE_OUTPUT_TOKENS = [201, 202, 203, 204] +FAKE_OUTPUT_STR = "42" +FAKE_PROMPT_STR = "What is the answer to life?" + + +def _simple_reward_fn( + prompt: str, + completions: str, + prompt_ids: list[int], + completion_ids: list[int], + **kwargs, +) -> float: + """Deterministic reward: 1.0 if completion contains the ground-truth answer, else 0.0.""" + answer = kwargs.get("answer", "") + return 1.0 if answer and answer in completions else 0.0 + + +class FakeGenerationController(Controller): + """Fills ``GenerationTask`` fields without calling any LLM backend.""" + + def __init__( + self, + output_str: str = FAKE_OUTPUT_STR, + output_tokens: list[int] | None = None, + input_tokens: list[int] | None = None, + output_logprobs: list[float] | None = None, + output_versions: list[int] | None = None, + ): + super().__init__() + self.output_str = output_str + self.output_tokens = output_tokens or FAKE_OUTPUT_TOKENS + self.input_tokens = input_tokens or FAKE_INPUT_TOKENS + self.output_logprobs = output_logprobs + self.output_versions = output_versions + + def process(self, tasks: list[Task], **kwargs) -> Any: + for task in tasks: + if isinstance(task, GenerationTask): + task.output_str = self.output_str + task.output_tokens = self.output_tokens + if task.input_tokens is None: + task.input_tokens = self.input_tokens + if self.output_logprobs is not None: + task.customized_result_fields["output_logprobs"] = ( + self.output_logprobs + ) + if self.output_versions is not None: + task.customized_result_fields["output_versions"] = ( + self.output_versions + ) + yield tasks + + +def _make_fake_completion(completion_id: str = "cmpl-001") -> MagicMock: + """Create a minimal fake ``ChatCompletion`` object.""" + completion = MagicMock() + completion.id = completion_id + completion.created = 1000 + completion.choices = [MagicMock()] + completion.choices[0].message.content = FAKE_OUTPUT_STR + completion.choices[0].finish_reason = "stop" + return completion + + +class FakeChatRolloutController(Controller): + """Simulates a multi-turn chat rollout by appending assistant messages. + + Each turn: + 1. Populates the ``ChatTask`` with a fake completion and tokens. + 2. Calls ``ChatTracer.before_yield`` / ``after_yield`` manually. + 3. Yields the tasks. + 4. Appends a follow-up user message before the next turn. + """ + + def __init__( + self, + n_turns: int = 2, + responses: list[str] | None = None, + output_tokens_per_turn: list[list[int]] | None = None, + input_tokens_per_turn: list[list[int]] | None = None, + ): + super().__init__() + self.n_turns = n_turns + self.responses = responses or [f"response_{i}" for i in range(n_turns)] + self.output_tokens_per_turn = output_tokens_per_turn or [ + [300 + i * 10 + j for j in range(4)] for i in range(n_turns) + ] + self.input_tokens_per_turn = input_tokens_per_turn or [ + FAKE_INPUT_TOKENS for _ in range(n_turns) + ] + # Set by the test to allow manual ChatTracer hook invocation. + self._tracer: ChatTracer | None = None + + def process(self, tasks: list[Task], **kwargs) -> Any: + for turn_idx in range(self.n_turns): + for task in tasks: + if isinstance(task, ChatTask): + task.messages.append( + AssistantMessage(content=self.responses[turn_idx]) + ) + task.output_tokens = self.output_tokens_per_turn[turn_idx] + task.input_tokens = self.input_tokens_per_turn[turn_idx] + task.completion = _make_fake_completion( + completion_id=f"cmpl-turn-{turn_idx}" + ) + + if self._tracer is not None: + self._tracer.before_yield(tasks) + + yield tasks + + if self._tracer is not None: + self._tracer.after_yield(tasks) + + if turn_idx < self.n_turns - 1: + for task in tasks: + if isinstance(task, ChatTask): + task.messages.append( + {"role": "user", "content": f"follow-up-{turn_idx}"} + ) + + +class FakeChatRewardController(Controller): + """Assigns predetermined rewards to ``ChatRewardTask`` objects.""" + + def __init__(self, rewards: list[float] | None = None, default_reward: float = 1.0): + super().__init__() + self.rewards = rewards + self.default_reward = default_reward + + def process(self, tasks: list[Task], **kwargs) -> Any: + for i, task in enumerate(tasks): + if isinstance(task, ChatRewardTask): + reward = ( + self.rewards[i] + if self.rewards is not None and i < len(self.rewards) + else self.default_reward + ) + task.reward = reward + if task.interaction is not None: + task.interaction.reward = reward + yield tasks + + +# --------------------------------------------------------------------------- +# Fixtures +# --------------------------------------------------------------------------- + + +@pytest.fixture(autouse=True) +def _reset_shared_tracer(): + """Reset the class-level ChatTracer before each test. + + The lightweight ``with_task_collection`` stores a single ``ChatTracer`` + instance on ``TraceTrajectoryMaker`` (class-level). We must clear it + between tests so that cached interactions don't leak. + """ + yield + tracer = TraceTrajectoryMaker.task_collections.get("chat_tracer") + if tracer is not None: + tracer.clear() + + +# =========================================================================== +# PipelineTrajectoryMaker tests +# =========================================================================== + + +class TestPipelineTrajectoryMaker: + """Tests for PipelineTrajectoryMaker (single-turn generation + reward).""" + + def test_basic_pipeline_correct_answer(self): + """Pipeline should produce interaction with reward=1.0 for a correct answer.""" + gen_ctrl = FakeGenerationController(output_str="The answer is 42.") + reward_ctrl = RLVRRewardController(_simple_reward_fn) + + maker = PipelineTrajectoryMaker( + generation_controller=gen_ctrl, + reward_controller=reward_ctrl, + task_data={"answer": "42"}, + prompt_str=FAKE_PROMPT_STR, + ) + + task = GenerationTask(input_str=FAKE_PROMPT_STR, input_tokens=FAKE_INPUT_TOKENS) + results = list(maker.process([task])) + + # Three yields: generation, reward, interactions dict + assert len(results) == 3 + interactions = results[-1] + assert isinstance(interactions, dict) + assert len(interactions) == 1 + + interaction = list(interactions.values())[0] + assert isinstance(interaction, InteractionWithTokenLogpReward) + assert interaction.reward == 1.0 + + def test_basic_pipeline_wrong_answer(self): + """Pipeline should produce interaction with reward=0.0 for a wrong answer.""" + gen_ctrl = FakeGenerationController(output_str="I don't know") + reward_ctrl = RLVRRewardController(_simple_reward_fn) + + maker = PipelineTrajectoryMaker( + generation_controller=gen_ctrl, + reward_controller=reward_ctrl, + task_data={"answer": "42"}, + prompt_str=FAKE_PROMPT_STR, + ) + + task = GenerationTask(input_str=FAKE_PROMPT_STR, input_tokens=FAKE_INPUT_TOKENS) + results = list(maker.process([task])) + + interactions = results[-1] + interaction = list(interactions.values())[0] + assert interaction.reward == 0.0 + + def test_pipeline_multiple_tasks(self): + """Pipeline should handle multiple GenerationTasks in a single batch.""" + gen_ctrl = FakeGenerationController(output_str="42") + reward_ctrl = RLVRRewardController(_simple_reward_fn) + + maker = PipelineTrajectoryMaker( + generation_controller=gen_ctrl, + reward_controller=reward_ctrl, + task_data={"answer": "42"}, + prompt_str=FAKE_PROMPT_STR, + ) + + tasks = [ + GenerationTask(input_str=FAKE_PROMPT_STR, input_tokens=FAKE_INPUT_TOKENS), + GenerationTask(input_str="Another prompt", input_tokens=[111, 112]), + ] + results = list(maker.process(tasks)) + + interactions = results[-1] + assert len(interactions) == 2 + for interaction in interactions.values(): + assert interaction.reward == 1.0 + + def test_pipeline_interaction_has_model_response(self): + """Each interaction should contain a valid ModelResponse with tokens.""" + gen_ctrl = FakeGenerationController( + output_str="42", + output_tokens=[201, 202], + output_logprobs=[0.1, 0.2], + output_versions=[1, 1], + ) + reward_ctrl = RLVRRewardController(_simple_reward_fn) + + maker = PipelineTrajectoryMaker( + generation_controller=gen_ctrl, + reward_controller=reward_ctrl, + task_data={"answer": "42"}, + prompt_str=FAKE_PROMPT_STR, + ) + + task = GenerationTask(input_str=FAKE_PROMPT_STR, input_tokens=FAKE_INPUT_TOKENS) + results = list(maker.process([task])) + + interaction = list(results[-1].values())[0] + mr = interaction.model_response + assert isinstance(mr, ModelResponse) + assert mr.input_tokens == FAKE_INPUT_TOKENS + assert mr.output_tokens == [201, 202] + assert mr.output_logprobs == [0.1, 0.2] + assert mr.output_versions == [1, 1] + + def test_pipeline_uses_constructor_prompt_str(self): + """Reward controller should receive the prompt_str provided at construction.""" + received_prompts = [] + + def _capture(prompt, completions, prompt_ids, completion_ids, **kw): + received_prompts.append(prompt) + return 1.0 + + gen_ctrl = FakeGenerationController(output_str="42") + reward_ctrl = RLVRRewardController(_capture) + + maker = PipelineTrajectoryMaker( + generation_controller=gen_ctrl, + reward_controller=reward_ctrl, + task_data={"answer": "42"}, + prompt_str="custom prompt", + ) + + task = GenerationTask(input_str="input str", input_tokens=FAKE_INPUT_TOKENS) + list(maker.process([task])) + + assert received_prompts == ["custom prompt"] + + def test_pipeline_falls_back_to_input_str(self): + """When prompt_str is empty, should fall back to task.input_str.""" + received_prompts = [] + + def _capture(prompt, completions, prompt_ids, completion_ids, **kw): + received_prompts.append(prompt) + return 1.0 + + gen_ctrl = FakeGenerationController(output_str="42") + reward_ctrl = RLVRRewardController(_capture) + + maker = PipelineTrajectoryMaker( + generation_controller=gen_ctrl, + reward_controller=reward_ctrl, + task_data={"answer": "42"}, + prompt_str="", + ) + + task = GenerationTask(input_str="fallback prompt", input_tokens=FAKE_INPUT_TOKENS) + list(maker.process([task])) + + assert received_prompts[0] == "fallback prompt" + + def test_pipeline_reward_scores_tracked(self): + """RLVRRewardController should track scores in self.scores.""" + gen_ctrl = FakeGenerationController(output_str="42") + reward_ctrl = RLVRRewardController(_simple_reward_fn) + + maker = PipelineTrajectoryMaker( + generation_controller=gen_ctrl, + reward_controller=reward_ctrl, + task_data={"answer": "42"}, + prompt_str=FAKE_PROMPT_STR, + ) + + tasks = [ + GenerationTask(input_str=FAKE_PROMPT_STR, input_tokens=FAKE_INPUT_TOKENS), + GenerationTask(input_str=FAKE_PROMPT_STR, input_tokens=FAKE_INPUT_TOKENS), + ] + list(maker.process(tasks)) + + assert reward_ctrl.scores == [1.0, 1.0] + + def test_pipeline_default_logprobs_and_versions(self): + """When no logprobs/versions provided, interaction should use placeholders.""" + gen_ctrl = FakeGenerationController( + output_str="42", + output_tokens=[201, 202], + # No logprobs/versions supplied + ) + reward_ctrl = RLVRRewardController(_simple_reward_fn) + + maker = PipelineTrajectoryMaker( + generation_controller=gen_ctrl, + reward_controller=reward_ctrl, + task_data={"answer": "42"}, + prompt_str=FAKE_PROMPT_STR, + ) + + task = GenerationTask(input_str=FAKE_PROMPT_STR, input_tokens=FAKE_INPUT_TOKENS) + results = list(maker.process([task])) + + interaction = list(results[-1].values())[0] + mr = interaction.model_response + # Placeholders: [0.0] * output_len + assert mr.output_logprobs == [0.0, 0.0] + assert mr.output_versions == [-1, -1] + + +# =========================================================================== +# ChatTracer tests +# =========================================================================== + + +class TestChatTracer: + """Tests for ChatTracer (TaskCollection for tracing multi-turn chats).""" + + def test_after_yield_creates_interaction(self): + """after_yield should create an interaction for each ChatTask.""" + tracer = ChatTracer(export_style="individual") + task = ChatTask(messages=[{"role": "user", "content": "hello"}]) + task.completion = _make_fake_completion("cmpl-001") + task.input_tokens = [1, 2, 3] + task.output_tokens = [4, 5] + + tracer.after_yield([task]) + + results = tracer.get_trace_results() + assert len(results) == 1 + assert "cmpl-001" in results + + def test_multiple_turns_traced(self): + """Calling after_yield with different completions should trace all.""" + tracer = ChatTracer(export_style="individual") + task = ChatTask(messages=[{"role": "user", "content": "hello"}]) + + # Turn 1 + task.completion = _make_fake_completion("cmpl-turn-0") + task.input_tokens = [1, 2] + task.output_tokens = [3, 4] + tracer.after_yield([task]) + + # Turn 2 (same ChatTask, new completion) + task.messages.append(AssistantMessage(content="first response")) + task.completion = _make_fake_completion("cmpl-turn-1") + task.input_tokens = [1, 2, 3, 4] + task.output_tokens = [5, 6] + tracer.after_yield([task]) + + results = tracer.get_trace_results() + assert len(results) == 2 + assert "cmpl-turn-0" in results + assert "cmpl-turn-1" in results + + def test_tracer_interaction_has_model_response(self): + """Traced interactions should have ModelResponse with correct tokens.""" + tracer = ChatTracer(export_style="individual") + task = ChatTask(messages=[{"role": "user", "content": "hello"}]) + task.completion = _make_fake_completion("cmpl-001") + task.input_tokens = [10, 20] + task.output_tokens = [30, 40, 50] + + tracer.after_yield([task]) + + interaction = tracer.get_trace_results()["cmpl-001"] + assert interaction.model_response is not None + assert interaction.model_response.input_tokens == [10, 20] + assert interaction.model_response.output_tokens == [30, 40, 50] + + def test_tracer_clear(self): + """clear() should remove all traced data.""" + tracer = ChatTracer(export_style="individual") + task = ChatTask(messages=[{"role": "user", "content": "hello"}]) + task.completion = _make_fake_completion("cmpl-001") + task.input_tokens = [1] + task.output_tokens = [2] + + tracer.after_yield([task]) + assert len(tracer.get_trace_results()) == 1 + + tracer.clear() + assert tracer.get_trace_results() == {} + + def test_tracer_ignores_non_chat_tasks(self): + """after_yield should skip non-ChatTask objects.""" + tracer = ChatTracer(export_style="individual") + tracer.after_yield([GenerationTask(input_str="hello")]) + assert tracer.get_trace_results() == {} + + def test_tracer_empty_returns_empty(self): + """get_trace_results on a fresh tracer should return empty dict.""" + tracer = ChatTracer(export_style="individual") + assert tracer.get_trace_results() == {} + + +# =========================================================================== +# TraceTrajectoryMaker tests +# =========================================================================== + + +class TestTraceTrajectoryMaker: + """Tests for TraceTrajectoryMaker (multi-turn tracing + reward pipeline).""" + + @staticmethod + def _make_trace_maker( + n_turns: int = 2, + responses: list[str] | None = None, + rewards: list[float] | None = None, + default_reward: float = 1.0, + ) -> tuple[TraceTrajectoryMaker, FakeChatRolloutController]: + """Build a ``TraceTrajectoryMaker`` with fake sub-controllers. + + Also wires the class-level ``ChatTracer`` into the fake rollout + controller so it can call ``before_yield`` / ``after_yield`` hooks. + """ + rollout_ctrl = FakeChatRolloutController( + n_turns=n_turns, responses=responses + ) + reward_ctrl = FakeChatRewardController( + rewards=rewards, default_reward=default_reward + ) + maker = TraceTrajectoryMaker( + rollout_controller=rollout_ctrl, + reward_controller=reward_ctrl, + ) + chat_tracer = maker.task_collections["chat_tracer"] + rollout_ctrl._tracer = chat_tracer + return maker, rollout_ctrl + + def test_basic_trace_single_turn(self): + """Single-turn trace should produce one traced interaction with reward.""" + maker, _ = self._make_trace_maker( + n_turns=1, + responses=["The answer is 42"], + default_reward=1.0, + ) + + task = TraceGenerationTask.create_from_prompt("What is 6*7?") + list(maker.process([task])) + + assert task.trace_results is not None + assert len(task.trace_results) == 1 + interaction = list(task.trace_results.values())[0] + assert interaction.reward == 1.0 + + def test_multi_turn_trace(self): + """Multi-turn trace should produce one interaction per turn.""" + maker, _ = self._make_trace_maker( + n_turns=3, + responses=["step 1", "step 2", "final answer: 42"], + rewards=[0.0, 0.0, 1.0], + ) + + task = TraceGenerationTask.create_from_prompt("Solve step by step") + list(maker.process([task])) + + assert task.trace_results is not None + assert len(task.trace_results) == 3 + + def test_trace_rewards_assigned_correctly(self): + """Each traced interaction should get its designated reward.""" + maker, _ = self._make_trace_maker( + n_turns=2, + responses=["thinking...", "42"], + rewards=[0.5, 1.0], + ) + + task = TraceGenerationTask.create_from_prompt("What is the answer?") + list(maker.process([task])) + + assert task.trace_results is not None + rewards = [i.reward for i in task.trace_results.values()] + assert rewards == [0.5, 1.0] + + def test_trace_results_stored_in_task(self): + """trace_results should be set on the TraceGenerationTask after processing.""" + maker, _ = self._make_trace_maker(n_turns=1, default_reward=0.0) + + task = TraceGenerationTask.create_from_prompt("hello") + assert task.trace_results is None + + list(maker.process([task])) + + assert task.trace_results is not None + assert isinstance(task.trace_results, dict) + + def test_trace_with_plain_task_fallback(self): + """process should not crash when given a plain ChatTask.""" + maker, _ = self._make_trace_maker(n_turns=1, default_reward=1.0) + + chat_task = ChatTask.create_from_prompt("direct chat task") + list(maker.process([chat_task])) + # No assertion on trace_results — plain ChatTask doesn't store them. + + def test_trace_generation_task_create_from_chat_task(self): + """TraceGenerationTask.create_from_chat_task should wrap correctly.""" + chat_task = ChatTask.create_from_prompt("hello") + trace_task = TraceGenerationTask.create_from_chat_task(chat_task) + + assert trace_task.generation_task is chat_task + assert trace_task.trace_results is None + + def test_trace_generation_task_scaffolding_output(self): + """create_scaffolding_output should reflect generation_task fields.""" + gen_task = GenerationTask( + output_str="result text", output_tokens=[10, 20, 30] + ) + trace_task = TraceGenerationTask(generation_task=gen_task) + + output = trace_task.create_scaffolding_output() + assert output.text == "result text" + assert output.token_ids == [10, 20, 30] + + def test_trace_generation_task_scaffolding_output_empty(self): + """create_scaffolding_output with no generation_task should return empty.""" + trace_task = TraceGenerationTask() + output = trace_task.create_scaffolding_output() + assert output.text == "" + assert output.token_ids == [] + + def test_no_reward_tasks_when_no_traces(self): + """If rollout produces no traceable outputs, reward step should be skipped.""" + + class EmptyRolloutController(Controller): + def process(self, tasks, **kwargs): + yield tasks + + reward_ctrl = FakeChatRewardController(default_reward=1.0) + maker = TraceTrajectoryMaker( + rollout_controller=EmptyRolloutController(), + reward_controller=reward_ctrl, + ) + assert maker.task_collections.get("chat_tracer") is not None + + task = TraceGenerationTask.create_from_prompt("hello") + list(maker.process([task])) + + assert task.trace_results is not None + assert len(task.trace_results) == 0 + + def test_trace_interaction_model_response_tokens(self): + """Traced interactions should carry correct per-turn tokens.""" + maker, _ = self._make_trace_maker( + n_turns=2, + responses=["r0", "r1"], + default_reward=1.0, + ) + # Use distinct per-turn tokens + rollout_ctrl = maker.rollout_controller + rollout_ctrl.output_tokens_per_turn = [[10, 11], [20, 21, 22]] + rollout_ctrl.input_tokens_per_turn = [[1, 2], [1, 2, 3]] + + task = TraceGenerationTask.create_from_prompt("prompt") + list(maker.process([task])) + + interactions = list(task.trace_results.values()) + assert interactions[0].model_response.output_tokens == [10, 11] + assert interactions[0].model_response.input_tokens == [1, 2] + assert interactions[1].model_response.output_tokens == [20, 21, 22] + assert interactions[1].model_response.input_tokens == [1, 2, 3] + + +# =========================================================================== +# RLVRRewardController tests +# =========================================================================== + + +class TestRLVRRewardController: + """Tests for RLVRRewardController (reward computation).""" + + def test_compute_reward_correct(self): + """Should return 1.0 when completion contains the answer.""" + ctrl = RLVRRewardController(_simple_reward_fn) + task = RLVRRewardTask( + prompt_str="What is 2+2?", + completion_str="The answer is 4", + input_tokens=[1, 2, 3], + output_tokens=[4, 5], + task_data={"answer": "4"}, + ) + + list(ctrl.process([task])) + assert task.reward == 1.0 + assert ctrl.scores == [1.0] + + def test_compute_reward_wrong(self): + """Should return 0.0 when completion does not contain the answer.""" + ctrl = RLVRRewardController(_simple_reward_fn) + task = RLVRRewardTask( + prompt_str="What is 2+2?", + completion_str="I think it's 5", + input_tokens=[1, 2, 3], + output_tokens=[4, 5], + task_data={"answer": "4"}, + ) + + list(ctrl.process([task])) + assert task.reward == 0.0 + + def test_compute_reward_updates_interaction(self): + """Reward should be propagated to the attached interaction object.""" + ctrl = RLVRRewardController(_simple_reward_fn) + interaction = InteractionWithTokenLogpReward( + model_response=ModelResponse( + input_tokens=[1], + output_tokens=[2], + output_logprobs=[0.0], + output_versions=[-1], + ), + messages=[], + ) + task = RLVRRewardTask( + prompt_str="Q", + completion_str="42", + input_tokens=[1], + output_tokens=[2], + task_data={"answer": "42"}, + interaction=interaction, + ) + + list(ctrl.process([task])) + assert interaction.reward == 1.0 + + def test_compute_reward_batch(self): + """Should process multiple tasks and track all scores.""" + ctrl = RLVRRewardController(_simple_reward_fn) + tasks = [ + RLVRRewardTask( + prompt_str="Q1", + completion_str="42", + task_data={"answer": "42"}, + ), + RLVRRewardTask( + prompt_str="Q2", + completion_str="wrong", + task_data={"answer": "42"}, + ), + RLVRRewardTask( + prompt_str="Q3", + completion_str="also 42", + task_data={"answer": "42"}, + ), + ] + + list(ctrl.process(tasks)) + assert ctrl.scores == [1.0, 0.0, 1.0] + + def test_reward_from_generation_task(self): + """Should handle GenerationTask via customized_result_fields path.""" + ctrl = RLVRRewardController(_simple_reward_fn) + gen_task = GenerationTask( + input_str="What is 2+2?", + output_str="4", + input_tokens=[1, 2], + output_tokens=[3], + ) + + list( + ctrl.process( + [gen_task], + task_data={"answer": "4"}, + prompt_str="What is 2+2?", + ) + ) + assert gen_task.customized_result_fields["reward"] == 1.0 + + +# =========================================================================== +# RLVRRewardTask tests +# =========================================================================== + + +class TestRLVRRewardTask: + """Tests for RLVRRewardTask creation and data flow.""" + + def test_create_from_generation_task(self): + """create_from_generation_task should correctly populate all fields.""" + gen_task = GenerationTask( + input_str="prompt", + output_str="completion text", + input_tokens=[1, 2, 3], + output_tokens=[4, 5], + ) + gen_task.customized_result_fields["output_logprobs"] = [0.1, 0.2] + gen_task.customized_result_fields["output_versions"] = [1, 1] + + reward_task = RLVRRewardTask.create_from_generation_task( + gen_task=gen_task, + prompt_str="original prompt", + task_data={"answer": "42"}, + ) + + assert reward_task.prompt_str == "original prompt" + assert reward_task.completion_str == "completion text" + assert reward_task.input_tokens == [1, 2, 3] + assert reward_task.output_tokens == [4, 5] + assert reward_task.output_logprobs == [0.1, 0.2] + assert reward_task.output_versions == [1, 1] + assert reward_task.task_data == {"answer": "42"} + assert reward_task.reward is None + + def test_create_from_generation_task_no_output(self): + """Should handle GenerationTask with None output gracefully.""" + gen_task = GenerationTask() + + reward_task = RLVRRewardTask.create_from_generation_task( + gen_task=gen_task, + prompt_str="prompt", + task_data={}, + ) + + assert reward_task.completion_str == "" + assert reward_task.input_tokens == [] + assert reward_task.output_tokens == [] + + +# =========================================================================== +# ChatRewardTask tests +# =========================================================================== + + +class TestChatRewardTask: + """Tests for ChatRewardTask creation.""" + + def test_create_from_trace_result(self): + """create_from_trace_result should wrap an interaction correctly.""" + interaction = InteractionWithTokenLogpReward( + model_response=ModelResponse( + input_tokens=[1, 2], + output_tokens=[3, 4], + output_logprobs=[0.0, 0.0], + output_versions=[-1, -1], + ), + messages=[{"role": "user", "content": "hello"}], + ) + + task = ChatRewardTask.create_from_trace_result("id-001", interaction) + + assert task.interaction is interaction + assert task.interaction_id == "id-001" + assert task.reward is None + + +# =========================================================================== +# End-to-end integration tests +# =========================================================================== + + +class TestEndToEnd: + """Integration tests that exercise the full scaffolding rollout pipeline.""" + + def test_pipeline_e2e_tensor_dict_compatible(self): + """PipelineTrajectoryMaker interactions should be convertible to tensor dicts.""" + gen_ctrl = FakeGenerationController( + output_str="42", + output_tokens=[201, 202], + ) + reward_ctrl = RLVRRewardController(_simple_reward_fn) + + maker = PipelineTrajectoryMaker( + generation_controller=gen_ctrl, + reward_controller=reward_ctrl, + task_data={"answer": "42"}, + prompt_str=FAKE_PROMPT_STR, + ) + + task = GenerationTask(input_str=FAKE_PROMPT_STR, input_tokens=FAKE_INPUT_TOKENS) + results = list(maker.process([task])) + interaction = list(results[-1].values())[0] + + td = interaction.to_tensor_dict() + assert "input_ids" in td + assert "loss_mask" in td + assert "logprobs" in td + assert "rewards" in td + assert td["rewards"].item() == 1.0 + + def test_trace_e2e_multi_turn_with_rewards(self): + """Full multi-turn TraceTrajectoryMaker E2E with per-turn rewards.""" + rollout_ctrl = FakeChatRolloutController( + n_turns=3, + responses=["Let me think...", "Calculating...", "The answer is 42"], + ) + reward_ctrl = FakeChatRewardController(rewards=[0.0, 0.5, 1.0]) + + maker = TraceTrajectoryMaker( + rollout_controller=rollout_ctrl, + reward_controller=reward_ctrl, + ) + chat_tracer = maker.task_collections["chat_tracer"] + rollout_ctrl._tracer = chat_tracer + + task = TraceGenerationTask.create_from_prompt("Solve step by step: 6*7") + list(maker.process([task])) + + assert task.trace_results is not None + assert len(task.trace_results) == 3 + rewards = [i.reward for i in task.trace_results.values()] + assert rewards == [0.0, 0.5, 1.0] + + def test_trace_e2e_single_turn_generates_output(self): + """Single-turn trace should produce a valid interaction with reward.""" + rollout_ctrl = FakeChatRolloutController(n_turns=1, responses=["42"]) + reward_ctrl = FakeChatRewardController(default_reward=1.0) + + maker = TraceTrajectoryMaker( + rollout_controller=rollout_ctrl, + reward_controller=reward_ctrl, + ) + chat_tracer = maker.task_collections["chat_tracer"] + rollout_ctrl._tracer = chat_tracer + + task = TraceGenerationTask.create_from_prompt("What is 6*7?") + list(maker.process([task])) + + assert task.trace_results is not None + assert len(task.trace_results) == 1 + interaction = list(task.trace_results.values())[0] + assert interaction.reward == 1.0 + assert interaction.model_response is not None + + +if __name__ == "__main__": + pytest.main([__file__, "-v"]) diff --git a/examples/scaffolding/gsm8k_rlvr_scaffolding.py b/examples/scaffolding/gsm8k_rlvr_scaffolding.py index bbc2ffdab2..529cdf4dfd 100644 --- a/examples/scaffolding/gsm8k_rlvr_scaffolding.py +++ b/examples/scaffolding/gsm8k_rlvr_scaffolding.py @@ -15,7 +15,7 @@ from areal.api.cli_args import GRPOConfig, load_expr_config from areal.dataset import get_custom_dataset -from areal.experimental.trainer import PPOTrainer +from areal.trainer import PPOTrainer from areal.utils.hf_utils import load_hf_tokenizer diff --git a/examples/scaffolding/gsm8k_rlvr_scaffolding.yaml b/examples/scaffolding/gsm8k_rlvr_scaffolding.yaml index c2e55954fe..b272679709 100644 --- a/examples/scaffolding/gsm8k_rlvr_scaffolding.yaml +++ b/examples/scaffolding/gsm8k_rlvr_scaffolding.yaml @@ -82,7 +82,7 @@ actor: port_count: 2 gpu: 1 mem: 32 - cmd: python3 -m areal.scheduler.rpc.rpc_server + cmd: python3 -m areal.infra.rpc.rpc_server env_vars: {} ref: @@ -107,8 +107,9 @@ sglang: skip_tokenizer_init: false dtype: ${actor.dtype} max_running_requests: null - context_length: 4096 - mem_fraction_static: 0.65 + context_length: 2048 + mem_fraction_static: 0.5 + attention_backend: flashinfer vllm: model: ${actor.path} From 39809834e1b4ca9365d51578f651586f94c00dfd Mon Sep 17 00:00:00 2001 From: narutolhy <582909902@qq.com> Date: Sat, 21 Feb 2026 00:36:11 -0800 Subject: [PATCH 03/18] add worker test --- areal/experimental/scaffolding/worker.py | 3 + .../experimental/scaffolding/test_worker.py | 204 ++++++++++++++++++ 2 files changed, 207 insertions(+) create mode 100644 areal/tests/experimental/scaffolding/test_worker.py diff --git a/areal/experimental/scaffolding/worker.py b/areal/experimental/scaffolding/worker.py index 479ae1dd88..7ce6f5423e 100644 --- a/areal/experimental/scaffolding/worker.py +++ b/areal/experimental/scaffolding/worker.py @@ -122,6 +122,9 @@ async def generation_handler(self, task: GenerationTask) -> TaskStatus: The status of the task execution. """ params = self.convert_task_params(task) + params["model"] = self.model + if task.input_str is not None: + params["prompt"] = task.input_str try: response = await self.async_client.completions.create(**params) diff --git a/areal/tests/experimental/scaffolding/test_worker.py b/areal/tests/experimental/scaffolding/test_worker.py new file mode 100644 index 0000000000..99b6c9f787 --- /dev/null +++ b/areal/tests/experimental/scaffolding/test_worker.py @@ -0,0 +1,204 @@ +"""Tests for SGLangWorker with a real SGLang server (requires GPU).""" + +from __future__ import annotations + +import subprocess +import sys +import time +from unittest.mock import MagicMock + +import openai +import pytest +import requests + +from areal.api.cli_args import SGLangConfig +from areal.experimental.scaffolding._compat import ( + GenerationTask, + TaskStatus, +) +from areal.experimental.scaffolding.worker import SGLangWorker +from areal.tests.utils import get_model_path +from areal.utils import network, seeding +from areal.utils.hf_utils import load_hf_tokenizer +from areal.utils.proc import kill_process_tree + +# --------------------------------------------------------------------------- +# Constants +# --------------------------------------------------------------------------- + +EXPR_NAME = "test_scaffolding_worker" +MODEL_PATH = get_model_path( + "/storage/openpsi/models/Qwen__Qwen3-0.6B/", "Qwen/Qwen3-0.6B" +) +PORT, DIST_PORT = network.find_free_ports(2) +HOST = network.gethostip() +RUN_SERVER_TIMEOUT = 180 + + +# --------------------------------------------------------------------------- +# Fixtures +# --------------------------------------------------------------------------- + + +def _check_server_health(base_url: str) -> bool: + try: + response = requests.get(f"{base_url}/health", timeout=30) + return response.status_code == 200 + except requests.exceptions.RequestException: + return False + + +@pytest.fixture(scope="module") +def sglang_server(): + """Launch a real SGLang server. + + Uses skip_tokenizer_init=True (AReaL default). Chat template is applied + client-side via the tokenizer, matching ScaffoldingWorkflow behavior. + """ + seeding.set_random_seed(1, EXPR_NAME) + cmd = SGLangConfig.build_cmd( + sglang_config=SGLangConfig( + skip_tokenizer_init=False, + model_path=MODEL_PATH, + mem_fraction_static=0.3, + ), + host=HOST, + port=PORT, + tp_size=1, + base_gpu_id=0, + dist_init_addr=f"{HOST}:{DIST_PORT}", + ) + process = subprocess.Popen( + cmd, + stdout=sys.stdout, + stderr=sys.stdout, + ) + base_url = f"http://{HOST}:{PORT}" + tik = time.time() + while time.time() - tik < RUN_SERVER_TIMEOUT: + if _check_server_health(base_url): + break + time.sleep(1) + if time.time() - tik > RUN_SERVER_TIMEOUT: + kill_process_tree(process.pid, graceful=True) + raise RuntimeError("SGLang server launch timed out") + yield base_url + kill_process_tree(process.pid, graceful=True) + + +@pytest.fixture(scope="module") +def tokenizer(): + """Load the tokenizer for client-side chat template application.""" + return load_hf_tokenizer(MODEL_PATH) + + +@pytest.fixture(scope="module") +def sglang_worker(sglang_server): + """Create an SGLangWorker connected to the real SGLang server.""" + base_url = sglang_server + if not base_url.endswith("/v1"): + base_url = f"{base_url}/v1" + async_client = openai.AsyncOpenAI(base_url=base_url, api_key="EMPTY") + mock_engine = MagicMock() + return SGLangWorker( + async_client=async_client, + model="default", + engine=mock_engine, + ) + + +# --------------------------------------------------------------------------- +# Tests — generation_handler (uses /v1/completions, no tokenizer needed) +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_generation_handler(sglang_worker): + """generation_handler should generate text from real server.""" + task = GenerationTask( + input_str="What is 1 + 1? Answer briefly:", + input_tokens=[], + ) + status = await sglang_worker.generation_handler(task) + + assert status == TaskStatus.SUCCESS + assert task.output_str is not None + assert len(task.output_str) > 0 + assert task.finish_reason in ("stop", "length") + + +@pytest.mark.asyncio +async def test_generation_max_tokens(sglang_worker): + """generation_handler should respect max_tokens and finish with 'length'.""" + task = GenerationTask( + input_str="Write a very long essay about the history of mathematics.", + input_tokens=[], + ) + original_create = sglang_worker.async_client.completions.create + + async def _create_with_limit(**kwargs): + kwargs["max_tokens"] = 5 + return await original_create(**kwargs) + + sglang_worker.async_client.completions.create = _create_with_limit + try: + status = await sglang_worker.generation_handler(task) + assert status == TaskStatus.SUCCESS + assert task.finish_reason == "length" + assert task.output_str is not None + finally: + sglang_worker.async_client.completions.create = original_create + + +# --------------------------------------------------------------------------- +# Tests — generation with chat template (client-side, like ScaffoldingWorkflow) +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_generation_with_chat_template(sglang_worker, tokenizer): + """Client-side chat template + completions API, matching ScaffoldingWorkflow.""" + messages = [{"role": "user", "content": "What is the capital of France?"}] + input_ids = tokenizer.apply_chat_template( + messages, tokenize=True, add_generation_prompt=True + ) + prompt_str = tokenizer.decode(input_ids) + + task = GenerationTask( + input_str=prompt_str, + input_tokens=input_ids, + ) + status = await sglang_worker.generation_handler(task) + + assert status == TaskStatus.SUCCESS + assert task.output_str is not None + assert len(task.output_str) > 0 + + +@pytest.mark.asyncio +async def test_multi_turn_generation(sglang_worker, tokenizer): + """Multi-turn via client-side chat template + completions API.""" + # Turn 1 + messages = [{"role": "user", "content": "My name is Alice."}] + input_ids = tokenizer.apply_chat_template( + messages, tokenize=True, add_generation_prompt=True + ) + prompt_str = tokenizer.decode(input_ids) + + task = GenerationTask(input_str=prompt_str, input_tokens=input_ids) + status = await sglang_worker.generation_handler(task) + assert status == TaskStatus.SUCCESS + assert task.output_str is not None + + # Turn 2 — append assistant reply and new user message + messages.append({"role": "assistant", "content": task.output_str}) + messages.append({"role": "user", "content": "What is my name?"}) + input_ids_2 = tokenizer.apply_chat_template( + messages, tokenize=True, add_generation_prompt=True + ) + prompt_str_2 = tokenizer.decode(input_ids_2) + + task2 = GenerationTask(input_str=prompt_str_2, input_tokens=input_ids_2) + status = await sglang_worker.generation_handler(task2) + assert status == TaskStatus.SUCCESS + assert task2.output_str is not None From f2fdc4908c896f90a9cbd5d81011239489f8dc7e Mon Sep 17 00:00:00 2001 From: narutolhy <582909902@qq.com> Date: Sat, 21 Feb 2026 00:37:37 -0800 Subject: [PATCH 04/18] format --- areal/experimental/scaffolding/workflow.py | 8 +--- areal/reward/__init__.py | 8 +--- examples/scaffolding/README.md | 50 +++++++++++----------- 3 files changed, 30 insertions(+), 36 deletions(-) diff --git a/areal/experimental/scaffolding/workflow.py b/areal/experimental/scaffolding/workflow.py index 35d0a9d736..16d1abf554 100644 --- a/areal/experimental/scaffolding/workflow.py +++ b/areal/experimental/scaffolding/workflow.py @@ -151,9 +151,7 @@ async def _generate_via_worker( output_str = response.choices[0].text # Tokenize to get output token IDs - output_token_ids = self.tokenizer.encode( - output_str, add_special_tokens=False - ) + output_token_ids = self.tokenizer.encode(output_str, add_special_tokens=False) # Package as a GenerationTask (scaffolding data structure) gen_task = GenerationTask( @@ -240,9 +238,7 @@ async def arun_episode( prompt_str = self.tokenizer.decode(input_ids) # Scaffolding pipeline: Worker (generate) + Controller (reward) - gen_task, reward = await self._collect_samples( - prompt_str, input_ids, data - ) + gen_task, reward = await self._collect_samples(prompt_str, input_ids, data) # Build tensor dict for PPO training output_tokens = list(gen_task.output_tokens or []) diff --git a/areal/reward/__init__.py b/areal/reward/__init__.py index f54bf05dd0..724fb1cb60 100644 --- a/areal/reward/__init__.py +++ b/areal/reward/__init__.py @@ -44,15 +44,11 @@ class MathVerifyWorker: def __init__(self, try_extract_without_anchor=True, precision: int = 6): self.gold_extraction_target = ( - ExprExtractionConfig( - try_extract_without_anchor=try_extract_without_anchor - ), + ExprExtractionConfig(try_extract_without_anchor=try_extract_without_anchor), LatexExtractionConfig(), ) self.pred_extraction_target = ( - ExprExtractionConfig( - try_extract_without_anchor=try_extract_without_anchor - ), + ExprExtractionConfig(try_extract_without_anchor=try_extract_without_anchor), LatexExtractionConfig(), ) self.precision = precision diff --git a/examples/scaffolding/README.md b/examples/scaffolding/README.md index 9840045275..e754edd8fb 100644 --- a/examples/scaffolding/README.md +++ b/examples/scaffolding/README.md @@ -5,25 +5,26 @@ framework with AReaL for reinforcement learning training. ## Overview -The scaffolding framework provides a modular and extensible way to compose inference-time -compute methods with RL training. It decouples the inference logic (Controllers) from the -execution backend (Workers), enabling flexible composition of different methods. +The scaffolding framework provides a modular and extensible way to compose +inference-time compute methods with RL training. It decouples the inference logic +(Controllers) from the execution backend (Workers), enabling flexible composition of +different methods. ### Key Components 1. **Controller**: Defines the inference-time compute logic (e.g., generation, reward computation) -2. **Worker**: Handles the actual execution of tasks (e.g., TRT-LLM, OpenAI API) -3. **ScaffoldingLlm**: Orchestrates controllers and workers together -4. **ScaffoldingWorkflow**: Wraps ScaffoldingLlm as a RolloutWorkflow for AReaL training +1. **Worker**: Handles the actual execution of tasks (e.g., TRT-LLM, OpenAI API) +1. **ScaffoldingLlm**: Orchestrates controllers and workers together +1. **ScaffoldingWorkflow**: Wraps ScaffoldingLlm as a RolloutWorkflow for AReaL training ### AReaL-Specific Components The following components are implemented in `areal/experimental/scaffolding/`: -- **`CreateWorkerFromEngine`**: Creates a scaffolding Worker from AReaL's InferenceEngine - (e.g., RemoteSGLangEngine). The returned Worker is similar to scaffolding's `OpenaiWorker` - but integrated with AReaL's engine. +- **`CreateWorkerFromEngine`**: Creates a scaffolding Worker from AReaL's + InferenceEngine (e.g., RemoteSGLangEngine). The returned Worker is similar to + scaffolding's `OpenaiWorker` but integrated with AReaL's engine. - **`RLVRRewardController`**: A Controller that computes rewards for generated samples using verifiable reward functions (e.g., math answer verification). @@ -31,8 +32,8 @@ The following components are implemented in `areal/experimental/scaffolding/`: - **`PipelineTrajectoryMaker`**: A Controller that composes generation and reward controllers into a pipeline that produces training trajectories. -- **`ScaffoldingWorkflow`**: A `RolloutWorkflow` implementation that wraps ScaffoldingLlm - for integration with AReaL's training pipeline. +- **`ScaffoldingWorkflow`**: A `RolloutWorkflow` implementation that wraps + ScaffoldingLlm for integration with AReaL's training pipeline. ## RLVR Example with GSM8K @@ -116,25 +117,26 @@ Data ───────────────────────── 1. **Engine Initialization**: `RemoteSGLangEngine` is initialized with the rollout configuration and connected to the model server. -2. **Worker Creation**: `CreateWorkerFromEngine(engine)` wraps the engine into a +1. **Worker Creation**: `CreateWorkerFromEngine(engine)` wraps the engine into a scaffolding-compatible Worker. This allows scaffolding controllers to use AReaL's inference backends. -3. **Controller Pipeline**: - - `NativeGenerationController()`: Handles text generation by yielding `GenerationTask` - objects to the Worker. - - `RLVRRewardController(reward_fn)`: Computes rewards for generated samples using - the provided reward function. - - `PipelineTrajectoryMaker(gen_ctrl, reward_ctrl)`: Composes these controllers into - a pipeline that produces training trajectories. +1. **Controller Pipeline**: -4. **ScaffoldingLlm**: Orchestrates the trajectory maker with the worker, handling the + - `NativeGenerationController()`: Handles text generation by yielding + `GenerationTask` objects to the Worker. + - `RLVRRewardController(reward_fn)`: Computes rewards for generated samples using the + provided reward function. + - `PipelineTrajectoryMaker(gen_ctrl, reward_ctrl)`: Composes these controllers into a + pipeline that produces training trajectories. + +1. **ScaffoldingLlm**: Orchestrates the trajectory maker with the worker, handling the async execution of tasks. -5. **ScaffoldingWorkflow**: Wraps the ScaffoldingLlm as a `RolloutWorkflow` that can be +1. **ScaffoldingWorkflow**: Wraps the ScaffoldingLlm as a `RolloutWorkflow` that can be used directly with AReaL's `PPOTrainer`. -6. **Training**: The trainer calls the workflow to generate trajectories, which are then +1. **Training**: The trainer calls the workflow to generate trajectories, which are then used for GRPO/PPO training. ### Configuration @@ -203,10 +205,10 @@ class CustomTrajectoryMaker(Controller): def process(self, tasks, **kwargs): # Run generation yield from self.generation_controller.process(tasks, **kwargs) - + # Run reward computation yield from self.reward_controller.process(tasks, **kwargs) - + # Build trajectories trajectories = [] for task in tasks: From 8cc49d7882be5198a6fec1b281cba9c74f56e369 Mon Sep 17 00:00:00 2001 From: Fred Wei <20514172+WeiHaocheng@users.noreply.github.com> Date: Mon, 23 Feb 2026 09:28:54 +0000 Subject: [PATCH 05/18] refactor workflow draft --- areal/experimental/scaffolding/workflow.py | 38 ++++++++++++++----- .../scaffolding/gsm8k_rlvr_scaffolding.py | 33 +++++++++++++++- 2 files changed, 60 insertions(+), 11 deletions(-) diff --git a/areal/experimental/scaffolding/workflow.py b/areal/experimental/scaffolding/workflow.py index 16d1abf554..5a08f42586 100644 --- a/areal/experimental/scaffolding/workflow.py +++ b/areal/experimental/scaffolding/workflow.py @@ -19,10 +19,10 @@ import torch from transformers import PreTrainedTokenizerFast +from areal import workflow_context from areal.api.cli_args import GenerationHyperparameters from areal.api.engine_api import InferenceEngine from areal.api.workflow_api import RolloutWorkflow -from areal import workflow_context from areal.experimental.scaffolding._compat import ( GenerationTask, NativeGenerationController, @@ -83,12 +83,10 @@ def __init__( self.gconfig = gconfig.new_with_stop_and_pad_token_ids(self.tokenizer) self.enable_thinking = enable_thinking - # Scaffolding controllers - self.reward_controller = RLVRRewardController(self.reward_fn) - self.gen_controller = NativeGenerationController() - - # Lazily created from engine server addresses + # Lazily created from engine server addresses via build_scaffolding_llm self.worker: SGLangWorker | None = None + self.gen_controller: NativeGenerationController | None = None + self.reward_controller: RLVRRewardController | None = None self.trajectory_maker: PipelineTrajectoryMaker | None = None self.scaffolding_llm: ScaffoldingLlm | None = None @@ -106,15 +104,37 @@ def _lazy_init_scaffolding(self, engine: InferenceEngine) -> None: async_client=async_client, model="default", engine=engine ) + self.scaffolding_llm = self.build_scaffolding_llm(engine) + logger.info(f"Initialized scaffolding components with server at {addr}") + + def build_scaffolding_llm(self, engine: InferenceEngine) -> ScaffoldingLlm: + """Build the ScaffoldingLlm instance. + + Override this method in subclasses to use different scaffolding + controllers or worker configurations. Subclasses should set + ``self.gen_controller`` and ``self.reward_controller`` here. + + When this method is called, ``self.worker`` is already initialized. + + Parameters + ---------- + engine : InferenceEngine + The inference engine (available for address lookup if needed). + + Returns + ------- + ScaffoldingLlm + The constructed ScaffoldingLlm instance. + """ + self.gen_controller = NativeGenerationController() + self.reward_controller = RLVRRewardController(self.reward_fn) self.trajectory_maker = PipelineTrajectoryMaker( self.gen_controller, self.reward_controller ) - - self.scaffolding_llm = ScaffoldingLlm( + return ScaffoldingLlm( self.trajectory_maker, {NativeGenerationController.WorkerTag.GENERATION: self.worker}, ) - logger.info(f"Initialized scaffolding components with server at {addr}") async def _generate_via_worker( self, prompt_str: str, input_ids: list[int] diff --git a/examples/scaffolding/gsm8k_rlvr_scaffolding.py b/examples/scaffolding/gsm8k_rlvr_scaffolding.py index 529cdf4dfd..41dda7da56 100644 --- a/examples/scaffolding/gsm8k_rlvr_scaffolding.py +++ b/examples/scaffolding/gsm8k_rlvr_scaffolding.py @@ -14,11 +14,40 @@ import sys from areal.api.cli_args import GRPOConfig, load_expr_config +from areal.api.engine_api import InferenceEngine from areal.dataset import get_custom_dataset +from areal.experimental.scaffolding._compat import ( + NativeGenerationController, + ScaffoldingLlm, +) +from areal.experimental.scaffolding.controllers import ( + PipelineTrajectoryMaker, + RLVRRewardController, +) +from areal.experimental.scaffolding.workflow import ScaffoldingWorkflow from areal.trainer import PPOTrainer from areal.utils.hf_utils import load_hf_tokenizer +class GSM8KScaffoldingWorkflow(ScaffoldingWorkflow): + """ScaffoldingWorkflow customized for GSM8K RLVR training. + + Demonstrates overriding ``build_scaffolding_llm`` to control how the + ScaffoldingLlm is constructed (e.g., swap controllers or workers). + """ + + def build_scaffolding_llm(self, engine: InferenceEngine) -> ScaffoldingLlm: + self.gen_controller = NativeGenerationController() + self.reward_controller = RLVRRewardController(self.reward_fn) + self.trajectory_maker = PipelineTrajectoryMaker( + self.gen_controller, self.reward_controller + ) + return ScaffoldingLlm( + self.trajectory_maker, + {NativeGenerationController.WorkerTag.GENERATION: self.worker}, + ) + + def main(args): """Main entry point for RLVR training with scaffolding.""" config, _ = load_expr_config(args, GRPOConfig) @@ -50,9 +79,9 @@ def main(args): valid_dataset=valid_dataset, ) as trainer: trainer.train( - workflow="areal.experimental.scaffolding.workflow.ScaffoldingWorkflow", + workflow=GSM8KScaffoldingWorkflow, workflow_kwargs=workflow_kwargs, - eval_workflow="areal.experimental.scaffolding.workflow.ScaffoldingWorkflow", + eval_workflow=GSM8KScaffoldingWorkflow, eval_workflow_kwargs=eval_workflow_kwargs, ) From f1910c4e0a1eaf6ce0a77baa41284d802cbbfe4e Mon Sep 17 00:00:00 2001 From: Fred Wei <20514172+WeiHaocheng@users.noreply.github.com> Date: Tue, 24 Feb 2026 15:16:36 +0000 Subject: [PATCH 06/18] add scaffolding core --- areal/experimental/scaffolding/__init__.py | 9 +- areal/experimental/scaffolding/_compat.py | 215 +---- areal/experimental/scaffolding/controllers.py | 10 +- .../experimental/scaffolding/core/__init__.py | 63 ++ .../scaffolding/core/controller.py | 207 +++++ .../scaffolding/core/math_utils.py | 89 ++ areal/experimental/scaffolding/core/result.py | 67 ++ .../scaffolding/core/scaffolding_llm.py | 234 ++++++ areal/experimental/scaffolding/core/task.py | 270 ++++++ .../scaffolding/core/task_collection.py | 492 +++++++++++ areal/experimental/scaffolding/core/worker.py | 171 ++++ areal/experimental/scaffolding/task.py | 2 +- areal/experimental/scaffolding/worker.py | 6 +- .../experimental/scaffolding_core/__init__.py | 0 .../scaffolding_core/test_self_contained.py | 786 ++++++++++++++++++ examples/scaffolding/README.md | 10 +- 16 files changed, 2440 insertions(+), 191 deletions(-) create mode 100644 areal/experimental/scaffolding/core/__init__.py create mode 100644 areal/experimental/scaffolding/core/controller.py create mode 100644 areal/experimental/scaffolding/core/math_utils.py create mode 100644 areal/experimental/scaffolding/core/result.py create mode 100644 areal/experimental/scaffolding/core/scaffolding_llm.py create mode 100644 areal/experimental/scaffolding/core/task.py create mode 100644 areal/experimental/scaffolding/core/task_collection.py create mode 100644 areal/experimental/scaffolding/core/worker.py create mode 100644 areal/tests/experimental/scaffolding_core/__init__.py create mode 100644 areal/tests/experimental/scaffolding_core/test_self_contained.py diff --git a/areal/experimental/scaffolding/__init__.py b/areal/experimental/scaffolding/__init__.py index 24cda41f7e..5e0246c8f2 100644 --- a/areal/experimental/scaffolding/__init__.py +++ b/areal/experimental/scaffolding/__init__.py @@ -1,8 +1,9 @@ """ Scaffolding Framework Integration for AReaL. -This module provides integration between TensorRT-LLM's Scaffolding framework -and AReaL's RL training pipeline. +This module provides the Scaffolding framework for composing inference-time +compute methods with AReaL's RL training pipeline. Core scaffolding primitives +are vendored from TensorRT-LLM under ``areal.experimental.scaffolding.core``. Key Components: - ScaffoldingWorkflow: RolloutWorkflow implementation that wraps ScaffoldingLlm @@ -15,11 +16,8 @@ - ChatRewardTask: Task for computing rewards on traced interactions - CreateWorkerFromEngine: Creates a scaffolding Worker from AReaL's InferenceEngine - SGLangWorker: Worker implementation for SGLang engines - -Note: Requires tensorrt_llm to be installed for full functionality. """ -from areal.experimental.scaffolding._compat import HAS_TENSORRT_LLM from areal.experimental.scaffolding.controllers import ( ChatTracer, PipelineTrajectoryMaker, @@ -35,7 +33,6 @@ from areal.experimental.scaffolding.workflow import ScaffoldingWorkflow __all__ = [ - "HAS_TENSORRT_LLM", "ScaffoldingWorkflow", "RLVRRewardTask", "RLVRRewardController", diff --git a/areal/experimental/scaffolding/_compat.py b/areal/experimental/scaffolding/_compat.py index 3b845056ee..d63112755c 100644 --- a/areal/experimental/scaffolding/_compat.py +++ b/areal/experimental/scaffolding/_compat.py @@ -1,192 +1,61 @@ -"""Compatibility layer for optional tensorrt_llm.scaffolding dependency. +"""Scaffolding framework primitives vendored from TensorRT-LLM. -Provides imports from tensorrt_llm.scaffolding when available, or lightweight -standalone implementations when not installed. +This module re-exports the core scaffolding classes from the vendored copy +at ``areal.experimental.scaffolding.core``, so the rest of AReaL can import +them from a single location. """ -from __future__ import annotations - -import enum -from dataclasses import dataclass, field -from typing import Any - -try: - from tensorrt_llm.scaffolding import ( - NativeGenerationController, - ScaffoldingLlm, - ) - from tensorrt_llm.scaffolding.controller import Controller - from tensorrt_llm.scaffolding.result import ScaffoldingOutput - from tensorrt_llm.scaffolding.task import ( - AssistantMessage, - ChatTask, - GenerationTask, - Task, - TaskStatus, - ) - from tensorrt_llm.scaffolding.task_collection import ( - TaskCollection, - with_task_collection, - ) - from tensorrt_llm.scaffolding.worker import OpenaiWorker, Worker - - HAS_TENSORRT_LLM = True - -except ImportError: - HAS_TENSORRT_LLM = False - - # ---- Standalone lightweight implementations ---- - # These provide the scaffolding interfaces so the framework works - # without tensorrt_llm installed. - - class Controller: - """Lightweight Controller base class.""" - - def process(self, tasks: list, **kwargs) -> Any: - yield tasks - - @dataclass - class Task: - """Lightweight Task base class.""" - - worker_tag: Any = None - - @dataclass - class GenerationTask(Task): - """Lightweight GenerationTask.""" - - input_str: str | None = None - output_str: str | None = None - input_tokens: list | None = None - output_tokens: list | None = None - logprobs: Any = None - finish_reason: str | None = None - perf_metrics: Any = None - customized_result_fields: dict = field(default_factory=dict) - - @dataclass - class ChatTask(Task): - """Lightweight ChatTask.""" - - messages: list = field(default_factory=list) - completion: Any = None - tools: list | None = None - finish_reason: str | None = None - input_tokens: list | None = None - output_tokens: list | None = None - enable_token_counting: bool = False - prompt_tokens_num: int = 0 - completion_tokens_num: int = 0 - reasoning_tokens_num: int = 0 - perf_metrics: Any = None - - @staticmethod - def create_from_prompt(prompt: str) -> ChatTask: - return ChatTask(messages=[{"role": "user", "content": prompt}]) - - def messages_to_dict_content(self) -> list: - return self.messages - - class TaskStatus(enum.Enum): - """Lightweight TaskStatus.""" - - SUCCESS = "success" - WORKER_EXECEPTION = "worker_exception" # noqa: S105 (matches upstream typo) - - class AssistantMessage: - """Lightweight AssistantMessage.""" - - def __init__( - self, - content: str | None = None, - reasoning: str | None = None, - reasoning_content: str | None = None, - tool_calls: list | None = None, - ): - self.content = content - self.reasoning = reasoning - self.reasoning_content = reasoning_content - self.tool_calls = tool_calls - - class TaskCollection: - """Lightweight TaskCollection base class.""" - - def before_yield(self, tasks: list) -> None: - pass - - def after_yield(self, tasks: list) -> None: - pass - - def with_task_collection(name: str, collection_cls: type): - """Decorator that attaches a TaskCollection to a Controller class.""" - - def decorator(cls): - if not hasattr(cls, "task_collections"): - cls.task_collections = {} - cls.task_collections[name] = collection_cls() - return cls - - return decorator - - class Worker: - """Lightweight Worker base class.""" - - class OpenaiWorker(Worker): - """Lightweight OpenaiWorker base class.""" - - def __init__(self, async_client: Any = None, model: str = "", **kwargs): - self.async_client = async_client - self.model = model - - def convert_task_params(self, task: Any) -> dict: - return {} - - @dataclass - class ScaffoldingOutput: - """Lightweight ScaffoldingOutput.""" - - text: str = "" - token_ids: list = field(default_factory=list) - - class NativeGenerationController(Controller): - """Lightweight NativeGenerationController.""" - - class WorkerTag(enum.Enum): - GENERATION = "generation" - - def process(self, tasks: list, **kwargs) -> Any: - yield tasks - - class ScaffoldingLlm: - """Lightweight ScaffoldingLlm.""" - - def __init__(self, controller: Controller, workers: dict | None = None): - self.controller = controller - self.workers = workers or {} - - def generate(self, prompt: str, **kwargs) -> Any: - return None - - async def generate_async(self, prompt: str, **kwargs) -> Any: - return None - - def shutdown(self) -> None: - pass - +from areal.experimental.scaffolding.core.controller import ( + BestOfNController, + Controller, + MajorityVoteController, + NativeChatController, + NativeGenerationController, + NativeRewardController, + ParallelProcess, +) +from areal.experimental.scaffolding.core.result import ScaffoldingOutput +from areal.experimental.scaffolding.core.scaffolding_llm import ScaffoldingLlm +from areal.experimental.scaffolding.core.task import ( + AssistantMessage, + ChatTask, + GenerationTask, + OpenAIToolDescription, + RoleMessage, + StreamGenerationTask, + SystemMessage, + Task, + TaskStatus, + UserMessage, +) +from areal.experimental.scaffolding.core.task_collection import ( + TaskCollection, + with_task_collection, +) +from areal.experimental.scaffolding.core.worker import OpenaiWorker, Worker __all__ = [ - "HAS_TENSORRT_LLM", "AssistantMessage", + "BestOfNController", "ChatTask", "Controller", "GenerationTask", + "MajorityVoteController", + "NativeChatController", "NativeGenerationController", + "NativeRewardController", + "OpenAIToolDescription", "OpenaiWorker", + "ParallelProcess", + "RoleMessage", "ScaffoldingLlm", "ScaffoldingOutput", + "StreamGenerationTask", + "SystemMessage", "Task", "TaskCollection", "TaskStatus", + "UserMessage", "Worker", "with_task_collection", ] diff --git a/areal/experimental/scaffolding/controllers.py b/areal/experimental/scaffolding/controllers.py index e8b68535fa..3f04ea7615 100644 --- a/areal/experimental/scaffolding/controllers.py +++ b/areal/experimental/scaffolding/controllers.py @@ -2,7 +2,7 @@ RLVR Controllers for Scaffolding Framework. This module provides controllers for RLVR (Reinforcement Learning with Verifiable -Rewards) that integrate with TensorRT-LLM's scaffolding framework. +Rewards) that integrate with the scaffolding framework. Key Components: - RLVRRewardController: Controller that processes reward computation @@ -307,6 +307,10 @@ def after_yield(self, tasks: list[Task]): if not isinstance(task, ChatTask): continue + # Skip tasks without a completion (e.g. not yet processed by worker) + if not hasattr(task, "completion") or task.completion is None: + continue + interaction = self._create_interaction_from_chat_task(task) # Use completion.id as the interaction key completion_id = task.completion.id @@ -411,7 +415,7 @@ class PipelineTrajectoryMaker(Controller): Example ------- ```python - from tensorrt_llm.scaffolding import NativeGenerationController + from areal.experimental.scaffolding._compat import NativeGenerationController gen_controller = NativeGenerationController() reward_controller = RLVRRewardController(gsm8k_reward_fn) @@ -560,7 +564,7 @@ class TraceTrajectoryMaker(Controller): Example ------- ```python - from tensorrt_llm.scaffolding import NativeGenerationController + from areal.experimental.scaffolding._compat import NativeGenerationController chat_controller = SomeChatController() reward_controller = RLVRRewardController(gsm8k_reward_fn) diff --git a/areal/experimental/scaffolding/core/__init__.py b/areal/experimental/scaffolding/core/__init__.py new file mode 100644 index 0000000000..e821f00a76 --- /dev/null +++ b/areal/experimental/scaffolding/core/__init__.py @@ -0,0 +1,63 @@ +# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +# Vendored from TensorRT-LLM scaffolding framework. +# Core scaffolding primitives adapted for standalone use in AReaL. + +from .controller import ( + BestOfNController, + Controller, + MajorityVoteController, + NativeChatController, + NativeGenerationController, + NativeRewardController, + ParallelProcess, +) +from .math_utils import ( + extract_answer_from_boxed, + extract_answer_with_regex, + get_digit_majority_vote_result, +) +from .result import ScaffoldingOutput +from .scaffolding_llm import ScaffoldingLlm +from .task import ( + AssistantMessage, + ChatTask, + GenerationTask, + OpenAIToolDescription, + RoleMessage, + StreamGenerationTask, + SystemMessage, + Task, + TaskStatus, + UserMessage, +) +from .task_collection import TaskCollection, with_task_collection +from .worker import OpenaiWorker, Worker + +__all__ = [ + "ScaffoldingLlm", + "ParallelProcess", + "Controller", + "NativeChatController", + "NativeGenerationController", + "NativeRewardController", + "MajorityVoteController", + "BestOfNController", + "Task", + "GenerationTask", + "StreamGenerationTask", + "ChatTask", + "OpenAIToolDescription", + "RoleMessage", + "UserMessage", + "SystemMessage", + "AssistantMessage", + "Worker", + "OpenaiWorker", + "TaskStatus", + "extract_answer_from_boxed", + "extract_answer_with_regex", + "get_digit_majority_vote_result", + "TaskCollection", + "with_task_collection", + "ScaffoldingOutput", +] diff --git a/areal/experimental/scaffolding/core/controller.py b/areal/experimental/scaffolding/core/controller.py new file mode 100644 index 0000000000..77fdc18f99 --- /dev/null +++ b/areal/experimental/scaffolding/core/controller.py @@ -0,0 +1,207 @@ +# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +# Vendored from tensorrt_llm.scaffolding.controller + +import copy +import logging +from abc import ABC +from collections.abc import Mapping +from enum import Enum +from typing import Any + +import torch + +from .math_utils import get_digit_majority_vote_result +from .task import ChatTask, GenerationTask, Task + +logger = logging.getLogger(__name__) + + +class Controller(ABC): + task_collections: dict = {} + + def __init__(self): + self.task_collections = {} + + def clone(self): + return copy.deepcopy(self) + + def generate(self, prompt: str, **kwargs): + task = GenerationTask.create_from_prompt(prompt) + + yield from self.process([task], **kwargs) + + return task.create_scaffolding_output() + + def process(self, tasks: list[Task], **kwargs): + raise NotImplementedError + + +class ParallelProcess: + def __init__( + self, + controllers: list[Controller], + tasks_list: list[list[Task]], + kwargs_list: list[Mapping[str, Any]], + ): + self.sub_gens = [] + for controller, tasks, kwargs in zip(controllers, tasks_list, kwargs_list): + gen = controller.process(tasks, **kwargs) + self.sub_gens.append(gen) + + +# Controller runs multiple generation tasks. +class NativeGenerationController(Controller): + class WorkerTag(Enum): + GENERATION = "generation" + + def __init__(self, sampling_params: dict = None, streaming: bool = False): + super().__init__() + if sampling_params is None: + sampling_params = {} + for key, value in list(sampling_params.items()): + if key not in GenerationTask.__annotations__: + logger.warning(f"{key} is not a supported field for GenerationTask") + sampling_params.pop(key) + self.sampling_params = sampling_params + self.streaming = streaming + + # [GenerationTask] -> [GenerationTask] | [ChatTask] -> [ChatTask] + def process(self, tasks: list[Task], **kwargs): + for task in tasks: + task.worker_tag = self.WorkerTag.GENERATION + for key, value in self.sampling_params.items(): + if getattr(task, key) is None: + setattr(task, key, value) + + task.streaming_output_flag = self.streaming + + yield tasks + + +class NativeChatController(NativeGenerationController): + def __init__(self, sampling_params: dict = None, streaming: bool = False): + super().__init__(sampling_params, streaming) + + def process(self, tasks: list[Task], **kwargs): + chat_tasks = [ChatTask.create_from_prompt(task.input_str) for task in tasks] + yield from super().process(chat_tasks, **kwargs) + + +class NativeRewardController(Controller): + def __init__(self): + self.scores = None + + class WorkerTag(Enum): + REWARD = "reward" + + def process(self, tasks: list[Task], **kwargs): + task = GenerationTask() + for task in tasks: + task.worker_tag = self.WorkerTag.REWARD + + yield tasks + + +class MajorityVoteController(Controller): + def __init__(self, generation_controller: Controller, default_sample_num: int = 1): + super().__init__() + self.generation_controller = generation_controller + self.default_sample_num = default_sample_num + + def clone(self): + generation_controller = self.generation_controller.clone() + return MajorityVoteController(generation_controller, self.default_sample_num) + + def process( + self, + tasks: list[Task], + sample_num: int = 1, + generation_kwargs: dict = {}, + majority_vote_kwargs: dict = {}, + ): + sample_num = max(sample_num, self.default_sample_num) + generation_controllers = [ + self.generation_controller.clone() for _ in range(sample_num) + ] + tasks_list = [copy.deepcopy(tasks) for _ in range(sample_num)] + generation_kwargs_list = [ + copy.deepcopy(generation_kwargs) for _ in range(sample_num) + ] + + yield ParallelProcess( + generation_controllers, tasks_list, generation_kwargs_list + ) + + majority_index, majority_answer = self.majority_vote( + tasks_list, **majority_vote_kwargs + ) + + assert isinstance(majority_answer, str), "majority_vote failed" + tasks[0].result = tasks_list[majority_index][0].result + + def majority_vote( + self, candidates_tasks: list[list[Task]], **kwargs + ) -> tuple[int, str]: + candidates = [tasks[0].output_str for tasks in candidates_tasks] + return get_digit_majority_vote_result(candidates) + + +class BestOfNController(Controller): + def __init__( + self, + generation_controller: Controller, + reward_controller: Controller, + default_sample_num: int = 4, + ): + super().__init__() + self.generation_controller = generation_controller + self.reward_controller = reward_controller + self.default_sample_num = default_sample_num + + def clone(self): + generation_controller = self.generation_controller.clone() + reward_controller = self.reward_controller.clone() + return BestOfNController( + generation_controller, reward_controller, self.default_sample_num + ) + + def process( + self, + tasks: list[Task], + sample_num: int = 4, + generation_kwargs: dict = {}, + reward_kwargs: dict = {}, + select_best_kwargs: dict = {}, + ): + assert len(tasks) == 1, "BestOfNController only supports one task" + task = tasks[0] + + sample_num = max(sample_num, self.default_sample_num) + generation_controllers = [self.generation_controller for _ in range(sample_num)] + generation_kwargs_list = [generation_kwargs for _ in range(sample_num)] + generation_tasks = [copy.deepcopy(task) for _ in range(sample_num)] + + yield ParallelProcess( + generation_controllers, + [[t] for t in generation_tasks], + generation_kwargs_list, + ) + + yield from self.reward_controller.process(generation_tasks, **reward_kwargs) + + assert self.reward_controller.scores is not None + reward_values = self.reward_controller.scores + + for i, gen_task, reward_value in zip( + range(sample_num), generation_tasks, reward_values + ): + logger.info(f"[output {i}, score {reward_value}]:\n{gen_task.output_str}") + + best_task, best_idx = self.select_best( + generation_tasks, reward_values, **select_best_kwargs + ) + task.result = best_task.result + + def select_best(self, tasks: list[Task], reward_values, **kwargs) -> Task: + max_index = torch.argmax(torch.tensor(reward_values)).item() + return tasks[max_index], max_index diff --git a/areal/experimental/scaffolding/core/math_utils.py b/areal/experimental/scaffolding/core/math_utils.py new file mode 100644 index 0000000000..e94319bacb --- /dev/null +++ b/areal/experimental/scaffolding/core/math_utils.py @@ -0,0 +1,89 @@ +# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +# Vendored from tensorrt_llm.scaffolding.math_utils + +import re + + +def extract_answer_with_regex( + string: str, extract_regex: str = r"The final answer is (.+)$" +): + match = re.search(extract_regex, string) + if match: + return match.group(1) + return None + + +def extract_answer_from_boxed(string: str): + """Extract Answer String from \\boxed expression or based on regex""" + + if "\\boxed" not in string: + return None + + idx = string.rfind("\\boxed") + if idx < 0: + idx = string.rfind("\\fbox") + if idx < 0: + return None + + i = idx + right_brace_idx = None + num_left_braces_open = 0 + while i < len(string): + if string[i] == "{": + num_left_braces_open += 1 + if string[i] == "}": + num_left_braces_open -= 1 + if num_left_braces_open == 0: + right_brace_idx = i + break + i += 1 + + if right_brace_idx is None: + retval = None + else: + retval = string[idx : right_brace_idx + 1] + + if retval: + left = "\\boxed{" + try: + assert retval[: len(left)] == left + assert retval[-1] == "}" + return retval[len(left) : -1] + except AssertionError: + return None + + return None + + +def get_majority_result( + results: list, + result_extractor=lambda x: x, + result_validator=lambda x: True, +): + extract_answers = [result_extractor(result) for result in results] + valid_answers = [ + result + for result in extract_answers + if result is not None and result_validator(result) is True + ] + if len(valid_answers) == 0: + return None, None + + answer_counts = {} + for answer in valid_answers: + answer_counts[answer] = answer_counts.get(answer, 0) + 1 + majority_answer = max(answer_counts, key=answer_counts.get) + majority_index = next( + filter(lambda x: x[1] == majority_answer, enumerate(extract_answers)) + )[0] + return majority_index, majority_answer + + +def get_digit_majority_vote_result(results: list[str]) -> str: + def is_digit(result: str): + return result.isdigit() + + index, extract_answer = get_majority_result( + results, result_extractor=extract_answer_from_boxed, result_validator=is_digit + ) + return (index, extract_answer) if extract_answer else (0, None) diff --git a/areal/experimental/scaffolding/core/result.py b/areal/experimental/scaffolding/core/result.py new file mode 100644 index 0000000000..cb0eaf6670 --- /dev/null +++ b/areal/experimental/scaffolding/core/result.py @@ -0,0 +1,67 @@ +# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +# Vendored from tensorrt_llm.scaffolding.result + +import asyncio +from collections.abc import Mapping +from dataclasses import dataclass +from typing import Any + + +@dataclass +class ScaffoldingOutput: + text: str + token_ids: list[int] + + +class ScaffoldingResult: + def __init__(self): + super().__init__() + self.aqueue = asyncio.Queue() + self.outputs = [] + # only support one output for now, so use an empty obj to init + self.outputs.append(ScaffoldingOutput("", [])) + self._done = False + self.task_collections = None + + def set_output(self, output: ScaffoldingOutput | Any): + if isinstance(output, ScaffoldingOutput): + self.set_output_streaming(output) + # terminate + self.set_output_streaming(None) + + def set_output_streaming(self, output: ScaffoldingOutput | Any): + self.aqueue.put_nowait(output) + + def set_task_collections(self, task_collections: Mapping[str, Any]): + self.task_collections = task_collections + + async def _aresult_step(self): + obj = await self.aqueue.get() + if obj is None: + self._done = True + else: # obj is ScaffoldingOutput + self.outputs[0] = obj + + def result(self, timeout: float | None = None) -> "ScaffoldingResult": + if not self._done: + loop = asyncio.get_event_loop() + asyncio.run_coroutine_threadsafe(self.aresult(), loop).result() + return self + + async def aresult(self) -> "ScaffoldingResult": + while not self._done: + await self._aresult_step() + return self + + def __await__(self): + return self.aresult().__await__() + + def __aiter__(self): + return self + + async def __anext__(self): + if self._done: + raise StopAsyncIteration + + await self._aresult_step() + return self diff --git a/areal/experimental/scaffolding/core/scaffolding_llm.py b/areal/experimental/scaffolding/core/scaffolding_llm.py new file mode 100644 index 0000000000..4e6018e145 --- /dev/null +++ b/areal/experimental/scaffolding/core/scaffolding_llm.py @@ -0,0 +1,234 @@ +# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +# Vendored from tensorrt_llm.scaffolding.scaffolding_llm + +import asyncio +import threading +import traceback +from collections import deque +from collections.abc import Generator, Mapping +from dataclasses import dataclass +from typing import Any + +from .controller import Controller, ParallelProcess +from .result import ScaffoldingResult +from .task import Task +from .worker import Worker + + +@dataclass(frozen=True) +class ScaffoldingRequest: + prompt: str + kwargs: Mapping[str, Any] + controller: Controller + result: "ScaffoldingResult" + + +class ScaffoldingLlm: + def __init__( + self, + prototype_controller: Controller, + workers: Mapping[str, Worker], # map of role to worker instance, + max_parallel_requests: int = 64, + ): + self.prototype_controller = prototype_controller + self.workers = workers + + self.loop = self._get_loop() + asyncio.set_event_loop(self.loop) + self.task_queue = asyncio.Queue() + self.main_loop_stop_event = asyncio.Event() + self.shutdown_event = asyncio.Event() + if self.own_loop: + self._run_main_loop_thread() + else: + self._run_main_loop_coroutine() + + # For top scheduler + self.running_req_count = 0 + self.max_parallel_requests = max_parallel_requests + self.pending_queue = deque() + + self.output_task_collection = False + + def __enter__(self): + return self + + def __exit__(self): + self.shutdown() + + def _get_loop(self): + try: + self.own_loop = False + return asyncio.get_running_loop() + except RuntimeError: + self.own_loop = True + return asyncio.new_event_loop() + return None + + async def _handle_controller_generator( + self, gen: Generator, request: ScaffoldingRequest = None + ): + """Handle a controller generator, processing tasks and parallel processes.""" + for obj in gen: + if isinstance(obj, ParallelProcess): + await self._handle_parallel_process(obj, request) + else: + await self._handle_task_list(obj, request) + + async def _handle_task_list( + self, tasks: list[Task], request: ScaffoldingRequest = None + ): + """Execute a list of tasks concurrently.""" + async_tasks = [ + asyncio.create_task(self.workers[task.worker_tag].run_task(task)) + for task in tasks + ] + await asyncio.gather(*async_tasks) + for task in tasks: + if task.streaming_output_flag: + for output in task.streaming_output_list: + request.result.set_output_streaming(output) + task.streaming_output_list = [] + + async def _handle_parallel_process( + self, tasks: ParallelProcess, request: ScaffoldingRequest = None + ): + """Handle parallel execution of multiple generators.""" + async_tasks = [ + asyncio.create_task(self._handle_controller_generator(sub_gen, request)) + for sub_gen in tasks.sub_gens + ] + await asyncio.gather(*async_tasks) + + async def _handle_single_request(self, request: ScaffoldingRequest): + """Process a single scaffolding request.""" + try: + gen = self._create_controller_generator(request) + await self._handle_controller_generator(gen, request) + except Exception as e: + print(f"ScaffoldingLLM request exception: {e}") + traceback.print_exc() + request.result.set_output(None) + raise + finally: + self.running_req_count -= 1 + self._maybe_schedule() + + def _create_controller_generator(self, request: ScaffoldingRequest): + """Create a generator wrapper for the controller.""" + scaffolding_output = yield from request.controller.generate( + request.prompt, **request.kwargs + ) + + if self.output_task_collection: + request.result.set_task_collections(request.controller.task_collections) + request.result.set_output(scaffolding_output) + + def _schedule_request(self, request: ScaffoldingRequest): + """Schedule a single request for execution.""" + asyncio.create_task(self._handle_single_request(request)) + self.running_req_count += 1 + + def _maybe_schedule(self, request: ScaffoldingRequest = None): + """Schedule pending requests if capacity allows.""" + if self.shutdown_event.is_set(): + return + + if request is not None: + self.pending_queue.append(request) + + while ( + self.running_req_count < self.max_parallel_requests and self.pending_queue + ): + next_request = self.pending_queue.popleft() + self._schedule_request(next_request) + + async def _handle_event_loop(self): + """Main event handling loop.""" + while True: + item = await self.task_queue.get() + + if item is None: + return + elif isinstance(item, ScaffoldingRequest): + self._maybe_schedule(item) + else: + raise ValueError(f"Unsupported task_queue item type: {type(item)}") + + async def _main_loop_async_func(self): + """Main async loop function.""" + handle_event_task = asyncio.create_task(self._handle_event_loop()) + await handle_event_task + self.main_loop_stop_event.set() + + def _run_main_loop_coroutine(self): + asyncio.run_coroutine_threadsafe(self._main_loop_async_func(), self.loop) + + def _run_main_loop_thread(self): + def main_loop_thread(): + self.loop.run_until_complete(self._main_loop_async_func()) + + self.main_loop_thread = threading.Thread(target=main_loop_thread) + self.main_loop_thread.start() + + def generate_async(self, prompt: str) -> ScaffoldingResult: + result = ScaffoldingResult() + + async def put_request(): + try: + request = ScaffoldingRequest( + prompt=prompt, + kwargs={}, + result=result, + controller=self.prototype_controller.clone(), + ) + except Exception as e: + self.task_queue.put(None) + print( + f"Error: build ScaffoldingRequest failed: {e} \n {traceback.format_exc()}" + ) + else: + await self.task_queue.put(request) + + asyncio.run_coroutine_threadsafe(put_request(), self.loop) + + return result + + def generate( + self, prompts: str | list[str] + ) -> ScaffoldingResult | list[ScaffoldingResult]: + unbatched = not isinstance(prompts, list) + batched_prompts = [prompts] if unbatched else prompts + + scaffolding_results = [] + for prompt in batched_prompts: + scaffolding_results.append(self.generate_async(prompt)) + + for scaffolding_result in scaffolding_results: + scaffolding_result.result() + + return scaffolding_results[0] if unbatched else scaffolding_results + + def enable_output_task_collection(self): + self.output_task_collection = True + + def shutdown(self, shutdown_workers=False): + def shutdown_workers_func(): + for worker in self.workers.values(): + worker.shutdown() + + async def stop_task_on_loop(): + await self.task_queue.put(None) + await self.main_loop_stop_event.wait() + for worker in self.workers.values(): + await worker.async_shutdown() + + asyncio.run_coroutine_threadsafe(stop_task_on_loop(), self.loop) + + if self.own_loop: + self.main_loop_thread.join() + else: + self.shutdown_event.set() + + if shutdown_workers: + shutdown_workers_func() diff --git a/areal/experimental/scaffolding/core/task.py b/areal/experimental/scaffolding/core/task.py new file mode 100644 index 0000000000..0e2b0c6bb4 --- /dev/null +++ b/areal/experimental/scaffolding/core/task.py @@ -0,0 +1,270 @@ +# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +# Vendored from tensorrt_llm.scaffolding.task + +import json +from collections.abc import Mapping +from dataclasses import dataclass, field +from enum import Enum +from typing import Any + +from .result import ScaffoldingOutput + + +@dataclass +class Task: + # Scaffolding delivers the task to the Worker by worker_tag. + worker_tag: str = field(default=None) + + # For streaming output. + streaming_output_flag: bool = field(default=False) + streaming_output_list: list[Any] = field(default_factory=list) + + # Reserve for custom input params. + custom_input_params: dict | None = None + + # Reserve for custom output params. + custom_output_params: dict | None = None + + @staticmethod + def create_from_prompt(prompt: str) -> "Task": + pass + + def create_scaffolding_output(self) -> ScaffoldingOutput: + pass + + def create_scaffolding_output_stream(self) -> list[ScaffoldingOutput]: + pass + + +class TaskStatus(Enum): + SUCCESS = "success" + WORKER_NOT_SUPPORTED = "worker_not_supported" + WORKER_EXECEPTION = "worker_exception" + + +@dataclass +class GenerationTask(Task): + # input field + input_tokens: list[int] | None = None + input_str: str | None = None + skip_tokenizer: bool = False + skip_detokenizer: bool = False + + # sampling params for openai + best_of: int | None = None + echo: bool | None = False + frequency_penalty: float | None = 0.0 + logit_bias: dict[str, float] | None = None + num_logprobs: int | None = None + max_tokens: int | None = None + n: int = 1 + presence_penalty: float | None = 0.0 + seed: int | None = None + stop: str | list[str] | None = field(default_factory=list) + suffix: str | None = None + temperature: float | None = None + top_p: float | None = None + user: str | None = None + ignore_eos: bool = False + + # sampling params + top_k: int | None = None + return_context_logits: bool | None = False + + # suggest to use Controller.WorkerTag + worker_tag: str | None = None + + # result field + output_str: str | None = None + output_tokens: list[int] | None = None + finish_reason: str | None = None + context_logits: Any = None + logprobs: Any = None + customized_result_fields: dict[str, Any] = field(default_factory=dict) + + perf_metrics: dict[str, float] | None = None + + @staticmethod + def create_from_prompt(prompt: str) -> "GenerationTask": + task = GenerationTask() + task.input_str = prompt + task.skip_tokenizer = False + task.skip_detokenizer = False + return task + + def create_scaffolding_output(self) -> ScaffoldingOutput: + return ScaffoldingOutput(self.output_str, self.output_tokens) + + +@dataclass +class StreamGenerationTask(GenerationTask): + # input field + cancel_flag: bool | None = field(default=False) + streaming_step: int | None = field(default=1) + + # result field + request_handle: Any = field(default=None) + end_flag: bool = field(default=False) + + @staticmethod + def create_from_generation_task( + task: GenerationTask, streaming_step + ) -> "StreamGenerationTask": + stream_task = StreamGenerationTask() + for k, v in task.__dict__.items(): + stream_task.__dict__[k] = v + stream_task.streaming_step = streaming_step + return stream_task + + +@dataclass +class RewardTask(Task): + # input field + input_tokens: list[int] | None = field(default=None) + input_str: str | None = field(default=None) + + +@dataclass +class RoleMessage: + role: str | None = field(default=None) + content: str | None = field(default=None) + prefix: str | None = field(default=None) + + def __str__(self) -> str: + return json.dumps( + { + "role": self.role, + "content": self.content, + } + ) + + def __repr__(self) -> str: + return f"{self.role}: {self.content}\n" + + def to_dict(self) -> dict[str, Any]: + return {"role": self.role, "content": self.content} + + @classmethod + def from_dict(cls, data: dict[str, Any]): + return cls(role=data["role"], content=data["content"]) + + +@dataclass +class UserMessage(RoleMessage): + def __init__(self, content: str, prefix: str | None = None): + super().__init__(role="user", content=content, prefix=prefix) + + +@dataclass +class AssistantMessage(RoleMessage): + reasoning: str | None = field(default=None) + reasoning_content: str | None = field(default=None) + tool_calls: list[Any] | None = field(default=None) + + def __init__( + self, + content: str, + reasoning: str | None = None, + reasoning_content: str | None = None, + tool_calls: list[Any] | None = None, + ): + super().__init__(role="assistant", content=content) + self.reasoning = reasoning + self.reasoning_content = reasoning_content + self.tool_calls = tool_calls + + def __str__(self) -> str: + return json.dumps( + { + "role": "assistant", + "content": self.content, + "reasoning": self.reasoning, + "reasoning_content": self.reasoning_content, + "tool_calls": [str(tool) for tool in self.tool_calls] + if self.tool_calls is not None + else None, + } + ) + + +@dataclass +class SystemMessage(RoleMessage): + def __init__(self, content: str, prefix: str | None = None): + super().__init__(role="system", content=content, prefix=prefix) + + +class ToolDescription: + def __init__(self, name: str, description: str, parameters: dict[str, Any]): + self.name = name + self.description = description + self.parameters = parameters + + def to_dict(self) -> dict[str, Any]: + pass + + +class OpenAIToolDescription(ToolDescription): + def to_dict(self) -> dict[str, Any]: + return { + "type": "function", + "function": { + "name": self.name, + "description": self.description, + "parameters": { + "type": "object", + "properties": self.parameters, + }, + }, + } + + +@dataclass +class ChatTask(StreamGenerationTask): + messages: list[RoleMessage] = field(default_factory=list) + tools: Any = field(default=None) + + # for token counting + enable_token_counting: bool = field(default=False) + prompt_tokens_num: int = field(default=0) + completion_tokens_num: int = field(default=0) + reasoning_tokens_num: int = field(default=0) + + # for sub request marker + sub_request_markers: list[tuple[str, int]] = field(default_factory=list) + unique_id: int | None = field(default=None) + + def messages_to_dict_content(self, start_index: int = 0) -> list[Mapping[str, str]]: + ret = [] + for message in self.messages[start_index:]: + if message.content is not None: + ret.append(message.to_dict()) + return ret + + def add_message(self, message: RoleMessage): + self.messages.append(message) + + def add_messages(self, messages: list[RoleMessage]): + self.messages.extend(messages) + + @staticmethod + def create_from_prompt( + user_prompt: str | None, + system_prompts: list[SystemMessage] | None = None, + tools: Any | None = None, + ) -> "ChatTask": + task = ChatTask() + if system_prompts is not None: + task.messages.extend(system_prompts) + if user_prompt is not None: + task.add_message(UserMessage(user_prompt)) + task.tools = tools + return task + + @staticmethod + def create_from_messages( + messages: list[RoleMessage], tools: Any | None = None + ) -> "ChatTask": + task = ChatTask() + task.messages = messages + task.tools = tools + return task diff --git a/areal/experimental/scaffolding/core/task_collection.py b/areal/experimental/scaffolding/core/task_collection.py new file mode 100644 index 0000000000..6ba6c303e7 --- /dev/null +++ b/areal/experimental/scaffolding/core/task_collection.py @@ -0,0 +1,492 @@ +# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +# Vendored from tensorrt_llm.scaffolding.task_collection + +import json +import time +from typing import Any + +from .controller import ParallelProcess +from .task import ChatTask, GenerationTask, Task + + +class TaskCollection: + def __init__(self): + # reserved for future use + pass + + def before_yield(self, tasks: list[Task]): + pass + + def after_yield(self, tasks: list[Task]): + pass + + @staticmethod + def get_global_info() -> Any: + pass + + +def with_task_collection( + name: str, task_collection_cls: type[TaskCollection], **task_collection_kwargs +): + def decorator(controller_cls: type): + original_init = controller_cls.__init__ + original_process = controller_cls.process + + # add task collection to controller + def new_init(self, *args, **kwargs): + original_init(self, *args, **kwargs) + self.task_collections[name] = task_collection_cls(**task_collection_kwargs) + + def new_process(self, tasks: list[Task], **kwargs): + class TaskCollectionWrapper: + def __init__(self, task_collection, gen): + self.task_collection = task_collection + self.gen = gen + + def __call__(self): + for obj in self.gen: + if isinstance(obj, ParallelProcess): + new_sub_gens = [] + for sub_gen in obj.sub_gens: + new_sub_gen = TaskCollectionWrapper( + self.task_collection, sub_gen + ) + new_sub_gens.append(new_sub_gen) + obj.sub_gens = new_sub_gens + + yield obj + else: # obj is a list of tasks + self.task_collection.before_yield(obj) + yield obj + self.task_collection.after_yield(obj) + + def __iter__(self): + return self.__call__() + + original_gen = original_process(self, tasks, **kwargs) + new_gen = TaskCollectionWrapper(self.task_collections[name], original_gen) + return new_gen() + + controller_cls.__init__ = new_init + controller_cls.process = new_process + + return controller_cls + + return decorator + + +class GenerationTokenCounter(TaskCollection): + def __init__(self): + super().__init__() + self.generation_token_count = 0 + self.pre_worker_token_sum = 0 + + def before_yield(self, tasks: list[Task]): + self.pre_worker_token_sum = 0 + for task in tasks: + if isinstance(task, GenerationTask) or issubclass( + type(task), GenerationTask + ): + if task.output_tokens: + self.pre_worker_token_sum += len(task.output_tokens) + + def after_yield(self, tasks: list[Task]): + post_worker_token_sum = 0 + for task in tasks: + if isinstance(task, GenerationTask) or issubclass( + type(task), GenerationTask + ): + if task.output_tokens: + post_worker_token_sum += len(task.output_tokens) + self.generation_token_count += post_worker_token_sum - self.pre_worker_token_sum + + +class ChatTokenCounter(TaskCollection): + # prompt tokens, completion tokens + statistics: dict[str, list[tuple[int, int]]] = {} + + def __init__(self, statistics_name: str): + super().__init__() + self.statistics_name = statistics_name + if statistics_name not in ChatTokenCounter.statistics: + ChatTokenCounter.statistics[statistics_name] = [] + + def before_yield(self, tasks: list[Task]): + for task in tasks: + if not isinstance(task, ChatTask): + continue + task.enable_token_counting = True + + def after_yield(self, tasks: list[Task]): + for task in tasks: + if not isinstance(task, ChatTask): + continue + ChatTokenCounter.statistics[self.statistics_name].append( + (task.prompt_tokens_num, task.completion_tokens_num) + ) + + def get_global_info() -> Any: + return ChatTokenCounter.statistics + + +class TaskTimer(TaskCollection): + statistics: dict[str, dict[type, list[float]]] = {} + + def __init__(self, statistics_name: str, task_types: list[type[Task]]): + super().__init__() + self.statistics_name = statistics_name + self.task_types = task_types + self.start_time_map = {} + if statistics_name not in TaskTimer.statistics: + TaskTimer.statistics[statistics_name] = {} + for task_type in task_types: + if task_type not in TaskTimer.statistics[statistics_name]: + TaskTimer.statistics[statistics_name][task_type] = [] + + def before_yield(self, tasks: list[Task]): + for task in tasks: + if type(task) not in self.task_types: + continue + + self.start_time_map[id(task)] = time.time() + + def after_yield(self, tasks: list[Task]): + for task in tasks: + if type(task) not in self.task_types: + continue + + end_time = time.time() + TaskTimer.statistics[self.statistics_name][type(task)].append( + end_time - self.start_time_map[id(task)] + ) + del self.start_time_map[id(task)] + + def get_global_info() -> Any: + return TaskTimer.statistics + + +class TaskMetricsCollector(TaskCollection): + """Task profiler that captures tasks at yield points.""" + + # Global statistics: controller_name -> List[task_info_dict] + statistics: dict[str, list[dict[str, Any]]] = {} + + def __init__( + self, + controller_name: str, + task_types: list[type[Task]] = None, + enable_print: bool = True, + capture_messages: bool = False, + ): + super().__init__() + self.controller_name = controller_name + self.task_types = task_types + self.enable_print = enable_print + self.capture_messages = capture_messages + self.start_time_map: dict[int, float] = {} + self.pre_message_count_map: dict[int, int] = {} + + if controller_name not in TaskMetricsCollector.statistics: + TaskMetricsCollector.statistics[controller_name] = [] + + def _should_process_task(self, task: Task) -> bool: + if self.task_types is not None and type(task) not in self.task_types: + return False + return True + + def _is_task_already_profiled(self, task: Task) -> bool: + return getattr(task, "_profiling_in_progress", False) + + def _mark_task_profiling_start(self, task: Task): + task._profiling_in_progress = True + + def _mark_task_profiling_end(self, task: Task): + task._profiling_in_progress = False + + def before_yield(self, tasks: list[Task]): + for task in tasks: + if not self._should_process_task(task): + continue + if self._is_task_already_profiled(task): + continue + + self._mark_task_profiling_start(task) + task_id = id(task) + self.start_time_map[task_id] = time.time() + + if isinstance(task, ChatTask): + task.enable_token_counting = True + if self.capture_messages: + self.pre_message_count_map[task_id] = len(task.messages) + + def after_yield(self, tasks: list[Task]): + for task in tasks: + task_id = id(task) + if task_id not in self.start_time_map: + continue + + end_time = time.time() + duration = end_time - self.start_time_map[task_id] + del self.start_time_map[task_id] + self._mark_task_profiling_end(task) + + task_info = { + "controller": self.controller_name, + "task_type": type(task).__name__, + "duration_ms": duration * 1000, + "timestamp": end_time, + } + + if isinstance(task, ChatTask): + task_info["prompt_tokens"] = getattr(task, "prompt_tokens_num", 0) + task_info["completion_tokens"] = getattr( + task, "completion_tokens_num", 0 + ) + task_info["reasoning_tokens"] = getattr(task, "reasoning_tokens_num", 0) + task_info["total_tokens"] = ( + task_info["prompt_tokens"] + task_info["completion_tokens"] + ) + task_info["finish_reason"] = getattr(task, "finish_reason", None) + task_info["unique_id"] = getattr(task, "unique_id", None) + task_info["sub_request_markers"] = getattr( + task, "sub_request_markers", [] + ) + task_info["perf_metrics"] = getattr(task, "perf_metrics", None) + + if self.capture_messages: + pre_message_count = self.pre_message_count_map.get(task_id, 0) + if task_id in self.pre_message_count_map: + del self.pre_message_count_map[task_id] + + task_info["message_count_before"] = pre_message_count + task_info["message_count_after"] = len(task.messages) + task_info["messages"] = [ + self._serialize_message(msg) for msg in task.messages + ] + if len(task.messages) > pre_message_count: + task_info["new_messages"] = [ + self._serialize_message(msg) + for msg in task.messages[pre_message_count:] + ] + else: + task_info["new_messages"] = [] + + TaskMetricsCollector.statistics[self.controller_name].append(task_info) + + if self.enable_print: + self._print_task_info(task_info) + + def _serialize_message(self, message) -> dict[str, Any]: + """Serialize a RoleMessage to a dictionary.""" + result = { + "role": getattr(message, "role", None), + "content": getattr(message, "content", None), + } + if hasattr(message, "reasoning") and message.reasoning is not None: + result["reasoning"] = message.reasoning + if ( + hasattr(message, "reasoning_content") + and message.reasoning_content is not None + ): + result["reasoning_content"] = message.reasoning_content + if hasattr(message, "tool_calls") and message.tool_calls is not None: + result["tool_calls"] = [str(tc) for tc in message.tool_calls] + return result + + def _print_task_info(self, task_info: dict[str, Any]): + log_parts = [ + f"[{task_info['controller']}]", + f"{task_info['task_type']}", + f"duration={task_info['duration_ms']:.2f}ms", + ] + + if "prompt_tokens" in task_info: + log_parts.append( + f"prompt={task_info['prompt_tokens']} " + f"completion={task_info['completion_tokens']} " + f"reasoning={task_info['reasoning_tokens']} " + f"total={task_info['total_tokens']}" + ) + + if task_info.get("perf_metrics"): + perf_str = ", ".join( + f"{k}={v:.2f}" if isinstance(v, float) else f"{k}={v}" + for k, v in task_info["perf_metrics"].items() + ) + log_parts.append(f"perf: {perf_str}") + + print(" | ".join(log_parts)) + + if "new_messages" in task_info and task_info["new_messages"]: + print( + f" Messages: {task_info['message_count_before']} -> {task_info['message_count_after']}" + ) + print(" New Messages:") + for msg in task_info["new_messages"]: + role = msg.get("role", "unknown") + content = msg.get("content", "") + if content and len(content) > 200: + content = content[:200] + "..." + print(f" [{role}]: {content}") + + @staticmethod + def _compute_stats(values: list[float]) -> dict[str, float]: + """Compute avg, median, min, max, sum for a list of values.""" + if not values: + return {"avg": 0, "median": 0, "min": 0, "max": 0, "sum": 0} + sorted_vals = sorted(values) + n = len(sorted_vals) + median = ( + sorted_vals[n // 2] + if n % 2 == 1 + else (sorted_vals[n // 2 - 1] + sorted_vals[n // 2]) / 2 + ) + return { + "avg": sum(values) / n, + "median": median, + "min": min(values), + "max": max(values), + "sum": sum(values), + } + + @staticmethod + def print_summary(): + """Print summary statistics for all controllers.""" + print("\n" + "=" * 80) + print("TASK METRICS SUMMARY") + print("=" * 80) + + for controller_name, task_list in TaskMetricsCollector.statistics.items(): + if not task_list: + continue + + print(f"\n{controller_name} ({len(task_list)} records)") + print("-" * 70) + + task_type_data: dict[str, dict[str, list[float]]] = {} + perf_metrics_agg: dict[str, dict[str, list[float]]] = {} + + for task_info in task_list: + task_type = task_info["task_type"] + if task_type not in task_type_data: + task_type_data[task_type] = { + "duration_ms": [], + "prompt_tokens": [], + "completion_tokens": [], + "reasoning_tokens": [], + "total_tokens": [], + } + perf_metrics_agg[task_type] = {} + + data = task_type_data[task_type] + data["duration_ms"].append(task_info["duration_ms"]) + data["prompt_tokens"].append(task_info.get("prompt_tokens", 0)) + data["completion_tokens"].append(task_info.get("completion_tokens", 0)) + data["reasoning_tokens"].append(task_info.get("reasoning_tokens", 0)) + data["total_tokens"].append(task_info.get("total_tokens", 0)) + + if task_info.get("perf_metrics"): + for key, value in task_info["perf_metrics"].items(): + if isinstance(value, (int, float)): + if key not in perf_metrics_agg[task_type]: + perf_metrics_agg[task_type][key] = [] + perf_metrics_agg[task_type][key].append(float(value)) + + for task_type, data in task_type_data.items(): + count = len(data["duration_ms"]) + print(f"\n {task_type} (count: {count})") + + duration_stats = TaskMetricsCollector._compute_stats( + data["duration_ms"] + ) + print( + f" Duration (ms): sum={duration_stats['sum']:.2f}, " + f"avg={duration_stats['avg']:.2f}, " + f"median={duration_stats['median']:.2f}, " + f"min={duration_stats['min']:.2f}, max={duration_stats['max']:.2f}" + ) + + if sum(data["total_tokens"]) > 0: + prompt_stats = TaskMetricsCollector._compute_stats( + data["prompt_tokens"] + ) + completion_stats = TaskMetricsCollector._compute_stats( + data["completion_tokens"] + ) + total_stats = TaskMetricsCollector._compute_stats( + data["total_tokens"] + ) + + print( + f" Prompt tokens: sum={prompt_stats['sum']:.0f}, " + f"avg={prompt_stats['avg']:.1f}, " + f"min={prompt_stats['min']:.0f}, max={prompt_stats['max']:.0f}" + ) + print( + f" Completion tokens: sum={completion_stats['sum']:.0f}, " + f"avg={completion_stats['avg']:.1f}, " + f"min={completion_stats['min']:.0f}, max={completion_stats['max']:.0f}" + ) + print( + f" Total tokens: sum={total_stats['sum']:.0f}, " + f"avg={total_stats['avg']:.1f}, " + f"min={total_stats['min']:.0f}, max={total_stats['max']:.0f}" + ) + + if perf_metrics_agg[task_type]: + print("\n Perf Metrics:") + for metric_name, values in sorted( + perf_metrics_agg[task_type].items() + ): + stats = TaskMetricsCollector._compute_stats(values) + print( + f" {metric_name}: sum={stats['sum']:.2f}, " + f"avg={stats['avg']:.2f}, " + f"min={stats['min']:.2f}, " + f"max={stats['max']:.2f}" + ) + + print("\n" + "=" * 80 + "\n") + + @staticmethod + def get_statistics(controller_name: str = None) -> dict[str, list[dict[str, Any]]]: + """Get statistics for a specific controller or all controllers.""" + if controller_name is not None: + return { + controller_name: TaskMetricsCollector.statistics.get( + controller_name, [] + ) + } + return TaskMetricsCollector.statistics + + @staticmethod + def get_all_records() -> list[dict[str, Any]]: + """Get all records across all controllers as a flat list.""" + all_records = [] + for records in TaskMetricsCollector.statistics.values(): + all_records.extend(records) + all_records.sort(key=lambda x: x.get("timestamp", 0)) + return all_records + + @staticmethod + def export_to_json(file_path: str, controller_name: str = None): + """Export metrics to a JSON file.""" + if controller_name is not None: + data = TaskMetricsCollector.statistics.get(controller_name, []) + else: + data = TaskMetricsCollector.statistics + with open(file_path, "w", encoding="utf-8") as f: + json.dump(data, f, indent=2, ensure_ascii=False, default=str) + + @staticmethod + def reset(controller_name: str = None): + """Reset statistics for a specific controller or all controllers.""" + if controller_name is not None: + if controller_name in TaskMetricsCollector.statistics: + TaskMetricsCollector.statistics[controller_name] = [] + else: + TaskMetricsCollector.statistics.clear() + + @staticmethod + def get_global_info() -> Any: + return TaskMetricsCollector.statistics diff --git a/areal/experimental/scaffolding/core/worker.py b/areal/experimental/scaffolding/core/worker.py new file mode 100644 index 0000000000..c2f76d06c7 --- /dev/null +++ b/areal/experimental/scaffolding/core/worker.py @@ -0,0 +1,171 @@ +# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +# Vendored from tensorrt_llm.scaffolding.worker +# TRTLLMWorker and MCPWorker omitted (require tensorrt_llm runtime). + +import os +from abc import ABC +from collections.abc import Callable + +import openai + +from .task import AssistantMessage, ChatTask, GenerationTask, Task, TaskStatus + + +class Worker(ABC): + def register_task_handler( + self, task_cls: type[Task], handler: Callable[[object, Task], TaskStatus] + ): + worker_cls = type(self) + worker_cls.task_handlers[task_cls] = handler + + async def run_task(self, task: Task) -> TaskStatus: + worker_cls = type(self) + if type(task) not in worker_cls.task_handlers: + return TaskStatus.WORKER_NOT_SUPPORTED + return await worker_cls.task_handlers[type(task)](self, task) + + task_handlers = {} + + def shutdown(self): + pass + + async def async_shutdown(self): + pass + + def __enter__(self): + return self + + def __exit__(self): + self.shutdown() + + +# helper function +def add_param_if_not_none(params, key, candidate_values): + for value in candidate_values: + if value is not None: + params[key] = value + return + + +# helper function +def add_attr_if_not_none(obj, attr, candidate_values): + for value in candidate_values: + if value is not None: + setattr(obj, attr, value) + return + + +def is_deterministic_mode(): + """Check if SCAFFOLDING_DETERMINISTIC environment variable is set.""" + return int(os.environ.get("SCAFFOLDING_DETERMINISTIC", 0)) == 1 + + +class OpenaiWorker(Worker): + def __init__( + self, + async_client: openai.AsyncOpenAI, + model: str, + kv_cache_hint_enabled: bool = False, + ): + self.model = model + self.async_client = async_client + self.kv_cache_hint_enabled = kv_cache_hint_enabled + + def convert_task_params(self, task: GenerationTask | ChatTask): + params = { + "model": self.model, + "extra_body": {}, + } + + if not isinstance(task, ChatTask): + params["prompt"] = task.input_str + add_param_if_not_none(params, "echo", [task.echo]) + + add_param_if_not_none(params, "best_of", [task.best_of]) + add_param_if_not_none(params, "frequency_penalty", [task.frequency_penalty]) + add_param_if_not_none(params, "logit_bias", [task.logit_bias]) + add_param_if_not_none(params, "logprobs", [task.num_logprobs]) + add_param_if_not_none(params, "max_tokens", [task.max_tokens]) + add_param_if_not_none(params, "n", [task.n]) + add_param_if_not_none(params, "presence_penalty", [task.presence_penalty]) + add_param_if_not_none(params, "seed", [task.seed]) + add_param_if_not_none(params, "stop", [task.stop]) + add_param_if_not_none(params, "suffix", [task.suffix]) + add_param_if_not_none(params, "temperature", [task.temperature]) + add_param_if_not_none(params, "top_p", [task.top_p]) + add_param_if_not_none(params, "user", [task.user]) + + # Override parameters for deterministic inference + if is_deterministic_mode(): + params["temperature"] = 0.0 + params["top_p"] = 1.0 + params["n"] = 1 + if "seed" not in params or params["seed"] is None: + params["seed"] = 42 + + if hasattr(task, "sub_request_markers") and len(task.sub_request_markers) > 0: + params["extra_body"]["agent_hierarchy"] = [task.sub_request_markers[-1]] + + return params + + def fill_generation_task_with_response( + self, task: GenerationTask, response: openai.Completion + ): + task.output_str = response.choices[0].text + task.output_tokens = response.choices[0].token_ids + task.finish_reason = response.choices[0].finish_reason + task.logprobs = response.choices[0].logprobs + task.perf_metrics = response.perf_metrics + + async def generation_handler(self, task: GenerationTask) -> TaskStatus: + params = self.convert_task_params(task) + + try: + response = await self.async_client.completions.create(**params) + self.fill_generation_task_with_response(task, response) + + return TaskStatus.SUCCESS + + except Exception as e: + print("Openai client get exception: " + str(e)) + return TaskStatus.WORKER_EXECEPTION + + async def chat_handler(self, task: ChatTask) -> TaskStatus: + params = self.convert_task_params(task) + params["messages"] = task.messages_to_dict_content() + params["model"] = self.model + if task.tools is not None: + params["tools"] = [tool.to_dict() for tool in task.tools] + + try: + response = await self.async_client.chat.completions.create(**params) + task.finish_reason = response.choices[0].finish_reason + task.perf_metrics = response.perf_metrics + content = response.choices[0].message.content + reasoning = response.choices[0].message.reasoning + reasoning_content = response.choices[0].message.reasoning_content + tool_calls = response.choices[0].message.tool_calls + task.messages.append( + AssistantMessage(content, reasoning, reasoning_content, tool_calls) + ) + if task.enable_token_counting: + task.prompt_tokens_num = response.usage.prompt_tokens + task.completion_tokens_num = response.usage.completion_tokens + if ( + hasattr(response.usage, "completion_tokens_details") + and response.usage.completion_tokens_details is not None + ): + task.reasoning_tokens_num = ( + response.usage.completion_tokens_details.reasoning_tokens + ) + + return TaskStatus.SUCCESS + + except Exception as e: + print("Openai chat client get exception: " + str(e)) + return TaskStatus.WORKER_EXECEPTION + + task_handlers = { + GenerationTask: generation_handler, + ChatTask: chat_handler, + } diff --git a/areal/experimental/scaffolding/task.py b/areal/experimental/scaffolding/task.py index 762d762311..3230608cee 100644 --- a/areal/experimental/scaffolding/task.py +++ b/areal/experimental/scaffolding/task.py @@ -2,7 +2,7 @@ RLVR Tasks for Scaffolding Framework. This module provides task definitions for RLVR (Reinforcement Learning with -Verifiable Rewards) that integrate with TensorRT-LLM's scaffolding framework. +Verifiable Rewards) that integrate with the scaffolding framework. """ from __future__ import annotations diff --git a/areal/experimental/scaffolding/worker.py b/areal/experimental/scaffolding/worker.py index 7ce6f5423e..2686a7c090 100644 --- a/areal/experimental/scaffolding/worker.py +++ b/areal/experimental/scaffolding/worker.py @@ -2,7 +2,7 @@ Worker implementations for Scaffolding Framework. This module provides Worker implementations that wrap AReaL inference engines -for use with TensorRT-LLM's scaffolding framework. +for use with the scaffolding framework. """ from __future__ import annotations @@ -158,7 +158,7 @@ def CreateWorkerFromEngine( """Create a scaffolding Worker from an AReaL SGLang engine. This function creates a Worker that wraps the given SGLang engine, - allowing it to be used with TensorRT-LLM's scaffolding framework. + allowing it to be used with the scaffolding framework. The worker uses the SGLang server's OpenAI-compatible API. Parameters @@ -187,7 +187,7 @@ def CreateWorkerFromEngine( worker = CreateWorkerFromEngine(engine) # Use with ScaffoldingLlm - from tensorrt_llm.scaffolding import ScaffoldingLlm, NativeGenerationController + from areal.experimental.scaffolding._compat import ScaffoldingLlm, NativeGenerationController llm = ScaffoldingLlm( controller, {NativeGenerationController.WorkerTag.GENERATION: worker}, diff --git a/areal/tests/experimental/scaffolding_core/__init__.py b/areal/tests/experimental/scaffolding_core/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/areal/tests/experimental/scaffolding_core/test_self_contained.py b/areal/tests/experimental/scaffolding_core/test_self_contained.py new file mode 100644 index 0000000000..4691074e90 --- /dev/null +++ b/areal/tests/experimental/scaffolding_core/test_self_contained.py @@ -0,0 +1,786 @@ +"""Tests proving the vendored scaffolding modules are self-contained. + +These tests verify that: +1. All scaffolding core modules import without tensorrt_llm installed. +2. Every public symbol in core/__init__.py and _compat.py is importable. +3. Core primitives (Task, Controller, Worker, ScaffoldingLlm, TaskCollection, + math_utils) function correctly in isolation. +4. AReaL wrapper modules (controllers, task, worker, workflow) import cleanly. +5. No source file under areal/experimental/scaffolding/ contains a live + ``import tensorrt_llm`` or ``from tensorrt_llm`` statement. +""" + +from __future__ import annotations + +import ast +import asyncio +import sys +from pathlib import Path +from unittest.mock import patch + +import pytest + +# ============================================================================ +# 1. Import isolation — tensorrt_llm must NOT be importable +# ============================================================================ + + +class TestNoTensorRTLLMDependency: + """Verify that tensorrt_llm is not required at runtime.""" + + def test_tensorrt_llm_not_installed(self): + """tensorrt_llm should not be importable in the test environment.""" + assert "tensorrt_llm" not in sys.modules or sys.modules["tensorrt_llm"] is None + + def test_no_live_tensorrt_llm_imports_in_source(self): + """No .py file under scaffolding/ should have a live import of tensorrt_llm. + + Comments and docstrings are allowed; only top-level or function-level + ``import tensorrt_llm`` / ``from tensorrt_llm import ...`` are flagged. + """ + scaffolding_root = ( + Path(__file__).resolve().parents[3] / "experimental" / "scaffolding" + ) + violations = [] + for py_file in scaffolding_root.rglob("*.py"): + try: + tree = ast.parse(py_file.read_text(), filename=str(py_file)) + except SyntaxError: + continue + for node in ast.walk(tree): + if isinstance(node, ast.Import): + for alias in node.names: + if alias.name == "tensorrt_llm" or alias.name.startswith( + "tensorrt_llm." + ): + violations.append( + f"{py_file.relative_to(scaffolding_root)}:{node.lineno}" + ) + elif isinstance(node, ast.ImportFrom): + if node.module and ( + node.module == "tensorrt_llm" + or node.module.startswith("tensorrt_llm.") + ): + violations.append( + f"{py_file.relative_to(scaffolding_root)}:{node.lineno}" + ) + + assert violations == [], "Found live tensorrt_llm imports in:\n" + "\n".join( + violations + ) + + +# ============================================================================ +# 2. Core module imports — every public symbol is importable +# ============================================================================ + + +class TestCoreImports: + """Verify every public symbol in core/ is importable.""" + + def test_core_init_imports(self): + """All symbols in core/__init__.py should import successfully.""" + from areal.experimental.scaffolding.core import ( # noqa: F401 + AssistantMessage, + BestOfNController, + ChatTask, + Controller, + GenerationTask, + MajorityVoteController, + NativeChatController, + NativeGenerationController, + NativeRewardController, + OpenAIToolDescription, + OpenaiWorker, + ParallelProcess, + RoleMessage, + ScaffoldingLlm, + ScaffoldingOutput, + StreamGenerationTask, + SystemMessage, + Task, + TaskCollection, + TaskStatus, + UserMessage, + Worker, + extract_answer_from_boxed, + extract_answer_with_regex, + get_digit_majority_vote_result, + with_task_collection, + ) + + # Spot-check a few are real classes/functions + assert callable(Controller) + assert callable(extract_answer_from_boxed) + assert callable(with_task_collection) + + def test_compat_reexports_match_core(self): + """_compat.py should re-export the same symbols as core/__init__.py + (minus math_utils which is not in _compat).""" + from areal.experimental.scaffolding import _compat, core + + compat_all = set(_compat.__all__) + # _compat intentionally omits math_utils functions + math_utils_names = { + "extract_answer_from_boxed", + "extract_answer_with_regex", + "get_digit_majority_vote_result", + } + core_all_minus_math = set(core.__all__) - math_utils_names + assert compat_all == core_all_minus_math + + def test_core_submodule_imports(self): + """Each core submodule should import independently.""" + import areal.experimental.scaffolding.core.controller # noqa: F401 + import areal.experimental.scaffolding.core.math_utils # noqa: F401 + import areal.experimental.scaffolding.core.result # noqa: F401 + import areal.experimental.scaffolding.core.scaffolding_llm # noqa: F401 + import areal.experimental.scaffolding.core.task # noqa: F401 + import areal.experimental.scaffolding.core.task_collection # noqa: F401 + import areal.experimental.scaffolding.core.worker # noqa: F401 + + +# ============================================================================ +# 3. Core primitives — functional tests +# ============================================================================ + + +class TestTask: + """Tests for Task, GenerationTask, ChatTask, and related dataclasses.""" + + def test_task_creation(self): + from areal.experimental.scaffolding.core.task import Task + + t = Task() + assert t.worker_tag is None + assert t.streaming_output_flag is False + assert t.streaming_output_list == [] + + def test_generation_task_create_from_prompt(self): + from areal.experimental.scaffolding.core.task import GenerationTask + + t = GenerationTask.create_from_prompt("Hello world") + assert t.input_str == "Hello world" + assert t.skip_tokenizer is False + assert t.skip_detokenizer is False + + def test_generation_task_scaffolding_output(self): + from areal.experimental.scaffolding.core.task import GenerationTask + + t = GenerationTask(output_str="result", output_tokens=[10, 20]) + output = t.create_scaffolding_output() + assert output.text == "result" + assert output.token_ids == [10, 20] + + def test_stream_generation_task_from_generation(self): + from areal.experimental.scaffolding.core.task import ( + GenerationTask, + StreamGenerationTask, + ) + + gen = GenerationTask(input_str="prompt", max_tokens=100) + stream = StreamGenerationTask.create_from_generation_task(gen, streaming_step=5) + assert stream.input_str == "prompt" + assert stream.max_tokens == 100 + assert stream.streaming_step == 5 + + def test_chat_task_create_from_prompt(self): + from areal.experimental.scaffolding.core.task import ( + ChatTask, + SystemMessage, + ) + + t = ChatTask.create_from_prompt( + "What is 2+2?", + system_prompts=[SystemMessage("You are a math tutor")], + ) + assert len(t.messages) == 2 + assert t.messages[0].role == "system" + assert t.messages[1].role == "user" + + def test_chat_task_create_from_messages(self): + from areal.experimental.scaffolding.core.task import ( + ChatTask, + UserMessage, + ) + + msgs = [UserMessage("hi"), UserMessage("hello")] + t = ChatTask.create_from_messages(msgs) + assert len(t.messages) == 2 + + def test_chat_task_add_message(self): + from areal.experimental.scaffolding.core.task import ( + AssistantMessage, + ChatTask, + UserMessage, + ) + + t = ChatTask() + t.add_message(UserMessage("Q")) + t.add_message(AssistantMessage("A")) + assert len(t.messages) == 2 + assert t.messages[0].role == "user" + assert t.messages[1].role == "assistant" + + def test_chat_task_messages_to_dict(self): + from areal.experimental.scaffolding.core.task import ( + ChatTask, + UserMessage, + ) + + t = ChatTask() + t.add_message(UserMessage("hello")) + dicts = t.messages_to_dict_content() + assert dicts == [{"role": "user", "content": "hello"}] + + def test_role_message_str_repr(self): + from areal.experimental.scaffolding.core.task import UserMessage + + m = UserMessage("hi") + assert '"role": "user"' in str(m) + assert "user" in repr(m) + + def test_role_message_from_dict(self): + from areal.experimental.scaffolding.core.task import RoleMessage + + m = RoleMessage.from_dict({"role": "user", "content": "test"}) + assert m.role == "user" + assert m.content == "test" + + def test_assistant_message_with_reasoning(self): + from areal.experimental.scaffolding.core.task import AssistantMessage + + m = AssistantMessage("answer", reasoning="chain of thought") + assert m.role == "assistant" + assert m.reasoning == "chain of thought" + assert '"reasoning"' in str(m) + + def test_openai_tool_description(self): + from areal.experimental.scaffolding.core.task import OpenAIToolDescription + + tool = OpenAIToolDescription( + "my_func", "Does something", {"x": {"type": "int"}} + ) + d = tool.to_dict() + assert d["type"] == "function" + assert d["function"]["name"] == "my_func" + assert d["function"]["description"] == "Does something" + + def test_task_status_values(self): + from areal.experimental.scaffolding.core.task import TaskStatus + + assert TaskStatus.SUCCESS.value == "success" + assert TaskStatus.WORKER_NOT_SUPPORTED.value == "worker_not_supported" + assert TaskStatus.WORKER_EXECEPTION.value == "worker_exception" + + +class TestController: + """Tests for Controller and built-in controller subclasses.""" + + def test_controller_is_abstract(self): + from areal.experimental.scaffolding.core.controller import Controller + + ctrl = Controller() + assert hasattr(ctrl, "task_collections") + with pytest.raises(NotImplementedError): + list(ctrl.process([])) + + def test_controller_clone(self): + from areal.experimental.scaffolding.core.controller import Controller + + class MyCtrl(Controller): + def __init__(self, val): + super().__init__() + self.val = val + + def process(self, tasks, **kw): + yield tasks + + original = MyCtrl(42) + cloned = original.clone() + assert cloned.val == 42 + assert cloned is not original + + def test_native_generation_controller(self): + from areal.experimental.scaffolding.core.controller import ( + NativeGenerationController, + ) + from areal.experimental.scaffolding.core.task import GenerationTask + + ctrl = NativeGenerationController( + sampling_params={"temperature": 0.7, "max_tokens": 100} + ) + task = GenerationTask(input_str="test") + results = list(ctrl.process([task])) + + assert len(results) == 1 + assert task.worker_tag == NativeGenerationController.WorkerTag.GENERATION + assert task.temperature == 0.7 + assert task.max_tokens == 100 + + def test_native_generation_controller_ignores_invalid_params(self): + from areal.experimental.scaffolding.core.controller import ( + NativeGenerationController, + ) + + ctrl = NativeGenerationController( + sampling_params={"invalid_param_xyz": 999, "temperature": 0.5} + ) + assert "temperature" in ctrl.sampling_params + assert "invalid_param_xyz" not in ctrl.sampling_params + + def test_native_chat_controller(self): + from areal.experimental.scaffolding.core.controller import ( + NativeChatController, + ) + from areal.experimental.scaffolding.core.task import ( + ChatTask, + GenerationTask, + ) + + ctrl = NativeChatController() + task = GenerationTask(input_str="What is 2+2?") + results = list(ctrl.process([task])) + + assert len(results) == 1 + # NativeChatController wraps in ChatTask + yielded_tasks = results[0] + assert isinstance(yielded_tasks[0], ChatTask) + + def test_native_reward_controller(self): + from areal.experimental.scaffolding.core.controller import ( + NativeRewardController, + ) + from areal.experimental.scaffolding.core.task import GenerationTask + + ctrl = NativeRewardController() + task = GenerationTask(input_str="test") + results = list(ctrl.process([task])) + + assert len(results) == 1 + assert task.worker_tag == NativeRewardController.WorkerTag.REWARD + + def test_parallel_process_creation(self): + from areal.experimental.scaffolding.core.controller import ( + NativeGenerationController, + ParallelProcess, + ) + from areal.experimental.scaffolding.core.task import GenerationTask + + ctrl1 = NativeGenerationController() + ctrl2 = NativeGenerationController() + tasks1 = [GenerationTask(input_str="a")] + tasks2 = [GenerationTask(input_str="b")] + + pp = ParallelProcess( + controllers=[ctrl1, ctrl2], + tasks_list=[tasks1, tasks2], + kwargs_list=[{}, {}], + ) + assert len(pp.sub_gens) == 2 + + +class TestWorker: + """Tests for Worker base class.""" + + def test_worker_is_abstract(self): + from areal.experimental.scaffolding.core.worker import Worker + + w = Worker() + assert hasattr(w, "task_handlers") + + @pytest.mark.asyncio + async def test_worker_unsupported_task(self): + from areal.experimental.scaffolding.core.task import Task, TaskStatus + from areal.experimental.scaffolding.core.worker import Worker + + class EmptyWorker(Worker): + task_handlers = {} + + w = EmptyWorker() + status = await w.run_task(Task()) + assert status == TaskStatus.WORKER_NOT_SUPPORTED + + @pytest.mark.asyncio + async def test_worker_register_handler(self): + from areal.experimental.scaffolding.core.task import ( + GenerationTask, + TaskStatus, + ) + from areal.experimental.scaffolding.core.worker import Worker + + class MyWorker(Worker): + task_handlers = {} + + async def my_handler(self, task): + task.output_str = "handled" + return TaskStatus.SUCCESS + + w = MyWorker() + w.register_task_handler(GenerationTask, my_handler) + task = GenerationTask(input_str="test") + status = await w.run_task(task) + assert status == TaskStatus.SUCCESS + assert task.output_str == "handled" + + def test_worker_context_manager(self): + from areal.experimental.scaffolding.core.worker import Worker + + w = Worker() + result = w.__enter__() + assert result is w + + def test_is_deterministic_mode(self): + from areal.experimental.scaffolding.core.worker import is_deterministic_mode + + assert is_deterministic_mode() is False + + with patch.dict("os.environ", {"SCAFFOLDING_DETERMINISTIC": "1"}): + assert is_deterministic_mode() is True + + +class TestResult: + """Tests for ScaffoldingOutput and ScaffoldingResult.""" + + def test_scaffolding_output(self): + from areal.experimental.scaffolding.core.result import ScaffoldingOutput + + o = ScaffoldingOutput(text="hello", token_ids=[1, 2, 3]) + assert o.text == "hello" + assert o.token_ids == [1, 2, 3] + + def test_scaffolding_result_set_output(self): + from areal.experimental.scaffolding.core.result import ( + ScaffoldingOutput, + ScaffoldingResult, + ) + + result = ScaffoldingResult() + assert not result._done + + output = ScaffoldingOutput("text", [1, 2]) + result.set_output(output) + + # After set_output, we should be able to get the result + # by draining the queue + loop = asyncio.new_event_loop() + try: + done = loop.run_until_complete(result.aresult()) + assert done._done + assert done.outputs[0].text == "text" + assert done.outputs[0].token_ids == [1, 2] + finally: + loop.close() + + def test_scaffolding_result_set_output_none(self): + from areal.experimental.scaffolding.core.result import ScaffoldingResult + + result = ScaffoldingResult() + result.set_output(None) + + loop = asyncio.new_event_loop() + try: + done = loop.run_until_complete(result.aresult()) + assert done._done + finally: + loop.close() + + def test_scaffolding_result_task_collections(self): + from areal.experimental.scaffolding.core.result import ScaffoldingResult + + result = ScaffoldingResult() + assert result.task_collections is None + result.set_task_collections({"key": "value"}) + assert result.task_collections == {"key": "value"} + + +class TestTaskCollection: + """Tests for TaskCollection and the with_task_collection decorator.""" + + def test_task_collection_base(self): + from areal.experimental.scaffolding.core.task_collection import TaskCollection + + tc = TaskCollection() + tc.before_yield([]) # Should not raise + tc.after_yield([]) # Should not raise + assert TaskCollection.get_global_info() is None + + def test_with_task_collection_decorator(self): + from areal.experimental.scaffolding.core.controller import Controller + from areal.experimental.scaffolding.core.task import Task + from areal.experimental.scaffolding.core.task_collection import ( + TaskCollection, + with_task_collection, + ) + + class MyCollection(TaskCollection): + def __init__(self): + super().__init__() + self.before_count = 0 + self.after_count = 0 + + def before_yield(self, tasks): + self.before_count += 1 + + def after_yield(self, tasks): + self.after_count += 1 + + @with_task_collection("my_tc", MyCollection) + class MyController(Controller): + def process(self, tasks, **kwargs): + yield tasks + + ctrl = MyController() + assert "my_tc" in ctrl.task_collections + tc = ctrl.task_collections["my_tc"] + assert isinstance(tc, MyCollection) + + list(ctrl.process([Task()])) + assert tc.before_count == 1 + assert tc.after_count == 1 + + def test_generation_token_counter(self): + from areal.experimental.scaffolding.core.task import GenerationTask + from areal.experimental.scaffolding.core.task_collection import ( + GenerationTokenCounter, + ) + + counter = GenerationTokenCounter() + task = GenerationTask(output_tokens=[1, 2, 3]) + + counter.before_yield([task]) + # Simulate worker adding tokens + task.output_tokens = [1, 2, 3, 4, 5] + counter.after_yield([task]) + + assert counter.generation_token_count == 2 # 5 - 3 = 2 new tokens + + def test_task_metrics_collector_reset(self): + from areal.experimental.scaffolding.core.task_collection import ( + TaskMetricsCollector, + ) + + TaskMetricsCollector.statistics["test_ctrl"] = [{"data": 1}] + TaskMetricsCollector.reset("test_ctrl") + assert TaskMetricsCollector.statistics["test_ctrl"] == [] + TaskMetricsCollector.reset() + assert TaskMetricsCollector.statistics == {} + + +class TestMathUtils: + """Tests for math_utils functions.""" + + def test_extract_answer_from_boxed_simple(self): + from areal.experimental.scaffolding.core.math_utils import ( + extract_answer_from_boxed, + ) + + assert extract_answer_from_boxed("The answer is \\boxed{42}") == "42" + + def test_extract_answer_from_boxed_nested(self): + from areal.experimental.scaffolding.core.math_utils import ( + extract_answer_from_boxed, + ) + + assert extract_answer_from_boxed("\\boxed{x^{2}}") == "x^{2}" + + def test_extract_answer_from_boxed_none(self): + from areal.experimental.scaffolding.core.math_utils import ( + extract_answer_from_boxed, + ) + + assert extract_answer_from_boxed("No boxed answer here") is None + + def test_extract_answer_with_regex(self): + from areal.experimental.scaffolding.core.math_utils import ( + extract_answer_with_regex, + ) + + result = extract_answer_with_regex("The final answer is 42") + assert result == "42" + + def test_extract_answer_with_regex_no_match(self): + from areal.experimental.scaffolding.core.math_utils import ( + extract_answer_with_regex, + ) + + assert extract_answer_with_regex("Nothing relevant") is None + + def test_get_digit_majority_vote_result(self): + from areal.experimental.scaffolding.core.math_utils import ( + get_digit_majority_vote_result, + ) + + results = [ + "The answer is \\boxed{42}", + "Therefore \\boxed{42}", + "I get \\boxed{99}", + ] + index, answer = get_digit_majority_vote_result(results) + assert answer == "42" + + def test_get_digit_majority_vote_no_valid(self): + from areal.experimental.scaffolding.core.math_utils import ( + get_digit_majority_vote_result, + ) + + results = ["no boxed", "nothing here"] + index, answer = get_digit_majority_vote_result(results) + assert answer is None + + +# ============================================================================ +# 4. AReaL wrapper modules — import and basic function +# ============================================================================ + + +class TestWrapperImports: + """Verify AReaL wrapper modules import without tensorrt_llm.""" + + def test_compat_module_imports(self): + from areal.experimental.scaffolding._compat import ( # noqa: F401 + AssistantMessage, + BestOfNController, + ChatTask, + Controller, + GenerationTask, + MajorityVoteController, + NativeChatController, + NativeGenerationController, + NativeRewardController, + OpenAIToolDescription, + OpenaiWorker, + ParallelProcess, + RoleMessage, + ScaffoldingLlm, + ScaffoldingOutput, + StreamGenerationTask, + SystemMessage, + Task, + TaskCollection, + TaskStatus, + UserMessage, + Worker, + with_task_collection, + ) + + def test_controllers_module_imports(self): + from areal.experimental.scaffolding.controllers import ( # noqa: F401 + ChatTracer, + PipelineTrajectoryMaker, + RLVRRewardController, + TraceTrajectoryMaker, + ) + + def test_task_module_imports(self): + from areal.experimental.scaffolding.task import ( # noqa: F401 + ChatRewardTask, + RLVRRewardTask, + TraceGenerationTask, + ) + + def test_worker_module_imports(self): + from areal.experimental.scaffolding.worker import ( # noqa: F401 + CreateWorkerFromEngine, + SGLangWorker, + ) + + def test_workflow_module_imports(self): + from areal.experimental.scaffolding.workflow import ( + ScaffoldingWorkflow, # noqa: F401 + ) + + def test_top_level_package_imports(self): + from areal.experimental.scaffolding import ( # noqa: F401 + ChatRewardTask, + ChatTracer, + CreateWorkerFromEngine, + PipelineTrajectoryMaker, + RLVRRewardController, + RLVRRewardTask, + ScaffoldingWorkflow, + SGLangWorker, + TraceGenerationTask, + TraceTrajectoryMaker, + ) + + +# ============================================================================ +# 5. Cross-module integration — core + wrappers work together +# ============================================================================ + + +class TestCrossModuleIntegration: + """Verify that core primitives and AReaL wrappers interoperate.""" + + def test_rlvr_reward_task_inherits_from_core_task(self): + from areal.experimental.scaffolding.core.task import Task + from areal.experimental.scaffolding.task import RLVRRewardTask + + t = RLVRRewardTask(prompt_str="Q", completion_str="A") + assert isinstance(t, Task) + + def test_sglang_worker_inherits_from_core_openai_worker(self): + from areal.experimental.scaffolding.core.worker import OpenaiWorker + from areal.experimental.scaffolding.worker import SGLangWorker + + assert issubclass(SGLangWorker, OpenaiWorker) + + def test_rlvr_reward_controller_inherits_from_core_controller(self): + from areal.experimental.scaffolding.controllers import RLVRRewardController + from areal.experimental.scaffolding.core.controller import Controller + + assert issubclass(RLVRRewardController, Controller) + + def test_pipeline_trajectory_maker_inherits_from_core_controller(self): + from areal.experimental.scaffolding.controllers import PipelineTrajectoryMaker + from areal.experimental.scaffolding.core.controller import Controller + + assert issubclass(PipelineTrajectoryMaker, Controller) + + def test_native_gen_controller_process_with_compat_import(self): + """Using _compat imports should produce the same result as core imports.""" + from areal.experimental.scaffolding._compat import ( + GenerationTask, + NativeGenerationController, + ) + + ctrl = NativeGenerationController(sampling_params={"temperature": 0.5}) + task = GenerationTask.create_from_prompt("test") + results = list(ctrl.process([task])) + assert len(results) == 1 + assert task.temperature == 0.5 + + def test_scaffolding_workflow_is_rollout_workflow(self): + from areal.api.workflow_api import RolloutWorkflow + from areal.experimental.scaffolding.workflow import ScaffoldingWorkflow + + assert issubclass(ScaffoldingWorkflow, RolloutWorkflow) + + def test_compat_classes_are_same_as_core_classes(self): + """Verify _compat re-exports are the exact same class objects as core.""" + from areal.experimental.scaffolding._compat import ( + Controller as CompatController, + ) + from areal.experimental.scaffolding._compat import ( + GenerationTask as CompatGenTask, + ) + from areal.experimental.scaffolding._compat import ( + ScaffoldingLlm as CompatLlm, + ) + from areal.experimental.scaffolding._compat import Task as CompatTask + from areal.experimental.scaffolding._compat import Worker as CompatWorker + from areal.experimental.scaffolding.core.controller import Controller + from areal.experimental.scaffolding.core.scaffolding_llm import ScaffoldingLlm + from areal.experimental.scaffolding.core.task import GenerationTask, Task + from areal.experimental.scaffolding.core.worker import Worker + + assert CompatTask is Task + assert CompatGenTask is GenerationTask + assert CompatController is Controller + assert CompatWorker is Worker + assert CompatLlm is ScaffoldingLlm + + +if __name__ == "__main__": + pytest.main([__file__, "-v"]) diff --git a/examples/scaffolding/README.md b/examples/scaffolding/README.md index e754edd8fb..b4ac39115d 100644 --- a/examples/scaffolding/README.md +++ b/examples/scaffolding/README.md @@ -1,7 +1,7 @@ # Scaffolding Framework Examples for AReaL -This directory contains examples demonstrating how to use the TensorRT-LLM Scaffolding -framework with AReaL for reinforcement learning training. +This directory contains examples demonstrating how to use the Scaffolding framework with +AReaL for reinforcement learning training. ## Overview @@ -83,7 +83,7 @@ scaffolding_workflow = ScaffoldingWorkflow(scaffolding_llm) │ │ │ │ │ │ │ │ │ ┌───────────────────────────────┐ │ │ │ Data ─────────────────────────┼──┼──┼──► NativeGenerationController │ │ │ │ - │ │ │ │ (from tensorrt_llm) │ │ │ │ + │ │ │ │ (from scaffolding.core) │ │ │ │ │ │ │ └───────────────┬───────────────┘ │ │ │ │ │ │ │ │ │ │ │ │ │ ▼ │ │ │ @@ -169,7 +169,7 @@ engine: You can create custom reward controllers by subclassing the base Controller: ```python -from tensorrt_llm.scaffolding import Controller +from areal.experimental.scaffolding._compat import Controller class CustomRewardController(Controller): def __init__(self, reward_fn): @@ -193,7 +193,7 @@ class CustomRewardController(Controller): For different RL algorithms, you may need different trajectory formats: ```python -from tensorrt_llm.scaffolding import Controller +from areal.experimental.scaffolding._compat import Controller import torch class CustomTrajectoryMaker(Controller): From 518eee962012f1b8e469db526232ab6b2ee24f20 Mon Sep 17 00:00:00 2001 From: Fred Wei <20514172+WeiHaocheng@users.noreply.github.com> Date: Tue, 3 Mar 2026 01:00:05 +0000 Subject: [PATCH 07/18] Move chat_scaffolding.py and debug --- examples/scaffolding/chat_scaffolding.py | 216 +++++++++++++++++++++ examples/scaffolding/chat_scaffolding.yaml | 179 +++++++++++++++++ 2 files changed, 395 insertions(+) create mode 100644 examples/scaffolding/chat_scaffolding.py create mode 100644 examples/scaffolding/chat_scaffolding.yaml diff --git a/examples/scaffolding/chat_scaffolding.py b/examples/scaffolding/chat_scaffolding.py new file mode 100644 index 0000000000..257cf58cb1 --- /dev/null +++ b/examples/scaffolding/chat_scaffolding.py @@ -0,0 +1,216 @@ +""" +Multi-turn Chat Scaffolding Example. + +This example demonstrates multi-turn chat-based RL training on GSM8K. +Each episode runs multiple generation turns with a reflection message +appended between turns to prompt the model to retry. + +Usage: + python examples/scaffolding/chat_scaffolding.py \ + --config examples/scaffolding/chat_scaffolding.yaml \ + +scheduler.type=local experiment_name=areal trial_name=chat_scaffolding +""" + +import sys +from collections.abc import Callable +from typing import Any + +import torch +from transformers import PreTrainedTokenizerFast + +from areal.api.cli_args import GenerationHyperparameters, GRPOConfig, load_expr_config +from areal.api.engine_api import InferenceEngine +from areal.dataset import get_custom_dataset +from areal.experimental.scaffolding._compat import GenerationTask +from areal.experimental.scaffolding.workflow import ScaffoldingWorkflow +from areal.trainer import PPOTrainer +from areal.utils import logging +from areal.utils.hf_utils import load_hf_tokenizer + +logger = logging.getLogger("ChatScaffoldingWorkflow") + +DEFAULT_REFLECTION_MESSAGE = ( + "Your answer is either wrong or not parsable to the reward function. " + "You may misunderstand the original question. " + "Please carefully read the original question, check the previous errors, " + "and try to answer it again." +) + + +class ChatScaffoldingWorkflow(ScaffoldingWorkflow): + """ScaffoldingWorkflow for multi-turn chat with reflection retry. + + Each episode runs up to ``max_turns`` generation turns. After each + non-final turn the ``reflection_message`` is appended as a user message + to prompt the model to retry. Reward is computed on the final turn only. + + Generation and reward use the same direct-API approach as the base class + (``_generate_via_worker`` + ``_compute_rewards_via_controller``). + + Parameters + ---------- + reward_fn : Callable | str + The reward function, or an importable string path. + gconfig : GenerationHyperparameters + Generation hyperparameters. + tokenizer : PreTrainedTokenizerFast | str + Tokenizer or path to load it. + enable_thinking : bool + Whether to enable thinking tokens. + max_turns : int + Maximum number of chat turns per episode. + reflection_message : str + Message appended after each non-final turn to prompt retry. + """ + + def __init__( + self, + reward_fn: Callable[..., Any] | str, + gconfig: GenerationHyperparameters, + tokenizer: PreTrainedTokenizerFast | str, + enable_thinking: bool = False, + max_turns: int = 2, + reflection_message: str = DEFAULT_REFLECTION_MESSAGE, + ): + super().__init__( + reward_fn=reward_fn, + gconfig=gconfig, + tokenizer=tokenizer, + enable_thinking=enable_thinking, + ) + self.max_turns = max_turns + self.reflection_message = reflection_message + + async def arun_episode( + self, engine: InferenceEngine, data: dict[str, Any] + ) -> dict[str, torch.Tensor]: + """Run a single multi-turn chat episode. + + Each turn: generate via SGLang OpenAI API, then optionally append a + reflection message and retry. After all turns, compute reward on the + final completion and build tensor dicts identical to the base class. + + Parameters + ---------- + engine : InferenceEngine + The inference engine. + data : dict[str, Any] + Input data containing messages and ground truth. + + Returns + ------- + dict[str, torch.Tensor] + Trajectory tensors for PPO training. + """ + if self.worker is None: + self._lazy_init_scaffolding(engine) + + # Start from the original messages + messages = list(data["messages"]) + last_output_str = "" + all_output_tokens: list[int] = [] + + for turn in range(self.max_turns): + # Build prompt for this turn + input_ids = list( + self.tokenizer.apply_chat_template( + messages, + tokenize=True, + add_generation_prompt=True, + enable_thinking=self.enable_thinking, + ) + ) + prompt_str = self.tokenizer.decode(input_ids) + + # Generate via the SGLang OpenAI API (same as base class) + gen_task = await self._generate_via_worker(prompt_str, input_ids) + last_output_str = gen_task.output_str or "" + all_output_tokens = list(gen_task.output_tokens or []) + + # Append the assistant response to messages + messages.append({"role": "assistant", "content": last_output_str}) + + # Append reflection message for non-final turns + if turn < self.max_turns - 1: + messages.append({"role": "user", "content": self.reflection_message}) + + # Compute reward on the final turn's output (same pattern as base class) + final_input_ids = list( + self.tokenizer.apply_chat_template( + data["messages"], + tokenize=True, + add_generation_prompt=True, + enable_thinking=self.enable_thinking, + ) + ) + final_prompt_str = self.tokenizer.decode(final_input_ids) + + final_gen_task = GenerationTask( + input_str=final_prompt_str, + input_tokens=final_input_ids, + output_str=last_output_str, + output_tokens=all_output_tokens, + ) + reward = await self._compute_rewards_via_controller( + final_gen_task, final_prompt_str, data + ) + + # Build tensor dict (same as base class) + seq = final_input_ids + all_output_tokens + logprobs = [0.0] * len(seq) + loss_mask = [0] * len(final_input_ids) + [1] * len(all_output_tokens) + versions = [-1] * len(seq) + + res = { + "input_ids": torch.tensor(seq, dtype=torch.int32), + "loss_mask": torch.tensor(loss_mask, dtype=torch.int32), + "logprobs": torch.tensor(logprobs, dtype=torch.float32), + "versions": torch.tensor(versions, dtype=torch.int32), + "attention_mask": torch.ones(len(seq), dtype=torch.bool), + "rewards": torch.tensor(reward, dtype=torch.float32), + } + return {k: v.unsqueeze(0) for k, v in res.items()} + + +def main(args): + """Main entry point for multi-turn chat scaffolding training.""" + config, _ = load_expr_config(args, GRPOConfig) + tokenizer = load_hf_tokenizer(config.tokenizer_path) + + train_dataset = get_custom_dataset( + split="train", + dataset_config=config.train_dataset, + tokenizer=tokenizer, + ) + valid_dataset = get_custom_dataset( + split="test", + dataset_config=config.valid_dataset, + tokenizer=tokenizer, + ) + + workflow_kwargs = dict( + reward_fn="areal.reward.gsm8k.gsm8k_reward_fn", + gconfig=config.gconfig, + tokenizer=config.tokenizer_path, + enable_thinking=False, + max_turns=2, + reflection_message=DEFAULT_REFLECTION_MESSAGE, + ) + eval_workflow_kwargs = workflow_kwargs.copy() + eval_workflow_kwargs["gconfig"] = config.gconfig.new(temperature=0.6) + + with PPOTrainer( + config, + train_dataset=train_dataset, + valid_dataset=valid_dataset, + ) as trainer: + trainer.train( + workflow="examples.scaffolding.chat_scaffolding.ChatScaffoldingWorkflow", + workflow_kwargs=workflow_kwargs, + eval_workflow="examples.scaffolding.chat_scaffolding.ChatScaffoldingWorkflow", + eval_workflow_kwargs=eval_workflow_kwargs, + ) + + +if __name__ == "__main__": + main(sys.argv[1:]) diff --git a/examples/scaffolding/chat_scaffolding.yaml b/examples/scaffolding/chat_scaffolding.yaml new file mode 100644 index 0000000000..4f9a1b8b4e --- /dev/null +++ b/examples/scaffolding/chat_scaffolding.yaml @@ -0,0 +1,179 @@ +# Multi-turn Chat Scaffolding Example Configuration for GSM8K +# Uses ChatScaffoldingWorkflow with TraceTrajectoryMaker + MultiTurnChatController +# Compatible with GRPOConfig, for 8-GPU setup (4 inference + 4 training) + +experiment_name: gsm8k-chat-scaffolding +trial_name: trial0 + +seed: 1 +enable_offload: false +total_train_epochs: 10 +tokenizer_path: ${actor.path} + +cluster: + n_nodes: 1 + n_gpus_per_node: 1 + fileroot: /tmp/areal/experiments + name_resolve: + type: nfs + nfs_record_root: /tmp/areal/name_resolve + +allocation_mode: sglang:d1+d1 + +rollout: + experiment_name: ${experiment_name} + trial_name: ${trial_name} + max_concurrent_rollouts: 256 + queue_size: null + consumer_batch_size: ${train_dataset.batch_size} + max_head_offpolicyness: 2 + enable_rollout_tracing: false + scheduling_spec: ${actor.scheduling_spec} + fileroot: ${cluster.fileroot} + tokenizer_path: ${tokenizer_path} + dump_to_file: true + +gconfig: + n_samples: 8 + min_new_tokens: 0 + max_new_tokens: 1024 + greedy: false + temperature: 1.0 + +actor: + experiment_name: ${experiment_name} + trial_name: ${trial_name} + path: Qwen/Qwen2.5-3B-Instruct + init_from_scratch: false + disable_dropout: true + gradient_checkpointing: true + dtype: bfloat16 + mb_spec: + max_tokens_per_mb: 10240 + optimizer: + type: adam + lr: 1.0e-6 + weight_decay: 0.01 + beta1: 0.9 + beta2: 0.999 + eps: 1e-8 + lr_scheduler_type: constant + gradient_clipping: 1.0 + warmup_steps_proportion: 0.001 + eps_clip: 0.2 + temperature: ${gconfig.temperature} + reward_scaling: 10.0 + reward_bias: -0.5 + kl_ctl: 0.0 + ppo_n_minibatches: 1 + recompute_logprob: true + use_decoupled_loss: true + behav_imp_weight_cap: 5.0 + reward_norm: + mean_level: group + std_level: group + group_size: ${gconfig.n_samples} + adv_norm: + mean_level: batch + std_level: batch + weight_update_mode: disk + max_new_tokens: ${gconfig.max_new_tokens} + scheduling_spec: + - task_type: worker + port_count: 2 + gpu: 1 + mem: 32 + cmd: python3 -m areal.infra.rpc.rpc_server + env_vars: {} + +ref: + experiment_name: ${experiment_name} + trial_name: ${trial_name} + path: ${actor.path} + init_from_scratch: false + disable_dropout: true + dtype: ${actor.dtype} + mb_spec: + max_tokens_per_mb: 10240 + optimizer: null + scheduling_strategy: + type: colocation + target: actor + scheduling_spec: ${actor.scheduling_spec} + +# SGLang +sglang: + model_path: ${actor.path} + random_seed: ${seed} + skip_tokenizer_init: false + dtype: ${actor.dtype} + max_running_requests: null + context_length: 2048 + mem_fraction_static: 0.5 + attention_backend: flashinfer + +vllm: + model: ${actor.path} + seed: ${seed} + skip_tokenizer_init: false + dtype: ${actor.dtype} + max_model_len: 4096 + gpu_memory_utilization: 0.8 + +# Datasets +train_dataset: + batch_size: 256 + shuffle: true + pin_memory: true + num_workers: 4 + path: openai/gsm8k + type: rl + max_length: 2048 + +valid_dataset: + batch_size: 256 + pin_memory: true + num_workers: 4 + path: openai/gsm8k + type: rl + +# Utilities +saver: + experiment_name: ${experiment_name} + trial_name: ${trial_name} + fileroot: ${cluster.fileroot} + freq_epochs: 1 + freq_steps: null + freq_secs: null + +recover: + mode: disabled + experiment_name: ${experiment_name} + trial_name: ${trial_name} + fileroot: ${cluster.fileroot} + freq_epochs: 1 + freq_steps: null + freq_secs: 3600 + +evaluator: + experiment_name: ${experiment_name} + trial_name: ${trial_name} + fileroot: ${cluster.fileroot} + freq_epochs: 1 + freq_steps: null + freq_secs: null + +stats_logger: + experiment_name: ${experiment_name} + trial_name: ${trial_name} + fileroot: ${cluster.fileroot} + wandb: + mode: disabled + +perf_tracer: + experiment_name: ${experiment_name} + trial_name: ${trial_name} + fileroot: ${cluster.fileroot} + enabled: false + session_tracer: + enabled: false From 78da25b97ef5278ce80d446d60e2ff09b9b25a4b Mon Sep 17 00:00:00 2001 From: Fred Wei <20514172+WeiHaocheng@users.noreply.github.com> Date: Tue, 3 Mar 2026 04:25:26 +0000 Subject: [PATCH 08/18] Modify ScaffoldingOutput --- areal/experimental/scaffolding/controllers.py | 2 +- areal/experimental/scaffolding/core/result.py | 1 + areal/experimental/scaffolding/task.py | 3 ++- .../scaffolding/test_controllers.py | 18 +++++++++--------- .../scaffolding_core/test_self_contained.py | 10 ++++++++++ 5 files changed, 23 insertions(+), 11 deletions(-) diff --git a/areal/experimental/scaffolding/controllers.py b/areal/experimental/scaffolding/controllers.py index 3f04ea7615..1d0082cb6d 100644 --- a/areal/experimental/scaffolding/controllers.py +++ b/areal/experimental/scaffolding/controllers.py @@ -681,4 +681,4 @@ def generate(self, prompt: str, **kwargs) -> Any: yield from self.process([task], **kwargs) - return task.create_scaffolding_output() + return task.trace_results diff --git a/areal/experimental/scaffolding/core/result.py b/areal/experimental/scaffolding/core/result.py index cb0eaf6670..f452634ae7 100644 --- a/areal/experimental/scaffolding/core/result.py +++ b/areal/experimental/scaffolding/core/result.py @@ -11,6 +11,7 @@ class ScaffoldingOutput: text: str token_ids: list[int] + data: Any = None class ScaffoldingResult: diff --git a/areal/experimental/scaffolding/task.py b/areal/experimental/scaffolding/task.py index 3230608cee..f986b6d09d 100644 --- a/areal/experimental/scaffolding/task.py +++ b/areal/experimental/scaffolding/task.py @@ -177,8 +177,9 @@ def create_scaffolding_output(self) -> ScaffoldingOutput: return ScaffoldingOutput( text=self.generation_task.output_str or "", token_ids=list(self.generation_task.output_tokens or []), + data=self.trace_results, ) - return ScaffoldingOutput(text="", token_ids=[]) + return ScaffoldingOutput(text="", token_ids=[], data=self.trace_results) @dataclass diff --git a/areal/tests/experimental/scaffolding/test_controllers.py b/areal/tests/experimental/scaffolding/test_controllers.py index 4fd2de3a86..5730563293 100644 --- a/areal/tests/experimental/scaffolding/test_controllers.py +++ b/areal/tests/experimental/scaffolding/test_controllers.py @@ -50,7 +50,6 @@ TraceGenerationTask, ) - # --------------------------------------------------------------------------- # Fake / stub helpers # --------------------------------------------------------------------------- @@ -363,7 +362,9 @@ def _capture(prompt, completions, prompt_ids, completion_ids, **kw): prompt_str="", ) - task = GenerationTask(input_str="fallback prompt", input_tokens=FAKE_INPUT_TOKENS) + task = GenerationTask( + input_str="fallback prompt", input_tokens=FAKE_INPUT_TOKENS + ) list(maker.process([task])) assert received_prompts[0] == "fallback prompt" @@ -520,9 +521,7 @@ def _make_trace_maker( Also wires the class-level ``ChatTracer`` into the fake rollout controller so it can call ``before_yield`` / ``after_yield`` hooks. """ - rollout_ctrl = FakeChatRolloutController( - n_turns=n_turns, responses=responses - ) + rollout_ctrl = FakeChatRolloutController(n_turns=n_turns, responses=responses) reward_ctrl = FakeChatRewardController( rewards=rewards, default_reward=default_reward ) @@ -608,15 +607,15 @@ def test_trace_generation_task_create_from_chat_task(self): assert trace_task.trace_results is None def test_trace_generation_task_scaffolding_output(self): - """create_scaffolding_output should reflect generation_task fields.""" - gen_task = GenerationTask( - output_str="result text", output_tokens=[10, 20, 30] - ) + """create_scaffolding_output should reflect generation_task fields and trace_results.""" + gen_task = GenerationTask(output_str="result text", output_tokens=[10, 20, 30]) trace_task = TraceGenerationTask(generation_task=gen_task) + trace_task.trace_results = {"id-1": "fake_interaction"} output = trace_task.create_scaffolding_output() assert output.text == "result text" assert output.token_ids == [10, 20, 30] + assert output.data == {"id-1": "fake_interaction"} def test_trace_generation_task_scaffolding_output_empty(self): """create_scaffolding_output with no generation_task should return empty.""" @@ -624,6 +623,7 @@ def test_trace_generation_task_scaffolding_output_empty(self): output = trace_task.create_scaffolding_output() assert output.text == "" assert output.token_ids == [] + assert output.data is None def test_no_reward_tasks_when_no_traces(self): """If rollout produces no traceable outputs, reward step should be skipped.""" diff --git a/areal/tests/experimental/scaffolding_core/test_self_contained.py b/areal/tests/experimental/scaffolding_core/test_self_contained.py index 4691074e90..b89da04d14 100644 --- a/areal/tests/experimental/scaffolding_core/test_self_contained.py +++ b/areal/tests/experimental/scaffolding_core/test_self_contained.py @@ -448,6 +448,16 @@ def test_scaffolding_output(self): o = ScaffoldingOutput(text="hello", token_ids=[1, 2, 3]) assert o.text == "hello" assert o.token_ids == [1, 2, 3] + assert o.data is None + + def test_scaffolding_output_with_data(self): + from areal.experimental.scaffolding.core.result import ScaffoldingOutput + + payload = {"key": "value", "nested": [1, 2, 3]} + o = ScaffoldingOutput(text="hello", token_ids=[1, 2, 3], data=payload) + assert o.text == "hello" + assert o.token_ids == [1, 2, 3] + assert o.data is payload def test_scaffolding_result_set_output(self): from areal.experimental.scaffolding.core.result import ( From 7538b44deb131d81182ff30349b637e35d74c3d2 Mon Sep 17 00:00:00 2001 From: Fred Wei <20514172+WeiHaocheng@users.noreply.github.com> Date: Tue, 3 Mar 2026 15:29:42 +0000 Subject: [PATCH 09/18] draft for trace chat --- areal/experimental/openai/cache.py | 15 + areal/experimental/scaffolding/controllers.py | 163 ++- .../scaffolding/core/scaffolding_llm.py | 43 +- areal/experimental/scaffolding/worker.py | 14 + areal/experimental/scaffolding/workflow.py | 144 +-- .../scaffolding/test_controllers.py | 29 +- .../test_scaffolding_llm_integration.py | 1066 +++++++++++++++++ examples/scaffolding/chat_scaffolding.py | 149 ++- 8 files changed, 1436 insertions(+), 187 deletions(-) create mode 100644 areal/tests/experimental/scaffolding/test_scaffolding_llm_integration.py diff --git a/areal/experimental/openai/cache.py b/areal/experimental/openai/cache.py index 1e46bb6933..06120bbf5b 100644 --- a/areal/experimental/openai/cache.py +++ b/areal/experimental/openai/cache.py @@ -15,6 +15,21 @@ def __init__(self, *args, **kwargs): self._total_reward = 0.0 self._lock = threading.Lock() + def __deepcopy__(self, memo): + """Create a fresh empty cache. + + ``threading.Lock`` cannot be deep-copied. Controllers that hold + an ``InteractionCache`` (e.g. ``ChatTracer``) are cloned via + ``Controller.clone()`` (``copy.deepcopy``). A cloned controller + should start with an empty cache, so we simply return a new + instance. + """ + import copy + + new = InteractionCache() + memo[id(self)] = new + return new + @property def last_interaction_id(self) -> str: return next(reversed(self)) diff --git a/areal/experimental/scaffolding/controllers.py b/areal/experimental/scaffolding/controllers.py index 1d0082cb6d..77db36a6b2 100644 --- a/areal/experimental/scaffolding/controllers.py +++ b/areal/experimental/scaffolding/controllers.py @@ -7,6 +7,7 @@ Key Components: - RLVRRewardController: Controller that processes reward computation - PipelineTrajectoryMaker: Controller that composes generation and reward pipelines +- MultiTurnChatController: Controller for multi-turn chat with reflection - ChatTracer: TaskCollection for tracing multi-turn chat conversations - TraceTrajectoryMaker: Controller that traces ChatTask objects during rollout """ @@ -24,8 +25,10 @@ ChatTask, Controller, GenerationTask, + RoleMessage, Task, TaskCollection, + UserMessage, with_task_collection, ) from areal.experimental.scaffolding.task import ( @@ -86,6 +89,19 @@ def __init__(self, reward_fn: Callable[..., Any]): self.async_reward_fn = AsyncRewardWrapper(reward_fn) self.scores: list[float] | None = None + def __deepcopy__(self, memo): + """Create a new RLVRRewardController with the same reward function. + + ``AsyncRewardWrapper`` contains ``ProcessPoolExecutor`` and + ``threading.Lock`` which cannot be deep-copied. Instead of + copying those objects, we create a fresh controller that shares + the same underlying executor pool via ``AsyncRewardWrapper``'s + class-level state. + """ + new = RLVRRewardController(self.reward_fn) + memo[id(self)] = new + return new + def process(self, tasks: list[Task], **kwargs) -> Any: """Process reward tasks and compute rewards. @@ -434,6 +450,7 @@ def __init__( reward_controller: RLVRRewardController, task_data: dict[str, Any] | None = None, prompt_str: str = "", + input_tokens: list[int] | None = None, ): """Initialize the pipeline trajectory maker. @@ -447,16 +464,24 @@ def __init__( Task data containing ground truth for reward computation. prompt_str : str, optional The prompt string used for generation. + input_tokens : list[int], optional + The tokenized input IDs for the prompt. """ super().__init__() self.generation_controller = generation_controller self.reward_controller = reward_controller self.task_data = task_data if task_data is not None else {} self.prompt_str = prompt_str + self.input_tokens = input_tokens if input_tokens is not None else [] def process(self, tasks: list[Task], **kwargs) -> Any: """Process tasks through the generation and reward pipeline. + Yields task lists only for generation (worker execution). Reward + computation is done locally without yielding to workers. + Interactions are stored in ``task.customized_result_fields["interactions"]`` + for retrieval in ``generate()``. + Parameters ---------- tasks : list[Task] @@ -466,12 +491,13 @@ def process(self, tasks: list[Task], **kwargs) -> Any: Yields ------ - dict[str, InteractionWithTokenLogpReward] - Dictionary mapping task IDs to their interaction results. + list[Task] + Task lists for worker execution (generation only). """ - # Step 1: Run generation + # Step 1: Run generation (yields task lists for worker execution) yield from self.generation_controller.process(tasks, **kwargs) + # Step 2: Create reward tasks and compute rewards locally reward_tasks = [] interactions = {} @@ -491,12 +517,47 @@ def process(self, tasks: list[Task], **kwargs) -> Any: ) reward_tasks.append(reward_task) - # Step 3: Process reward tasks - yield from self.reward_controller.process(reward_tasks, **kwargs) + # Compute rewards locally (no yield to workers) + for _ in self.reward_controller.process(reward_tasks, **kwargs): + pass - # The interactions now have rewards set - # Return as the final result - yield interactions + # Store interactions on tasks for retrieval in generate() + for task in tasks: + task.customized_result_fields["interactions"] = interactions + + def generate(self, prompt: str, **kwargs) -> Any: + """Generate with the full pipeline and return interactions in output. + + Overrides the base ``Controller.generate()`` to: + 1. Set ``input_tokens`` on the task so ``to_tensor_dict()`` works. + 2. Reset ``stop`` to ``None`` so ``NativeGenerationController`` can + set it from ``sampling_params``. + 3. Return a ``ScaffoldingOutput`` with interactions in ``data``. + + Parameters + ---------- + prompt : str + The input prompt string. + **kwargs + Additional keyword arguments. + + Returns + ------- + ScaffoldingOutput + Output with ``data`` containing the interactions dict. + """ + + task = GenerationTask.create_from_prompt(prompt) + if self.input_tokens: + task.input_tokens = self.input_tokens + # Reset stop to None so NativeGenerationController can set from sampling_params + task.stop = None + + yield from self.process([task], **kwargs) + + output = task.create_scaffolding_output() + output.data = task.customized_result_fields.get("interactions") + return output def _create_interaction_from_task( self, task: GenerationTask @@ -542,6 +603,81 @@ def _create_interaction_from_task( return interaction +class MultiTurnChatController(Controller): + """Controller for multi-turn chat with reflection between turns. + + Handles the chat loop: for each turn, yields a ChatTask to the worker + for generation, then appends a reflection message for non-final turns. + + Per-episode data (``messages``, ``input_tokens``) should be set before + calling ``generate()`` or ``process()``. ``ScaffoldingLlm`` deep-copies + the controller via ``clone()`` so each request gets its own copy. + + Parameters + ---------- + generation_controller : Controller + The controller for text generation (e.g., NativeGenerationController). + max_turns : int + Maximum number of chat turns per episode. + reflection_message : str + Message appended after each non-final turn to prompt retry. + tokenizer : Any + Tokenizer for encoding the final output to token IDs. + messages : list[dict], optional + The original chat messages to create the ChatTask from (set per-episode). + input_tokens : list[int], optional + The tokenized input IDs for the original prompt (set per-episode). + """ + + def __init__( + self, + generation_controller: Controller, + max_turns: int = 2, + reflection_message: str = "", + tokenizer: Any = None, + messages: list[dict] | None = None, + input_tokens: list[int] | None = None, + ): + super().__init__() + self.generation_controller = generation_controller + self.max_turns = max_turns + self.reflection_message = reflection_message + self.tokenizer = tokenizer + self.messages = messages if messages is not None else [] + self.input_tokens = input_tokens if input_tokens is not None else [] + + def process(self, tasks: list[Task], **kwargs) -> Any: + """Run multi-turn chat generation. + + Creates a ChatTask from the stored messages and yields it to the + worker for each turn. Between turns, appends a reflection message. + + Parameters + ---------- + tasks : list[Task] + Ignored; the ChatTask is created from ``self.messages``. + **kwargs + Additional keyword arguments. + + Yields + ------ + list[Task] + Task lists for worker execution (one per turn). + """ + role_messages = [RoleMessage.from_dict(m) for m in self.messages] + chat_task = ChatTask.create_from_messages(role_messages) + # Reset stop so NativeGenerationController can set from sampling_params + chat_task.stop = None + if self.input_tokens: + chat_task.input_tokens = self.input_tokens + + for turn in range(self.max_turns): + yield from self.generation_controller.process([chat_task], **kwargs) + + if turn < self.max_turns - 1: + chat_task.add_message(UserMessage(self.reflection_message)) + + @with_task_collection("chat_tracer", ChatTracer) class TraceTrajectoryMaker(Controller): """Controller that traces ChatTask objects during rollout using ChatTracer. @@ -646,9 +782,10 @@ def process(self, tasks: list[Task], **kwargs) -> Any: for interaction_id, interaction in trace_results.items() ] - # Run reward computation + # Run reward computation locally (no yield to workers) if reward_tasks: - yield from self.reward_controller.process(reward_tasks, **kwargs) + for _ in self.reward_controller.process(reward_tasks, **kwargs): + pass # Update trace_results with computed rewards for reward_task in reward_tasks: @@ -674,11 +811,11 @@ def generate(self, prompt: str, **kwargs) -> Any: Returns ------- - Any - The scaffolding output. + ScaffoldingOutput + Output with trace results in ``data``. """ task = TraceGenerationTask.create_from_prompt(prompt) yield from self.process([task], **kwargs) - return task.trace_results + return task.create_scaffolding_output() diff --git a/areal/experimental/scaffolding/core/scaffolding_llm.py b/areal/experimental/scaffolding/core/scaffolding_llm.py index 4e6018e145..6af58e5ccf 100644 --- a/areal/experimental/scaffolding/core/scaffolding_llm.py +++ b/areal/experimental/scaffolding/core/scaffolding_llm.py @@ -9,11 +9,15 @@ from dataclasses import dataclass from typing import Any +from areal.utils import logging as areal_logging + from .controller import Controller, ParallelProcess from .result import ScaffoldingResult from .task import Task from .worker import Worker +_logger = areal_logging.getLogger("ScaffoldingLlm") + @dataclass(frozen=True) class ScaffoldingRequest: @@ -65,15 +69,47 @@ def _get_loop(self): return asyncio.new_event_loop() return None + def _schedule_on_loop(self, coro): + """Schedule a coroutine on self.loop. + + Uses create_task when called from within the event loop thread + (to avoid deadlocks with run_coroutine_threadsafe on uvloop), + and falls back to run_coroutine_threadsafe for cross-thread calls. + """ + try: + running = asyncio.get_running_loop() + except RuntimeError: + running = None + + if running is self.loop: + _logger.info("_schedule_on_loop: same loop, create_task for %s", coro.__name__) + self.loop.create_task(coro) + elif running is not None: + _logger.warning( + "_schedule_on_loop: different loop! running=%s self.loop=%s", + id(running), id(self.loop), + ) + asyncio.run_coroutine_threadsafe(coro, self.loop) + else: + _logger.info("_schedule_on_loop: no running loop, run_coroutine_threadsafe for %s", coro.__name__) + asyncio.run_coroutine_threadsafe(coro, self.loop) + async def _handle_controller_generator( self, gen: Generator, request: ScaffoldingRequest = None ): """Handle a controller generator, processing tasks and parallel processes.""" + step = 0 for obj in gen: if isinstance(obj, ParallelProcess): await self._handle_parallel_process(obj, request) else: + step += 1 + _logger.info( + "Dispatching task list step=%d, tasks=%d, types=%s", + step, len(obj), [type(t).__name__ for t in obj], + ) await self._handle_task_list(obj, request) + _logger.info("Task list step=%d completed.", step) async def _handle_task_list( self, tasks: list[Task], request: ScaffoldingRequest = None @@ -157,12 +193,13 @@ async def _handle_event_loop(self): async def _main_loop_async_func(self): """Main async loop function.""" + _logger.info("_main_loop_async_func STARTED (loop=%s)", id(asyncio.get_running_loop())) handle_event_task = asyncio.create_task(self._handle_event_loop()) await handle_event_task self.main_loop_stop_event.set() def _run_main_loop_coroutine(self): - asyncio.run_coroutine_threadsafe(self._main_loop_async_func(), self.loop) + self._schedule_on_loop(self._main_loop_async_func()) def _run_main_loop_thread(self): def main_loop_thread(): @@ -190,7 +227,7 @@ async def put_request(): else: await self.task_queue.put(request) - asyncio.run_coroutine_threadsafe(put_request(), self.loop) + self._schedule_on_loop(put_request()) return result @@ -223,7 +260,7 @@ async def stop_task_on_loop(): for worker in self.workers.values(): await worker.async_shutdown() - asyncio.run_coroutine_threadsafe(stop_task_on_loop(), self.loop) + self._schedule_on_loop(stop_task_on_loop()) if self.own_loop: self.main_loop_thread.join() diff --git a/areal/experimental/scaffolding/worker.py b/areal/experimental/scaffolding/worker.py index 2686a7c090..3b0c9076d8 100644 --- a/areal/experimental/scaffolding/worker.py +++ b/areal/experimental/scaffolding/worker.py @@ -18,10 +18,13 @@ OpenaiWorker, TaskStatus, ) +from areal.utils import logging if TYPE_CHECKING: from areal.engine.sglang_remote import RemoteSGLangEngine +worker_logger = logging.getLogger("SGLangWorker") + class SGLangWorker(OpenaiWorker): """Worker that wraps an SGLang engine for scaffolding. @@ -71,7 +74,13 @@ async def chat_handler(self, task: ChatTask) -> TaskStatus: params["tools"] = [tool.to_dict() for tool in task.tools] try: + worker_logger.info( + "Sending chat request to %s (messages=%d) ...", + self.async_client.base_url, + len(params.get("messages", [])), + ) response = await self.async_client.chat.completions.create(**params) + worker_logger.info("Chat response received.") # Store the completion in the task for tracing task.completion = response @@ -127,7 +136,12 @@ async def generation_handler(self, task: GenerationTask) -> TaskStatus: params["prompt"] = task.input_str try: + worker_logger.info( + "Sending generation request to %s ...", + self.async_client.base_url, + ) response = await self.async_client.completions.create(**params) + worker_logger.info("Generation response received.") task.output_str = response.choices[0].text if hasattr(response.choices[0], "token_ids"): diff --git a/areal/experimental/scaffolding/workflow.py b/areal/experimental/scaffolding/workflow.py index 5a08f42586..19b46ea2b5 100644 --- a/areal/experimental/scaffolding/workflow.py +++ b/areal/experimental/scaffolding/workflow.py @@ -19,12 +19,10 @@ import torch from transformers import PreTrainedTokenizerFast -from areal import workflow_context from areal.api.cli_args import GenerationHyperparameters from areal.api.engine_api import InferenceEngine from areal.api.workflow_api import RolloutWorkflow from areal.experimental.scaffolding._compat import ( - GenerationTask, NativeGenerationController, ScaffoldingLlm, ) @@ -32,11 +30,9 @@ PipelineTrajectoryMaker, RLVRRewardController, ) -from areal.experimental.scaffolding.task import RLVRRewardTask from areal.experimental.scaffolding.worker import SGLangWorker -from areal.utils import logging, stats_tracker +from areal.utils import logging from areal.utils.dynamic_import import import_from_string -from areal.utils.perf_tracer import session_context, trace_session logger = logging.getLogger("ScaffoldingWorkflow") @@ -126,34 +122,7 @@ def build_scaffolding_llm(self, engine: InferenceEngine) -> ScaffoldingLlm: ScaffoldingLlm The constructed ScaffoldingLlm instance. """ - self.gen_controller = NativeGenerationController() - self.reward_controller = RLVRRewardController(self.reward_fn) - self.trajectory_maker = PipelineTrajectoryMaker( - self.gen_controller, self.reward_controller - ) - return ScaffoldingLlm( - self.trajectory_maker, - {NativeGenerationController.WorkerTag.GENERATION: self.worker}, - ) - - async def _generate_via_worker( - self, prompt_str: str, input_ids: list[int] - ) -> GenerationTask: - """Run generation through scaffolding Worker (SGLang OpenAI API). - - Parameters - ---------- - prompt_str : str - The prompt string. - input_ids : list[int] - The tokenized input IDs. - - Returns - ------- - GenerationTask - Completed task with output_str and output_tokens. - """ - # Build generation params for SGLang completions API + # Convert gconfig to sampling params for NativeGenerationController stop_strings = [] if self.gconfig.stop_token_ids: for tid in self.gconfig.stop_token_ids: @@ -161,72 +130,34 @@ async def _generate_via_worker( if decoded: stop_strings.append(decoded) - response = await self.worker.async_client.completions.create( - model=self.worker.model, - prompt=prompt_str, - max_tokens=self.gconfig.max_new_tokens, - temperature=self.gconfig.temperature or 1.0, - stop=stop_strings or None, - ) + sampling_params = { + "max_tokens": self.gconfig.max_new_tokens, + "temperature": self.gconfig.temperature or 1.0, + } + if stop_strings: + sampling_params["stop"] = stop_strings - output_str = response.choices[0].text - # Tokenize to get output token IDs - output_token_ids = self.tokenizer.encode(output_str, add_special_tokens=False) - - # Package as a GenerationTask (scaffolding data structure) - gen_task = GenerationTask( - input_str=prompt_str, - input_tokens=input_ids, - output_str=output_str, - output_tokens=output_token_ids, - finish_reason=response.choices[0].finish_reason, + self.gen_controller = NativeGenerationController( + sampling_params=sampling_params ) - return gen_task - - @trace_session("reward") - async def _compute_rewards_via_controller( - self, - gen_task: GenerationTask, - prompt_str: str, - task_data: dict[str, Any], - ) -> float: - """Compute reward via scaffolding RLVRRewardController.""" - reward_task = RLVRRewardTask( - prompt_str=prompt_str, - completion_str=gen_task.output_str or "", - input_tokens=list(gen_task.input_tokens or []), - output_tokens=list(gen_task.output_tokens or []), - task_data=task_data, + self.reward_controller = RLVRRewardController(self.reward_fn) + self.trajectory_maker = PipelineTrajectoryMaker( + self.gen_controller, self.reward_controller ) - for _ in self.reward_controller.process([reward_task]): - pass - return float(reward_task.reward) - - @session_context() - async def _collect_samples( - self, - prompt_str: str, - input_ids: list[int], - task_data: dict[str, Any], - ) -> tuple[GenerationTask, float]: - """Generate via Worker, compute reward via Controller.""" - gen_task = await self._generate_via_worker(prompt_str, input_ids) - - reward = await self._compute_rewards_via_controller( - gen_task, prompt_str, task_data + return ScaffoldingLlm( + self.trajectory_maker, + {NativeGenerationController.WorkerTag.GENERATION: self.worker}, ) - stats_tracker.get(workflow_context.stat_scope()).scalar(reward=reward) - - return gen_task, reward async def arun_episode( self, engine: InferenceEngine, data: dict[str, Any] ) -> dict[str, torch.Tensor]: """Run a single episode via scaffolding pipeline. - 1. Generation: SGLangWorker -> SGLang completions API - 2. Reward: RLVRRewardController.process() - 3. Output: tensor dict for PPO training + Delegates the full episode (generation + reward) to + ``self.scaffolding_llm``, which wraps a ``PipelineTrajectoryMaker``. + The result is an ``InteractionWithTokenLogpReward`` whose + ``to_tensor_dict()`` produces the training tensors. Note: logprobs are placeholders (0.0). Set ``recompute_logprob: true`` in actor config so the training engine computes exact logprobs. @@ -257,23 +188,18 @@ async def arun_episode( ) prompt_str = self.tokenizer.decode(input_ids) - # Scaffolding pipeline: Worker (generate) + Controller (reward) - gen_task, reward = await self._collect_samples(prompt_str, input_ids, data) - - # Build tensor dict for PPO training - output_tokens = list(gen_task.output_tokens or []) - seq = input_ids + output_tokens - # Placeholder logprobs — recompute_logprob=true will replace these - logprobs = [0.0] * len(seq) - loss_mask = [0] * len(input_ids) + [1] * len(output_tokens) - versions = [-1] * len(seq) - - res = { - "input_ids": torch.tensor(seq, dtype=torch.int32), - "loss_mask": torch.tensor(loss_mask, dtype=torch.int32), - "logprobs": torch.tensor(logprobs, dtype=torch.float32), - "versions": torch.tensor(versions, dtype=torch.int32), - "attention_mask": torch.ones(len(seq), dtype=torch.bool), - "rewards": torch.tensor(reward, dtype=torch.float32), - } - return {k: v.unsqueeze(0) for k, v in res.items()} + # Configure per-episode data on trajectory maker + # (clone() in scaffolding_llm will deep-copy these) + self.trajectory_maker.task_data = data + self.trajectory_maker.prompt_str = prompt_str + self.trajectory_maker.input_tokens = input_ids + + # Run full pipeline via scaffolding_llm + result = self.scaffolding_llm.generate_async(prompt_str) + await result + + # Extract interaction and convert to tensor dict + scaffolding_output = result.outputs[0] + interactions = scaffolding_output.data + interaction = next(iter(interactions.values())) + return interaction.to_tensor_dict() diff --git a/areal/tests/experimental/scaffolding/test_controllers.py b/areal/tests/experimental/scaffolding/test_controllers.py index 5730563293..31ce716bc9 100644 --- a/areal/tests/experimental/scaffolding/test_controllers.py +++ b/areal/tests/experimental/scaffolding/test_controllers.py @@ -241,9 +241,11 @@ def test_basic_pipeline_correct_answer(self): task = GenerationTask(input_str=FAKE_PROMPT_STR, input_tokens=FAKE_INPUT_TOKENS) results = list(maker.process([task])) - # Three yields: generation, reward, interactions dict - assert len(results) == 3 - interactions = results[-1] + # Only generation yields (reward computed locally, no dict yield) + assert len(results) == 1 + + # Interactions stored on task + interactions = task.customized_result_fields["interactions"] assert isinstance(interactions, dict) assert len(interactions) == 1 @@ -264,9 +266,9 @@ def test_basic_pipeline_wrong_answer(self): ) task = GenerationTask(input_str=FAKE_PROMPT_STR, input_tokens=FAKE_INPUT_TOKENS) - results = list(maker.process([task])) + list(maker.process([task])) - interactions = results[-1] + interactions = task.customized_result_fields["interactions"] interaction = list(interactions.values())[0] assert interaction.reward == 0.0 @@ -286,9 +288,10 @@ def test_pipeline_multiple_tasks(self): GenerationTask(input_str=FAKE_PROMPT_STR, input_tokens=FAKE_INPUT_TOKENS), GenerationTask(input_str="Another prompt", input_tokens=[111, 112]), ] - results = list(maker.process(tasks)) + list(maker.process(tasks)) - interactions = results[-1] + # Both tasks should have the same interactions dict + interactions = tasks[0].customized_result_fields["interactions"] assert len(interactions) == 2 for interaction in interactions.values(): assert interaction.reward == 1.0 @@ -311,9 +314,9 @@ def test_pipeline_interaction_has_model_response(self): ) task = GenerationTask(input_str=FAKE_PROMPT_STR, input_tokens=FAKE_INPUT_TOKENS) - results = list(maker.process([task])) + list(maker.process([task])) - interaction = list(results[-1].values())[0] + interaction = list(task.customized_result_fields["interactions"].values())[0] mr = interaction.model_response assert isinstance(mr, ModelResponse) assert mr.input_tokens == FAKE_INPUT_TOKENS @@ -406,9 +409,9 @@ def test_pipeline_default_logprobs_and_versions(self): ) task = GenerationTask(input_str=FAKE_PROMPT_STR, input_tokens=FAKE_INPUT_TOKENS) - results = list(maker.process([task])) + list(maker.process([task])) - interaction = list(results[-1].values())[0] + interaction = list(task.customized_result_fields["interactions"].values())[0] mr = interaction.model_response # Placeholders: [0.0] * output_len assert mr.output_logprobs == [0.0, 0.0] @@ -872,8 +875,8 @@ def test_pipeline_e2e_tensor_dict_compatible(self): ) task = GenerationTask(input_str=FAKE_PROMPT_STR, input_tokens=FAKE_INPUT_TOKENS) - results = list(maker.process([task])) - interaction = list(results[-1].values())[0] + list(maker.process([task])) + interaction = list(task.customized_result_fields["interactions"].values())[0] td = interaction.to_tensor_dict() assert "input_ids" in td diff --git a/areal/tests/experimental/scaffolding/test_scaffolding_llm_integration.py b/areal/tests/experimental/scaffolding/test_scaffolding_llm_integration.py new file mode 100644 index 0000000000..3b337ed4ea --- /dev/null +++ b/areal/tests/experimental/scaffolding/test_scaffolding_llm_integration.py @@ -0,0 +1,1066 @@ +"""Integration tests for ScaffoldingLlm with full controller pipelines. + +Tests verify that ScaffoldingLlm can correctly drive: +1. PipelineTrajectoryMaker (single-turn generation + reward) +2. TraceTrajectoryMaker + MultiTurnChatController (multi-turn chat + tracing + reward) + +These tests use a FakeChatWorker that handles ChatTask and GenerationTask +via the Worker.task_handlers dispatch mechanism, simulating an LLM backend. +""" + +from __future__ import annotations + +import asyncio +from unittest.mock import MagicMock + +import pytest + +from areal.experimental.scaffolding._compat import ( + AssistantMessage, + ChatTask, + GenerationTask, + NativeGenerationController, + ScaffoldingLlm, + TaskStatus, + Worker, +) +from areal.experimental.scaffolding.controllers import ( + MultiTurnChatController, + PipelineTrajectoryMaker, + RLVRRewardController, + TraceTrajectoryMaker, +) + +# --------------------------------------------------------------------------- +# Test constants +# --------------------------------------------------------------------------- + +FAKE_INPUT_TOKENS = [101, 102, 103] +FAKE_OUTPUT_TOKENS = [201, 202, 203, 204] +FAKE_OUTPUT_STR = "42" +FAKE_PROMPT_STR = "What is the answer to life?" + + +def _simple_reward_fn( + prompt: str, + completions: str, + prompt_ids: list[int], + completion_ids: list[int], + **kwargs, +) -> float: + """Deterministic reward: 1.0 if completion contains the answer, else 0.0.""" + answer = kwargs.get("answer", "") + return 1.0 if answer and answer in completions else 0.0 + + +# --------------------------------------------------------------------------- +# Fake Worker +# --------------------------------------------------------------------------- + + +def _make_fake_completion(completion_id: str = "cmpl-001") -> MagicMock: + """Create a minimal fake ChatCompletion object.""" + completion = MagicMock() + completion.id = completion_id + completion.created = 1000 + completion.choices = [MagicMock()] + completion.choices[0].message.content = FAKE_OUTPUT_STR + completion.choices[0].finish_reason = "stop" + return completion + + +_chat_handler_call_count = 0 + + +class FakeChatWorker(Worker): + """Worker that handles ChatTask and GenerationTask without any backend. + + For ChatTask: appends an AssistantMessage and sets completion/tokens. + For GenerationTask: fills output_str and output_tokens. + """ + + def __init__( + self, + response_text: str = FAKE_OUTPUT_STR, + output_tokens: list[int] | None = None, + ): + self.response_text = response_text + self.output_tokens = output_tokens or FAKE_OUTPUT_TOKENS + + async def _handle_generation_task(self, task: GenerationTask) -> TaskStatus: + task.output_str = self.response_text + task.output_tokens = list(self.output_tokens) + if task.input_tokens is None: + task.input_tokens = FAKE_INPUT_TOKENS + task.finish_reason = "stop" + return TaskStatus.SUCCESS + + async def _handle_chat_task(self, task: ChatTask) -> TaskStatus: + global _chat_handler_call_count + _chat_handler_call_count += 1 + + completion_id = f"cmpl-{_chat_handler_call_count:03d}" + task.completion = _make_fake_completion(completion_id) + task.completion.choices[0].message.content = self.response_text + task.output_tokens = list(self.output_tokens) + if task.input_tokens is None: + task.input_tokens = FAKE_INPUT_TOKENS + task.finish_reason = "stop" + + # Mimic what OpenaiWorker.chat_handler does: append assistant message + task.messages.append(AssistantMessage(content=self.response_text)) + return TaskStatus.SUCCESS + + task_handlers = { + GenerationTask: _handle_generation_task, + ChatTask: _handle_chat_task, + } + + +# --------------------------------------------------------------------------- +# Fixtures +# --------------------------------------------------------------------------- + + +@pytest.fixture(autouse=True) +def _reset_chat_handler_count(): + """Reset the global call counter before each test.""" + global _chat_handler_call_count + _chat_handler_call_count = 0 + yield + + +@pytest.fixture(autouse=True) +def _reset_shared_tracer(): + """Reset the class-level ChatTracer before each test.""" + yield + tracer = TraceTrajectoryMaker.task_collections.get("chat_tracer") + if tracer is not None: + tracer.clear() + + +# =========================================================================== +# PipelineTrajectoryMaker + ScaffoldingLlm +# =========================================================================== + + +class TestPipelineViaScaffoldingLlm: + """Test PipelineTrajectoryMaker running through ScaffoldingLlm.""" + + def test_single_generation_sync(self): + """ScaffoldingLlm.generate() should produce a result with interactions.""" + worker = FakeChatWorker(response_text="The answer is 42.") + reward_ctrl = RLVRRewardController(_simple_reward_fn) + gen_ctrl = NativeGenerationController( + sampling_params={"max_tokens": 100, "temperature": 1.0} + ) + trajectory_maker = PipelineTrajectoryMaker( + generation_controller=gen_ctrl, + reward_controller=reward_ctrl, + task_data={"answer": "42"}, + prompt_str=FAKE_PROMPT_STR, + input_tokens=FAKE_INPUT_TOKENS, + ) + + llm = ScaffoldingLlm( + trajectory_maker, + {NativeGenerationController.WorkerTag.GENERATION: worker}, + ) + try: + result = llm.generate(FAKE_PROMPT_STR) + + assert result is not None + assert result._done is True + output = result.outputs[0] + assert output.data is not None + + interactions = output.data + assert len(interactions) == 1 + interaction = list(interactions.values())[0] + assert interaction.reward == 1.0 + assert interaction.model_response is not None + assert interaction.model_response.input_tokens == FAKE_INPUT_TOKENS + finally: + llm.shutdown() + + def test_single_generation_async(self): + """ScaffoldingLlm.generate_async() + await should work.""" + worker = FakeChatWorker(response_text="42") + reward_ctrl = RLVRRewardController(_simple_reward_fn) + gen_ctrl = NativeGenerationController(sampling_params={"max_tokens": 100}) + trajectory_maker = PipelineTrajectoryMaker( + generation_controller=gen_ctrl, + reward_controller=reward_ctrl, + task_data={"answer": "42"}, + prompt_str=FAKE_PROMPT_STR, + input_tokens=FAKE_INPUT_TOKENS, + ) + + llm = ScaffoldingLlm( + trajectory_maker, + {NativeGenerationController.WorkerTag.GENERATION: worker}, + ) + try: + result = llm.generate_async(FAKE_PROMPT_STR) + # Use result.result() which blocks via the event loop + result.result(timeout=10.0) + + assert result._done is True + output = result.outputs[0] + interactions = output.data + assert len(interactions) == 1 + interaction = list(interactions.values())[0] + assert interaction.reward == 1.0 + finally: + llm.shutdown() + + def test_batch_generation(self): + """ScaffoldingLlm.generate() with a list of prompts should work.""" + worker = FakeChatWorker(response_text="42") + reward_ctrl = RLVRRewardController(_simple_reward_fn) + gen_ctrl = NativeGenerationController(sampling_params={"max_tokens": 50}) + trajectory_maker = PipelineTrajectoryMaker( + generation_controller=gen_ctrl, + reward_controller=reward_ctrl, + task_data={"answer": "42"}, + prompt_str=FAKE_PROMPT_STR, + input_tokens=FAKE_INPUT_TOKENS, + ) + + llm = ScaffoldingLlm( + trajectory_maker, + {NativeGenerationController.WorkerTag.GENERATION: worker}, + ) + try: + results = llm.generate(["prompt1", "prompt2", "prompt3"]) + assert len(results) == 3 + for result in results: + assert result._done is True + assert result.outputs[0].data is not None + finally: + llm.shutdown() + + +# =========================================================================== +# MultiTurnChatController + TraceTrajectoryMaker + ScaffoldingLlm +# =========================================================================== + + +class TestMultiTurnViaScaffoldingLlm: + """Test MultiTurnChatController + TraceTrajectoryMaker via ScaffoldingLlm. + + This reproduces the exact architecture used by ChatScaffoldingWorkflow. + """ + + @staticmethod + def _build_llm( + response_text: str = FAKE_OUTPUT_STR, + max_turns: int = 2, + reward_fn=None, + ) -> tuple[ScaffoldingLlm, MultiTurnChatController, TraceTrajectoryMaker]: + """Build the full ChatScaffoldingWorkflow pipeline.""" + worker = FakeChatWorker(response_text=response_text) + gen_ctrl = NativeGenerationController( + sampling_params={"max_tokens": 100, "temperature": 1.0} + ) + reward_ctrl = RLVRRewardController(reward_fn or _simple_reward_fn) + multi_turn_ctrl = MultiTurnChatController( + generation_controller=gen_ctrl, + max_turns=max_turns, + reflection_message="Try again.", + ) + trace_maker = TraceTrajectoryMaker( + rollout_controller=multi_turn_ctrl, + reward_controller=reward_ctrl, + ) + llm = ScaffoldingLlm( + trace_maker, + {NativeGenerationController.WorkerTag.GENERATION: worker}, + ) + return llm, multi_turn_ctrl, trace_maker + + def test_single_turn_sync(self): + """Single-turn chat via ScaffoldingLlm.generate() should work.""" + llm, multi_turn_ctrl, _ = self._build_llm(response_text="42", max_turns=1) + try: + # Set per-episode data (like ChatScaffoldingWorkflow.arun_episode does) + multi_turn_ctrl.messages = [{"role": "user", "content": "What is 6*7?"}] + multi_turn_ctrl.input_tokens = FAKE_INPUT_TOKENS + + result = llm.generate(FAKE_PROMPT_STR) + + assert result is not None + assert result._done is True + output = result.outputs[0] + # TraceTrajectoryMaker.generate() returns ScaffoldingOutput with + # trace_results in data + assert output.data is not None or output.text is not None + finally: + llm.shutdown() + + def test_multi_turn_sync(self): + """Multi-turn (2-turn) chat via ScaffoldingLlm.generate() should work.""" + llm, multi_turn_ctrl, _ = self._build_llm(response_text="42", max_turns=2) + try: + multi_turn_ctrl.messages = [{"role": "user", "content": "Solve: 6*7"}] + multi_turn_ctrl.input_tokens = FAKE_INPUT_TOKENS + + result = llm.generate(FAKE_PROMPT_STR) + + assert result is not None + assert result._done is True + + # Verify the worker was called twice (2 turns) + global _chat_handler_call_count + assert _chat_handler_call_count >= 2 + finally: + llm.shutdown() + + def test_multi_turn_async_result(self): + """generate_async() + result() should complete for multi-turn.""" + llm, multi_turn_ctrl, _ = self._build_llm(response_text="42", max_turns=2) + try: + multi_turn_ctrl.messages = [{"role": "user", "content": "Solve: 6*7"}] + multi_turn_ctrl.input_tokens = FAKE_INPUT_TOKENS + + result = llm.generate_async(FAKE_PROMPT_STR) + result.result(timeout=10.0) + + assert result._done is True + finally: + llm.shutdown() + + def test_trace_results_available(self): + """Trace results should be accessible via ScaffoldingOutput.data.""" + llm, multi_turn_ctrl, _ = self._build_llm(response_text="42", max_turns=1) + try: + multi_turn_ctrl.messages = [{"role": "user", "content": "What is 6*7?"}] + multi_turn_ctrl.input_tokens = FAKE_INPUT_TOKENS + + result = llm.generate(FAKE_PROMPT_STR) + output = result.outputs[0] + + # TraceTrajectoryMaker stores trace_results in + # TraceGenerationTask.trace_results, which is returned via + # create_scaffolding_output().data + trace_data = output.data + if trace_data is not None: + # If tracing worked, we should have at least one interaction + assert len(trace_data) >= 1 + finally: + llm.shutdown() + + def test_multiple_concurrent_requests(self): + """Multiple concurrent requests should all complete.""" + llm, multi_turn_ctrl, _ = self._build_llm(response_text="42", max_turns=1) + try: + multi_turn_ctrl.messages = [{"role": "user", "content": "What is 6*7?"}] + multi_turn_ctrl.input_tokens = FAKE_INPUT_TOKENS + + results = llm.generate([f"prompt_{i}" for i in range(5)]) + + assert len(results) == 5 + for result in results: + assert result._done is True + finally: + llm.shutdown() + + def test_clone_isolation(self): + """Each request should get an independent controller clone.""" + call_messages = [] + + def _capture_reward(prompt, completions, prompt_ids, completion_ids, **kw): + call_messages.append(completions) + return 1.0 + + llm, multi_turn_ctrl, _ = self._build_llm( + response_text="42", max_turns=1, reward_fn=_capture_reward + ) + try: + multi_turn_ctrl.messages = [{"role": "user", "content": "Q1"}] + multi_turn_ctrl.input_tokens = FAKE_INPUT_TOKENS + + # Two sequential requests + r1 = llm.generate("prompt1") + r2 = llm.generate("prompt2") + + assert r1._done is True + assert r2._done is True + finally: + llm.shutdown() + + +# =========================================================================== +# Edge cases +# =========================================================================== + + +class TestEdgeCases: + """Edge cases and error scenarios.""" + + def test_empty_messages(self): + """Controller with empty messages should not crash.""" + worker = FakeChatWorker(response_text="hello") + gen_ctrl = NativeGenerationController(sampling_params={"max_tokens": 10}) + reward_ctrl = RLVRRewardController(_simple_reward_fn) + multi_turn_ctrl = MultiTurnChatController( + generation_controller=gen_ctrl, + max_turns=1, + reflection_message="retry", + messages=[], + input_tokens=[], + ) + trace_maker = TraceTrajectoryMaker( + rollout_controller=multi_turn_ctrl, + reward_controller=reward_ctrl, + ) + llm = ScaffoldingLlm( + trace_maker, + {NativeGenerationController.WorkerTag.GENERATION: worker}, + ) + try: + result = llm.generate("test prompt") + assert result._done is True + finally: + llm.shutdown() + + def test_shutdown_is_safe(self): + """Calling shutdown() should not raise.""" + worker = FakeChatWorker() + gen_ctrl = NativeGenerationController(sampling_params={"max_tokens": 10}) + reward_ctrl = RLVRRewardController(_simple_reward_fn) + trajectory_maker = PipelineTrajectoryMaker( + generation_controller=gen_ctrl, + reward_controller=reward_ctrl, + ) + llm = ScaffoldingLlm( + trajectory_maker, + {NativeGenerationController.WorkerTag.GENERATION: worker}, + ) + llm.shutdown() + # No exception = success + + +# =========================================================================== +# Async context tests (reproducing the training framework scenario) +# =========================================================================== + + +class TestAsyncContextScaffoldingLlm: + """Tests that run ScaffoldingLlm from within an existing asyncio event loop. + + This reproduces the actual deployment scenario: the rollout framework's + arun_episode is called from within an asyncio context that already has a + running event loop. ScaffoldingLlm._get_loop() will detect the running + loop (own_loop=False) and schedule its main loop on it. + """ + + @pytest.mark.asyncio + async def test_pipeline_in_async_context(self): + """PipelineTrajectoryMaker should work when called from async context.""" + worker = FakeChatWorker(response_text="The answer is 42.") + gen_ctrl = NativeGenerationController(sampling_params={"max_tokens": 100}) + reward_ctrl = RLVRRewardController(_simple_reward_fn) + trajectory_maker = PipelineTrajectoryMaker( + generation_controller=gen_ctrl, + reward_controller=reward_ctrl, + task_data={"answer": "42"}, + prompt_str=FAKE_PROMPT_STR, + input_tokens=FAKE_INPUT_TOKENS, + ) + + llm = ScaffoldingLlm( + trajectory_maker, + {NativeGenerationController.WorkerTag.GENERATION: worker}, + ) + try: + result = llm.generate_async(FAKE_PROMPT_STR) + await asyncio.wait_for(result, timeout=10.0) + + assert result._done is True + output = result.outputs[0] + assert output.data is not None + interactions = output.data + assert len(interactions) == 1 + interaction = list(interactions.values())[0] + assert interaction.reward == 1.0 + finally: + llm.shutdown() + + @pytest.mark.asyncio + async def test_multi_turn_in_async_context(self): + """MultiTurnChatController + TraceTrajectoryMaker in async context.""" + worker = FakeChatWorker(response_text="42") + gen_ctrl = NativeGenerationController(sampling_params={"max_tokens": 100}) + reward_ctrl = RLVRRewardController(_simple_reward_fn) + multi_turn_ctrl = MultiTurnChatController( + generation_controller=gen_ctrl, + max_turns=2, + reflection_message="Try again.", + messages=[{"role": "user", "content": "What is 6*7?"}], + input_tokens=FAKE_INPUT_TOKENS, + ) + trace_maker = TraceTrajectoryMaker( + rollout_controller=multi_turn_ctrl, + reward_controller=reward_ctrl, + ) + + llm = ScaffoldingLlm( + trace_maker, + {NativeGenerationController.WorkerTag.GENERATION: worker}, + ) + try: + result = llm.generate_async(FAKE_PROMPT_STR) + await asyncio.wait_for(result, timeout=10.0) + + assert result._done is True + finally: + llm.shutdown() + + @pytest.mark.asyncio + async def test_multiple_async_requests(self): + """Multiple concurrent async requests from within async context.""" + worker = FakeChatWorker(response_text="42") + gen_ctrl = NativeGenerationController(sampling_params={"max_tokens": 100}) + reward_ctrl = RLVRRewardController(_simple_reward_fn) + trajectory_maker = PipelineTrajectoryMaker( + generation_controller=gen_ctrl, + reward_controller=reward_ctrl, + task_data={"answer": "42"}, + prompt_str=FAKE_PROMPT_STR, + input_tokens=FAKE_INPUT_TOKENS, + ) + + llm = ScaffoldingLlm( + trajectory_maker, + {NativeGenerationController.WorkerTag.GENERATION: worker}, + ) + try: + # Launch multiple requests concurrently (like the rollout framework) + results = [] + for i in range(5): + results.append(llm.generate_async(f"prompt_{i}")) + + # Await all concurrently + await asyncio.wait_for( + asyncio.gather(*[r.aresult() for r in results]), + timeout=10.0, + ) + + for result in results: + assert result._done is True + assert result.outputs[0].data is not None + finally: + llm.shutdown() + + +# =========================================================================== +# Uvloop thread tests (reproducing the exact AsyncTaskRunner scenario) +# =========================================================================== + + +class TestUvloopThreadScaffoldingLlm: + """Tests that reproduce the exact production deployment architecture. + + AsyncTaskRunner runs a uvloop in a separate thread. Multiple coroutines + (one per arun_episode) run concurrently on this loop. They share one + ScaffoldingLlm instance. The ScaffoldingLlm is lazily initialized inside + the first coroutine (so it captures the uvloop as its event loop). + + This is the exact scenario where the deadlock was observed. + """ + + @staticmethod + def _build_shared_components( + response_text: str = FAKE_OUTPUT_STR, + max_turns: int = 2, + ): + """Build shared components (simulating ChatScaffoldingWorkflow.__init__).""" + worker = FakeChatWorker(response_text=response_text) + gen_ctrl = NativeGenerationController( + sampling_params={"max_tokens": 100, "temperature": 1.0} + ) + reward_ctrl = RLVRRewardController(_simple_reward_fn) + multi_turn_ctrl = MultiTurnChatController( + generation_controller=gen_ctrl, + max_turns=max_turns, + reflection_message="Try again.", + ) + trace_maker = TraceTrajectoryMaker( + rollout_controller=multi_turn_ctrl, + reward_controller=reward_ctrl, + ) + return worker, multi_turn_ctrl, trace_maker + + def test_uvloop_lazy_init_single_request(self): + """ScaffoldingLlm lazily initialized on uvloop should handle one request.""" + import threading + + import uvloop + + worker, multi_turn_ctrl, trace_maker = self._build_shared_components( + max_turns=1 + ) + multi_turn_ctrl.messages = [{"role": "user", "content": "What is 6*7?"}] + multi_turn_ctrl.input_tokens = FAKE_INPUT_TOKENS + + result_holder = {} + error_holder = {} + + async def run_on_uvloop(): + try: + # Lazy init ScaffoldingLlm inside the uvloop (like _lazy_init_scaffolding) + llm = ScaffoldingLlm( + trace_maker, + {NativeGenerationController.WorkerTag.GENERATION: worker}, + ) + try: + result = llm.generate_async(FAKE_PROMPT_STR) + await asyncio.wait_for(result, timeout=10.0) + result_holder["result"] = result + finally: + llm.shutdown() + except Exception as e: + error_holder["error"] = e + + def thread_fn(): + loop = uvloop.new_event_loop() + asyncio.set_event_loop(loop) + loop.run_until_complete(run_on_uvloop()) + loop.close() + + t = threading.Thread(target=thread_fn) + t.start() + t.join(timeout=30) + assert not t.is_alive(), "Thread deadlocked!" + + if "error" in error_holder: + raise error_holder["error"] + assert "result" in result_holder + assert result_holder["result"]._done is True + + def test_uvloop_concurrent_coroutines_shared_llm(self): + """Multiple concurrent coroutines on uvloop sharing one ScaffoldingLlm. + + This is the exact scenario from the training pipeline: + - AsyncTaskRunner creates a uvloop in a background thread + - Multiple _execute_workflow coroutines run concurrently + - They all share the same ChatScaffoldingWorkflow instance + - The ScaffoldingLlm is lazily initialized on first call + """ + import threading + + import uvloop + + worker, multi_turn_ctrl, trace_maker = self._build_shared_components( + max_turns=2 + ) + + num_concurrent = 10 + results = {} + errors = {} + + async def run_concurrent(): + # Lazy init ScaffoldingLlm inside uvloop (like _lazy_init_scaffolding) + llm = ScaffoldingLlm( + trace_maker, + {NativeGenerationController.WorkerTag.GENERATION: worker}, + ) + try: + + async def simulate_arun_episode(idx: int): + """Simulate what arun_episode does.""" + try: + # Each coroutine sets per-episode data and calls generate_async + # Note: in real code, clone() inside ScaffoldingLlm deep-copies + # the prototype controller, so the race on shared state is safe + # as long as the prototype is set before generate_async. + multi_turn_ctrl.messages = [ + {"role": "user", "content": f"Question {idx}"} + ] + multi_turn_ctrl.input_tokens = FAKE_INPUT_TOKENS + + result = llm.generate_async(f"prompt_{idx}") + await asyncio.wait_for(result, timeout=10.0) + results[idx] = result + except Exception as e: + errors[idx] = e + + # Run all concurrently (like AsyncTaskRunner does) + await asyncio.gather( + *[simulate_arun_episode(i) for i in range(num_concurrent)] + ) + finally: + llm.shutdown() + + def thread_fn(): + loop = uvloop.new_event_loop() + asyncio.set_event_loop(loop) + loop.run_until_complete(run_concurrent()) + loop.close() + + t = threading.Thread(target=thread_fn) + t.start() + t.join(timeout=60) + assert not t.is_alive(), "Thread deadlocked with concurrent coroutines!" + + if errors: + first_error = next(iter(errors.values())) + raise first_error + + assert len(results) == num_concurrent + for idx, result in results.items(): + assert result._done is True, f"Result {idx} not done" + + def test_uvloop_high_concurrency(self): + """Stress test: 50 concurrent coroutines on uvloop (closer to 256 in prod).""" + import threading + + import uvloop + + worker, multi_turn_ctrl, trace_maker = self._build_shared_components( + max_turns=1 + ) + multi_turn_ctrl.messages = [{"role": "user", "content": "test"}] + multi_turn_ctrl.input_tokens = FAKE_INPUT_TOKENS + + num_concurrent = 50 + done_count = {"value": 0} + errors = [] + + async def run_stress(): + llm = ScaffoldingLlm( + trace_maker, + {NativeGenerationController.WorkerTag.GENERATION: worker}, + ) + try: + + async def single_episode(idx: int): + try: + result = llm.generate_async(f"prompt_{idx}") + await asyncio.wait_for(result, timeout=15.0) + assert result._done is True + done_count["value"] += 1 + except Exception as e: + errors.append((idx, e)) + + await asyncio.gather( + *[single_episode(i) for i in range(num_concurrent)] + ) + finally: + llm.shutdown() + + def thread_fn(): + loop = uvloop.new_event_loop() + asyncio.set_event_loop(loop) + loop.run_until_complete(run_stress()) + loop.close() + + t = threading.Thread(target=thread_fn) + t.start() + t.join(timeout=60) + assert not t.is_alive(), "Thread deadlocked under high concurrency!" + + if errors: + raise errors[0][1] + assert done_count["value"] == num_concurrent + + +# =========================================================================== +# AsyncTaskRunner integration test (exact production scenario) +# =========================================================================== + + +class TestAsyncTaskRunnerScaffoldingLlm: + """Tests using the real AsyncTaskRunner to reproduce the production architecture. + + This is the closest reproduction of the actual training pipeline: + - AsyncTaskRunner runs a uvloop in a background thread + - A shared workflow object is used across all tasks + - ScaffoldingLlm is lazily initialized inside the first coroutine + - Multiple concurrent coroutines share one ScaffoldingLlm instance + """ + + def test_async_task_runner_with_scaffolding_llm(self): + """Full integration test using AsyncTaskRunner + shared ScaffoldingLlm.""" + from areal.infra.async_task_runner import AsyncTaskRunner + + worker = FakeChatWorker(response_text="42") + gen_ctrl = NativeGenerationController( + sampling_params={"max_tokens": 100, "temperature": 1.0} + ) + reward_ctrl = RLVRRewardController(_simple_reward_fn) + multi_turn_ctrl = MultiTurnChatController( + generation_controller=gen_ctrl, + max_turns=2, + reflection_message="Try again.", + ) + trace_maker = TraceTrajectoryMaker( + rollout_controller=multi_turn_ctrl, + reward_controller=reward_ctrl, + ) + + # Shared mutable state, like in the real workflow + shared_state = { + "scaffolding_llm": None, + "worker": worker, + "multi_turn_ctrl": multi_turn_ctrl, + "trace_maker": trace_maker, + } + + async def simulate_arun_episode(episode_idx: int) -> dict: + """Simulate ChatScaffoldingWorkflow.arun_episode.""" + # Lazy init (like _lazy_init_scaffolding) + if shared_state["scaffolding_llm"] is None: + shared_state["scaffolding_llm"] = ScaffoldingLlm( + shared_state["trace_maker"], + { + NativeGenerationController.WorkerTag.GENERATION: shared_state[ + "worker" + ] + }, + ) + + llm = shared_state["scaffolding_llm"] + + # Set per-episode data (race condition in production!) + shared_state["multi_turn_ctrl"].messages = [ + {"role": "user", "content": f"Question {episode_idx}"} + ] + shared_state["multi_turn_ctrl"].input_tokens = FAKE_INPUT_TOKENS + + # Generate + result = llm.generate_async(f"prompt_{episode_idx}") + await result + + return {"episode": episode_idx, "done": result._done} + + # Use AsyncTaskRunner like the real WorkflowExecutor + runner = AsyncTaskRunner(max_queue_size=64) + runner.initialize() + + num_tasks = 10 + try: + for i in range(num_tasks): + runner.submit(simulate_arun_episode, i, task_id=i) + + results = runner.wait(count=num_tasks, timeout=30.0) + + assert len(results) == num_tasks + for result in results: + assert result is not None + assert result["done"] is True + finally: + # Shutdown ScaffoldingLlm if initialized + if shared_state["scaffolding_llm"] is not None: + shared_state["scaffolding_llm"].shutdown() + runner.destroy() + + def test_async_task_runner_high_concurrency(self): + """Stress test: 50 concurrent tasks via AsyncTaskRunner.""" + from areal.infra.async_task_runner import AsyncTaskRunner + + worker = FakeChatWorker(response_text="42") + gen_ctrl = NativeGenerationController(sampling_params={"max_tokens": 100}) + reward_ctrl = RLVRRewardController(_simple_reward_fn) + multi_turn_ctrl = MultiTurnChatController( + generation_controller=gen_ctrl, + max_turns=1, + reflection_message="Try again.", + ) + trace_maker = TraceTrajectoryMaker( + rollout_controller=multi_turn_ctrl, + reward_controller=reward_ctrl, + ) + + shared_state = { + "scaffolding_llm": None, + "worker": worker, + "multi_turn_ctrl": multi_turn_ctrl, + "trace_maker": trace_maker, + } + + async def simulate_arun_episode(episode_idx: int) -> dict: + if shared_state["scaffolding_llm"] is None: + shared_state["scaffolding_llm"] = ScaffoldingLlm( + shared_state["trace_maker"], + { + NativeGenerationController.WorkerTag.GENERATION: shared_state[ + "worker" + ] + }, + ) + + llm = shared_state["scaffolding_llm"] + + shared_state["multi_turn_ctrl"].messages = [ + {"role": "user", "content": f"Q{episode_idx}"} + ] + shared_state["multi_turn_ctrl"].input_tokens = FAKE_INPUT_TOKENS + + result = llm.generate_async(f"prompt_{episode_idx}") + await result + + return {"episode": episode_idx, "done": result._done} + + runner = AsyncTaskRunner(max_queue_size=128) + runner.initialize() + + num_tasks = 50 + try: + for i in range(num_tasks): + runner.submit(simulate_arun_episode, i, task_id=i) + + results = runner.wait(count=num_tasks, timeout=60.0) + + assert len(results) == num_tasks + for result in results: + assert result is not None + assert result["done"] is True + finally: + if shared_state["scaffolding_llm"] is not None: + shared_state["scaffolding_llm"].shutdown() + runner.destroy() + + +# =========================================================================== +# Per-task ScaffoldingLlm instances (RolloutController path) +# =========================================================================== + + +class TestPerTaskScaffoldingLlm: + """Tests where EACH task creates its own ScaffoldingLlm instance. + + In the RolloutController path, RemoteInfEngine._resolve_workflow() creates + a NEW ChatScaffoldingWorkflow per submit() call. So 256 tasks create 256 + workflow instances, each with its own ScaffoldingLlm. All 256 ScaffoldingLlm + instances share the same uvloop (own_loop=False) and each schedules its own + _main_loop_async_func on it. + + This is the ACTUAL production architecture with scheduler.type=local. + """ + + def test_multiple_llm_instances_on_async_task_runner(self): + """Each task creates its own ScaffoldingLlm — deadlock reproduction.""" + from areal.infra.async_task_runner import AsyncTaskRunner + + worker = FakeChatWorker(response_text="42") + llm_instances = [] + + async def simulate_arun_episode_per_instance(episode_idx: int) -> dict: + """Each task creates its OWN ScaffoldingLlm (like _resolve_workflow).""" + # Create fresh controller hierarchy (like workflow.__init__ + build_scaffolding_llm) + gen_ctrl = NativeGenerationController( + sampling_params={"max_tokens": 100, "temperature": 1.0} + ) + reward_ctrl = RLVRRewardController(_simple_reward_fn) + multi_turn_ctrl = MultiTurnChatController( + generation_controller=gen_ctrl, + max_turns=1, + reflection_message="Try again.", + messages=[{"role": "user", "content": f"Q{episode_idx}"}], + input_tokens=FAKE_INPUT_TOKENS, + ) + trace_maker = TraceTrajectoryMaker( + rollout_controller=multi_turn_ctrl, + reward_controller=reward_ctrl, + ) + + # Each task creates its OWN ScaffoldingLlm + llm = ScaffoldingLlm( + trace_maker, + {NativeGenerationController.WorkerTag.GENERATION: worker}, + ) + llm_instances.append(llm) + + result = llm.generate_async(f"prompt_{episode_idx}") + await result + + return {"episode": episode_idx, "done": result._done} + + runner = AsyncTaskRunner(max_queue_size=64) + runner.initialize() + + num_tasks = 10 + try: + for i in range(num_tasks): + runner.submit( + simulate_arun_episode_per_instance, i, task_id=i + ) + + results = runner.wait(count=num_tasks, timeout=30.0) + + assert len(results) == num_tasks + for result in results: + assert result is not None + assert result["done"] is True + finally: + for llm in llm_instances: + try: + llm.shutdown() + except Exception: + pass + runner.destroy() + + def test_many_llm_instances_on_async_task_runner(self): + """50 per-task ScaffoldingLlm instances — stress test.""" + from areal.infra.async_task_runner import AsyncTaskRunner + + worker = FakeChatWorker(response_text="42") + llm_instances = [] + + async def simulate_arun_episode_per_instance(episode_idx: int) -> dict: + gen_ctrl = NativeGenerationController( + sampling_params={"max_tokens": 50} + ) + reward_ctrl = RLVRRewardController(_simple_reward_fn) + multi_turn_ctrl = MultiTurnChatController( + generation_controller=gen_ctrl, + max_turns=1, + reflection_message="retry", + messages=[{"role": "user", "content": f"Q{episode_idx}"}], + input_tokens=FAKE_INPUT_TOKENS, + ) + trace_maker = TraceTrajectoryMaker( + rollout_controller=multi_turn_ctrl, + reward_controller=reward_ctrl, + ) + + llm = ScaffoldingLlm( + trace_maker, + {NativeGenerationController.WorkerTag.GENERATION: worker}, + ) + llm_instances.append(llm) + + result = llm.generate_async(f"prompt_{episode_idx}") + await result + + return {"episode": episode_idx, "done": result._done} + + runner = AsyncTaskRunner(max_queue_size=128) + runner.initialize() + + num_tasks = 50 + try: + for i in range(num_tasks): + runner.submit( + simulate_arun_episode_per_instance, i, task_id=i + ) + + results = runner.wait(count=num_tasks, timeout=60.0) + + assert len(results) == num_tasks + for result in results: + assert result is not None + assert result["done"] is True + finally: + for llm in llm_instances: + try: + llm.shutdown() + except Exception: + pass + runner.destroy() + + +if __name__ == "__main__": + pytest.main([__file__, "-v"]) diff --git a/examples/scaffolding/chat_scaffolding.py b/examples/scaffolding/chat_scaffolding.py index 257cf58cb1..3c724c5333 100644 --- a/examples/scaffolding/chat_scaffolding.py +++ b/examples/scaffolding/chat_scaffolding.py @@ -21,7 +21,15 @@ from areal.api.cli_args import GenerationHyperparameters, GRPOConfig, load_expr_config from areal.api.engine_api import InferenceEngine from areal.dataset import get_custom_dataset -from areal.experimental.scaffolding._compat import GenerationTask +from areal.experimental.scaffolding._compat import ( + NativeGenerationController, + ScaffoldingLlm, +) +from areal.experimental.scaffolding.controllers import ( + MultiTurnChatController, + RLVRRewardController, + TraceTrajectoryMaker, +) from areal.experimental.scaffolding.workflow import ScaffoldingWorkflow from areal.trainer import PPOTrainer from areal.utils import logging @@ -44,8 +52,8 @@ class ChatScaffoldingWorkflow(ScaffoldingWorkflow): non-final turn the ``reflection_message`` is appended as a user message to prompt the model to retry. Reward is computed on the final turn only. - Generation and reward use the same direct-API approach as the base class - (``_generate_via_worker`` + ``_compute_rewards_via_controller``). + Generation and reward are delegated to ``scaffolding_llm`` which wraps a + ``TraceTrajectoryMaker`` with a ``MultiTurnChatController``. Parameters ---------- @@ -81,14 +89,60 @@ def __init__( self.max_turns = max_turns self.reflection_message = reflection_message + def build_scaffolding_llm(self, engine: InferenceEngine) -> ScaffoldingLlm: + """Build ScaffoldingLlm with MultiTurnChatController + TraceTrajectoryMaker. + + Parameters + ---------- + engine : InferenceEngine + The inference engine. + + Returns + ------- + ScaffoldingLlm + The constructed ScaffoldingLlm instance. + """ + stop_strings = [] + if self.gconfig.stop_token_ids: + for tid in self.gconfig.stop_token_ids: + decoded = self.tokenizer.decode([tid]) + if decoded: + stop_strings.append(decoded) + + sampling_params = { + "max_tokens": self.gconfig.max_new_tokens, + "temperature": self.gconfig.temperature or 1.0, + } + if stop_strings: + sampling_params["stop"] = stop_strings + + self.gen_controller = NativeGenerationController( + sampling_params=sampling_params + ) + self.reward_controller = RLVRRewardController(self.reward_fn) + self.multi_turn_controller = MultiTurnChatController( + generation_controller=self.gen_controller, + max_turns=self.max_turns, + reflection_message=self.reflection_message, + tokenizer=self.tokenizer, + ) + self.trajectory_maker = TraceTrajectoryMaker( + rollout_controller=self.multi_turn_controller, + reward_controller=self.reward_controller, + ) + return ScaffoldingLlm( + self.trajectory_maker, + {NativeGenerationController.WorkerTag.GENERATION: self.worker}, + ) + async def arun_episode( self, engine: InferenceEngine, data: dict[str, Any] ) -> dict[str, torch.Tensor]: """Run a single multi-turn chat episode. - Each turn: generate via SGLang OpenAI API, then optionally append a - reflection message and retry. After all turns, compute reward on the - final completion and build tensor dicts identical to the base class. + Delegates the full episode (multi-turn generation + reward) to + ``self.scaffolding_llm``, which wraps a ``TraceTrajectoryMaker`` + with a ``MultiTurnChatController``. Parameters ---------- @@ -103,39 +157,12 @@ async def arun_episode( Trajectory tensors for PPO training. """ if self.worker is None: + logger.info("Calling _lazy_init_scaffolding ...") self._lazy_init_scaffolding(engine) + logger.info("_lazy_init_scaffolding done.") - # Start from the original messages - messages = list(data["messages"]) - last_output_str = "" - all_output_tokens: list[int] = [] - - for turn in range(self.max_turns): - # Build prompt for this turn - input_ids = list( - self.tokenizer.apply_chat_template( - messages, - tokenize=True, - add_generation_prompt=True, - enable_thinking=self.enable_thinking, - ) - ) - prompt_str = self.tokenizer.decode(input_ids) - - # Generate via the SGLang OpenAI API (same as base class) - gen_task = await self._generate_via_worker(prompt_str, input_ids) - last_output_str = gen_task.output_str or "" - all_output_tokens = list(gen_task.output_tokens or []) - - # Append the assistant response to messages - messages.append({"role": "assistant", "content": last_output_str}) - - # Append reflection message for non-final turns - if turn < self.max_turns - 1: - messages.append({"role": "user", "content": self.reflection_message}) - - # Compute reward on the final turn's output (same pattern as base class) - final_input_ids = list( + # Tokenize the original prompt (before multi-turn) + input_ids = list( self.tokenizer.apply_chat_template( data["messages"], tokenize=True, @@ -143,22 +170,46 @@ async def arun_episode( enable_thinking=self.enable_thinking, ) ) - final_prompt_str = self.tokenizer.decode(final_input_ids) + prompt_str = self.tokenizer.decode(input_ids) + + # Configure per-episode data on multi-turn controller + # (clone() in scaffolding_llm will deep-copy these) + self.multi_turn_controller.messages = data["messages"] + self.multi_turn_controller.input_tokens = input_ids + + # Run full pipeline via scaffolding_llm + logger.info("Calling generate_async ...") + result = self.scaffolding_llm.generate_async(prompt_str) + logger.info("generate_async returned, awaiting result ...") + await result + logger.info("Result received.") + + # Extract trace results from ScaffoldingOutput + scaffolding_output = result.outputs[0] + trace_results = scaffolding_output.data + + # Get the final output text from the last traced interaction + if trace_results: + last_interaction = list(trace_results.values())[-1] + output_str = "" + if last_interaction.completion is not None: + output_str = ( + last_interaction.completion.choices[0].message.content or "" + ) + else: + output_str = scaffolding_output.text or "" - final_gen_task = GenerationTask( - input_str=final_prompt_str, - input_tokens=final_input_ids, - output_str=last_output_str, - output_tokens=all_output_tokens, - ) - reward = await self._compute_rewards_via_controller( - final_gen_task, final_prompt_str, data + output_tokens = self.tokenizer.encode(output_str, add_special_tokens=False) + + # Compute reward on the final turn's output + reward = float( + self.reward_fn(prompt_str, output_str, input_ids, output_tokens, **data) ) - # Build tensor dict (same as base class) - seq = final_input_ids + all_output_tokens + # Build tensor dict for PPO training + seq = input_ids + output_tokens logprobs = [0.0] * len(seq) - loss_mask = [0] * len(final_input_ids) + [1] * len(all_output_tokens) + loss_mask = [0] * len(input_ids) + [1] * len(output_tokens) versions = [-1] * len(seq) res = { From 098c8a1107342d278ea2e3e1b5eb510e01a75044 Mon Sep 17 00:00:00 2001 From: Fred Wei <20514172+WeiHaocheng@users.noreply.github.com> Date: Wed, 4 Mar 2026 06:17:21 +0000 Subject: [PATCH 10/18] Fix scaffolding holding issue --- areal/experimental/scaffolding/core/result.py | 28 ++++-- .../scaffolding/core/scaffolding_llm.py | 87 +++++-------------- examples/scaffolding/chat_scaffolding.py | 5 -- 3 files changed, 43 insertions(+), 77 deletions(-) diff --git a/areal/experimental/scaffolding/core/result.py b/areal/experimental/scaffolding/core/result.py index f452634ae7..52f786bef6 100644 --- a/areal/experimental/scaffolding/core/result.py +++ b/areal/experimental/scaffolding/core/result.py @@ -2,6 +2,7 @@ # Vendored from tensorrt_llm.scaffolding.result import asyncio +import queue from collections.abc import Mapping from dataclasses import dataclass from typing import Any @@ -15,9 +16,16 @@ class ScaffoldingOutput: class ScaffoldingResult: + """Result object for scaffolding requests. + + Uses a thread-safe ``queue.Queue`` for cross-thread communication so + that producers (ScaffoldingLlm loop thread) and consumers + (caller's event loop) can safely exchange data. + """ + def __init__(self): super().__init__() - self.aqueue = asyncio.Queue() + self._queue: queue.Queue = queue.Queue() self.outputs = [] # only support one output for now, so use an empty obj to init self.outputs.append(ScaffoldingOutput("", [])) @@ -31,22 +39,30 @@ def set_output(self, output: ScaffoldingOutput | Any): self.set_output_streaming(None) def set_output_streaming(self, output: ScaffoldingOutput | Any): - self.aqueue.put_nowait(output) + self._queue.put_nowait(output) def set_task_collections(self, task_collections: Mapping[str, Any]): self.task_collections = task_collections async def _aresult_step(self): - obj = await self.aqueue.get() + """Asynchronously wait for the next item from the thread-safe queue.""" + loop = asyncio.get_running_loop() + obj = await loop.run_in_executor(None, self._queue.get) if obj is None: self._done = True else: # obj is ScaffoldingOutput self.outputs[0] = obj def result(self, timeout: float | None = None) -> "ScaffoldingResult": - if not self._done: - loop = asyncio.get_event_loop() - asyncio.run_coroutine_threadsafe(self.aresult(), loop).result() + while not self._done: + try: + obj = self._queue.get(timeout=timeout) + except queue.Empty: + break + if obj is None: + self._done = True + else: + self.outputs[0] = obj return self async def aresult(self) -> "ScaffoldingResult": diff --git a/areal/experimental/scaffolding/core/scaffolding_llm.py b/areal/experimental/scaffolding/core/scaffolding_llm.py index 6af58e5ccf..ca5e621581 100644 --- a/areal/experimental/scaffolding/core/scaffolding_llm.py +++ b/areal/experimental/scaffolding/core/scaffolding_llm.py @@ -9,15 +9,11 @@ from dataclasses import dataclass from typing import Any -from areal.utils import logging as areal_logging - from .controller import Controller, ParallelProcess from .result import ScaffoldingResult from .task import Task from .worker import Worker -_logger = areal_logging.getLogger("ScaffoldingLlm") - @dataclass(frozen=True) class ScaffoldingRequest: @@ -37,15 +33,17 @@ def __init__( self.prototype_controller = prototype_controller self.workers = workers - self.loop = self._get_loop() - asyncio.set_event_loop(self.loop) - self.task_queue = asyncio.Queue() - self.main_loop_stop_event = asyncio.Event() - self.shutdown_event = asyncio.Event() - if self.own_loop: - self._run_main_loop_thread() - else: - self._run_main_loop_coroutine() + # Always create a dedicated event loop in a separate thread. + # This avoids deadlocks when ScaffoldingLlm is used inside another + # event loop (e.g. AsyncTaskRunner's uvloop), where fire-and-forget + # tasks created via create_task() would never get executed. + # + # asyncio primitives (Queue, Event) are created inside the loop + # thread to ensure they are bound to the correct event loop. + self.loop = asyncio.new_event_loop() + self._ready = threading.Event() + self._run_main_loop_thread() + self._ready.wait() # For top scheduler self.running_req_count = 0 @@ -60,56 +58,15 @@ def __enter__(self): def __exit__(self): self.shutdown() - def _get_loop(self): - try: - self.own_loop = False - return asyncio.get_running_loop() - except RuntimeError: - self.own_loop = True - return asyncio.new_event_loop() - return None - - def _schedule_on_loop(self, coro): - """Schedule a coroutine on self.loop. - - Uses create_task when called from within the event loop thread - (to avoid deadlocks with run_coroutine_threadsafe on uvloop), - and falls back to run_coroutine_threadsafe for cross-thread calls. - """ - try: - running = asyncio.get_running_loop() - except RuntimeError: - running = None - - if running is self.loop: - _logger.info("_schedule_on_loop: same loop, create_task for %s", coro.__name__) - self.loop.create_task(coro) - elif running is not None: - _logger.warning( - "_schedule_on_loop: different loop! running=%s self.loop=%s", - id(running), id(self.loop), - ) - asyncio.run_coroutine_threadsafe(coro, self.loop) - else: - _logger.info("_schedule_on_loop: no running loop, run_coroutine_threadsafe for %s", coro.__name__) - asyncio.run_coroutine_threadsafe(coro, self.loop) - async def _handle_controller_generator( self, gen: Generator, request: ScaffoldingRequest = None ): """Handle a controller generator, processing tasks and parallel processes.""" - step = 0 for obj in gen: if isinstance(obj, ParallelProcess): await self._handle_parallel_process(obj, request) else: - step += 1 - _logger.info( - "Dispatching task list step=%d, tasks=%d, types=%s", - step, len(obj), [type(t).__name__ for t in obj], - ) await self._handle_task_list(obj, request) - _logger.info("Task list step=%d completed.", step) async def _handle_task_list( self, tasks: list[Task], request: ScaffoldingRequest = None @@ -193,19 +150,21 @@ async def _handle_event_loop(self): async def _main_loop_async_func(self): """Main async loop function.""" - _logger.info("_main_loop_async_func STARTED (loop=%s)", id(asyncio.get_running_loop())) handle_event_task = asyncio.create_task(self._handle_event_loop()) await handle_event_task self.main_loop_stop_event.set() - def _run_main_loop_coroutine(self): - self._schedule_on_loop(self._main_loop_async_func()) - def _run_main_loop_thread(self): def main_loop_thread(): + asyncio.set_event_loop(self.loop) + # Create asyncio primitives inside the loop thread. + self.task_queue = asyncio.Queue() + self.main_loop_stop_event = asyncio.Event() + self.shutdown_event = asyncio.Event() + self._ready.set() self.loop.run_until_complete(self._main_loop_async_func()) - self.main_loop_thread = threading.Thread(target=main_loop_thread) + self.main_loop_thread = threading.Thread(target=main_loop_thread, daemon=True) self.main_loop_thread.start() def generate_async(self, prompt: str) -> ScaffoldingResult: @@ -227,7 +186,7 @@ async def put_request(): else: await self.task_queue.put(request) - self._schedule_on_loop(put_request()) + asyncio.run_coroutine_threadsafe(put_request(), self.loop) return result @@ -260,12 +219,8 @@ async def stop_task_on_loop(): for worker in self.workers.values(): await worker.async_shutdown() - self._schedule_on_loop(stop_task_on_loop()) - - if self.own_loop: - self.main_loop_thread.join() - else: - self.shutdown_event.set() + asyncio.run_coroutine_threadsafe(stop_task_on_loop(), self.loop) + self.main_loop_thread.join() if shutdown_workers: shutdown_workers_func() diff --git a/examples/scaffolding/chat_scaffolding.py b/examples/scaffolding/chat_scaffolding.py index 3c724c5333..cacdc5903e 100644 --- a/examples/scaffolding/chat_scaffolding.py +++ b/examples/scaffolding/chat_scaffolding.py @@ -157,9 +157,7 @@ async def arun_episode( Trajectory tensors for PPO training. """ if self.worker is None: - logger.info("Calling _lazy_init_scaffolding ...") self._lazy_init_scaffolding(engine) - logger.info("_lazy_init_scaffolding done.") # Tokenize the original prompt (before multi-turn) input_ids = list( @@ -178,11 +176,8 @@ async def arun_episode( self.multi_turn_controller.input_tokens = input_ids # Run full pipeline via scaffolding_llm - logger.info("Calling generate_async ...") result = self.scaffolding_llm.generate_async(prompt_str) - logger.info("generate_async returned, awaiting result ...") await result - logger.info("Result received.") # Extract trace results from ScaffoldingOutput scaffolding_output = result.outputs[0] From eb01b10e9f6a694b172d28ef2b25369d146446e2 Mon Sep 17 00:00:00 2001 From: narutolhy Date: Wed, 4 Mar 2026 20:31:29 -0800 Subject: [PATCH 11/18] fix dead lock --- areal/experimental/scaffolding/core/scaffolding_llm.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/areal/experimental/scaffolding/core/scaffolding_llm.py b/areal/experimental/scaffolding/core/scaffolding_llm.py index ca5e621581..7e7ac751e2 100644 --- a/areal/experimental/scaffolding/core/scaffolding_llm.py +++ b/areal/experimental/scaffolding/core/scaffolding_llm.py @@ -179,7 +179,7 @@ async def put_request(): controller=self.prototype_controller.clone(), ) except Exception as e: - self.task_queue.put(None) + await self.task_queue.put(None) print( f"Error: build ScaffoldingRequest failed: {e} \n {traceback.format_exc()}" ) From 973c79115947d6226085ab46547521306b278b4d Mon Sep 17 00:00:00 2001 From: Fred Wei <20514172+WeiHaocheng@users.noreply.github.com> Date: Thu, 5 Mar 2026 09:19:39 +0000 Subject: [PATCH 12/18] fix gsm8k_rlvr_scaffolding --- areal/experimental/scaffolding/workflow.py | 13 +++++++++++++ examples/scaffolding/gsm8k_rlvr_scaffolding.py | 4 ++-- examples/scaffolding/gsm8k_rlvr_scaffolding.yaml | 14 +++++++------- 3 files changed, 22 insertions(+), 9 deletions(-) diff --git a/areal/experimental/scaffolding/workflow.py b/areal/experimental/scaffolding/workflow.py index 19b46ea2b5..343e57ca89 100644 --- a/areal/experimental/scaffolding/workflow.py +++ b/areal/experimental/scaffolding/workflow.py @@ -202,4 +202,17 @@ async def arun_episode( scaffolding_output = result.outputs[0] interactions = scaffolding_output.data interaction = next(iter(interactions.values())) + + # If output_tokens is missing (e.g., SGLang didn't return token_ids), + # tokenize the output_str as a fallback so loss_mask is non-zero. + resp = interaction.model_response + if resp is not None and not resp.output_tokens: + output_text = scaffolding_output.text or "" + output_tokens = self.tokenizer.encode( + output_text, add_special_tokens=False + ) + resp.output_tokens = output_tokens + resp.output_logprobs = [0.0] * len(output_tokens) + resp.output_versions = [-1] * len(output_tokens) + return interaction.to_tensor_dict() diff --git a/examples/scaffolding/gsm8k_rlvr_scaffolding.py b/examples/scaffolding/gsm8k_rlvr_scaffolding.py index 41dda7da56..9cdbf8817d 100644 --- a/examples/scaffolding/gsm8k_rlvr_scaffolding.py +++ b/examples/scaffolding/gsm8k_rlvr_scaffolding.py @@ -79,9 +79,9 @@ def main(args): valid_dataset=valid_dataset, ) as trainer: trainer.train( - workflow=GSM8KScaffoldingWorkflow, + workflow="examples.scaffolding.gsm8k_rlvr_scaffolding.GSM8KScaffoldingWorkflow", workflow_kwargs=workflow_kwargs, - eval_workflow=GSM8KScaffoldingWorkflow, + eval_workflow="examples.scaffolding.gsm8k_rlvr_scaffolding.GSM8KScaffoldingWorkflow", eval_workflow_kwargs=eval_workflow_kwargs, ) diff --git a/examples/scaffolding/gsm8k_rlvr_scaffolding.yaml b/examples/scaffolding/gsm8k_rlvr_scaffolding.yaml index b272679709..18ce7ac759 100644 --- a/examples/scaffolding/gsm8k_rlvr_scaffolding.yaml +++ b/examples/scaffolding/gsm8k_rlvr_scaffolding.yaml @@ -1,5 +1,5 @@ # RLVR Scaffolding Example Configuration for GSM8K -# Compatible with GRPOConfig, for 8-GPU setup (4 inference + 4 training) +# Compatible with GRPOConfig, single-GPU setup (inference + training share 1 GPU) experiment_name: gsm8k-rlvr-scaffolding trial_name: trial0 @@ -11,18 +11,18 @@ tokenizer_path: ${actor.path} cluster: n_nodes: 1 - n_gpus_per_node: 8 + n_gpus_per_node: 1 fileroot: /tmp/areal/experiments name_resolve: type: nfs nfs_record_root: /tmp/areal/name_resolve -allocation_mode: sglang:d4p1t1+d4p1t1 +allocation_mode: sglang:d1+d1 rollout: experiment_name: ${experiment_name} trial_name: ${trial_name} - max_concurrent_rollouts: 256 + max_concurrent_rollouts: 64 queue_size: null consumer_batch_size: ${train_dataset.batch_size} max_head_offpolicyness: 2 @@ -108,7 +108,7 @@ sglang: dtype: ${actor.dtype} max_running_requests: null context_length: 2048 - mem_fraction_static: 0.5 + mem_fraction_static: 0.3 attention_backend: flashinfer vllm: @@ -121,7 +121,7 @@ vllm: # Datasets train_dataset: - batch_size: 256 + batch_size: 64 shuffle: true pin_memory: true num_workers: 4 @@ -130,7 +130,7 @@ train_dataset: max_length: 2048 valid_dataset: - batch_size: 256 + batch_size: 64 pin_memory: true num_workers: 4 path: openai/gsm8k From ebb674dc33956ea5bcddd0da31819899c4716b3d Mon Sep 17 00:00:00 2001 From: "luhongyu.4869" Date: Mon, 9 Mar 2026 21:53:05 -0700 Subject: [PATCH 13/18] feat: add IPv6 network utilities - Add is_ipv6_address() to detect IPv6 address strings - Add format_addr() to bracket IPv6 in host:port strings (RFC 2732) - Extend gethostip() with IPv6 fallback for IPv6-only environments - Update is_port_free() to check both AF_INET and AF_INET6 Co-Authored-By: Claude Sonnet 4.6 --- areal/utils/network.py | 78 +++++++++++++++++++++++++++++------------- 1 file changed, 55 insertions(+), 23 deletions(-) diff --git a/areal/utils/network.py b/areal/utils/network.py index 199998a7ab..8985d37fb3 100644 --- a/areal/utils/network.py +++ b/areal/utils/network.py @@ -6,11 +6,26 @@ def gethostname(): return socket.gethostname() +def is_ipv6_address(ip: str) -> bool: + """Return True if *ip* is an IPv6 address string.""" + try: + socket.inet_pton(socket.AF_INET6, ip) + return True + except OSError: + return False + + +def format_addr(host: str, port: int) -> str: + """Format host:port, wrapping IPv6 addresses in brackets as required by URLs.""" + if is_ipv6_address(host): + return f"[{host}]:{port}" + return f"{host}:{port}" + + def gethostip(probe_host: str = "8.8.8.8", probe_port: int = 80) -> str: """ - Find the local IPv4 address for outbound route to `probe_host:probe_port` (typically - a LAN/private IP). Use hostname resolution first; if it fails or returns loopback (127.*), - fall back to a UDP connect. + Find the local IP address for outbound traffic. Tries IPv4 first, then falls back + to IPv6 for IPv6-only environments. Args: probe_host: Remote IPv4 address used to trigger route selection, default to Google @@ -18,10 +33,10 @@ def gethostip(probe_host: str = "8.8.8.8", probe_port: int = 80) -> str: probe_port: Remote port used for the UDP probe. Returns: - The selected local IPv4 address as a string + The selected local IP address as a string (IPv4 or IPv6) Raises: - RuntimeError: If no suitable IPv4 address can be determined + RuntimeError: If no suitable IP address can be determined """ try: ip = socket.gethostbyname(socket.gethostname()) @@ -34,6 +49,23 @@ def gethostip(probe_host: str = "8.8.8.8", probe_port: int = 80) -> str: with socket.socket(socket.AF_INET, socket.SOCK_DGRAM) as sock: sock.connect((probe_host, probe_port)) return sock.getsockname()[0] + except OSError: + pass + + # IPv6 fallback for IPv6-only environments + try: + infos = socket.getaddrinfo(socket.gethostname(), None, socket.AF_INET6) + for info in infos: + ip = info[4][0] + if ip and not ip.startswith("::1"): + return ip + except socket.gaierror: + pass + + try: + with socket.socket(socket.AF_INET6, socket.SOCK_DGRAM) as sock: + sock.connect(("2001:4860:4860::8888", 80)) + return sock.getsockname()[0] except OSError as e: raise RuntimeError("Could not determine host IP") from e @@ -101,27 +133,27 @@ def find_free_ports( def is_port_free(port: int) -> bool: """ - Check if a port is free by attempting to bind to it. + Check if a port is free by attempting to bind to it on both IPv4 and IPv6. Args: port: Port number to check Returns: - True if port is free, False otherwise + True if port is free on all address families, False otherwise """ - # Check TCP - sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM) - try: - sock.bind(("", port)) - sock.close() - except OSError: - return False - - # Check UDP - sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM) - try: - sock.bind(("", port)) - sock.close() - return True - except OSError: - return False + for family in (socket.AF_INET, socket.AF_INET6): + for sock_type in (socket.SOCK_STREAM, socket.SOCK_DGRAM): + try: + sock = socket.socket(family, sock_type) + sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) + if family == socket.AF_INET6: + sock.setsockopt(socket.IPPROTO_IPV6, socket.IPV6_V6ONLY, 1) + sock.bind(("", port)) + sock.close() + except OSError: + try: + sock.close() + except Exception: + pass + return False + return True From caa5a49f3fe126c6abfc76a943de38e89f588679 Mon Sep 17 00:00:00 2001 From: "luhongyu.4869" Date: Mon, 9 Mar 2026 21:53:10 -0700 Subject: [PATCH 14/18] feat: support IPv6 hosts in infra server address handling Use format_addr() everywhere host:port strings are constructed so that IPv6 addresses are properly bracketed in URLs. Also: - remote_inf_engine: accept per-worker port_range to avoid concurrent find_free_ports() TOCTOU races when multiple workers start in parallel - rollout_controller: partition port space by worker rank so each worker picks from a non-overlapping range, eliminating bind conflicts Co-Authored-By: Claude Sonnet 4.6 --- areal/infra/controller/rollout_controller.py | 33 +++++++++++++++----- areal/infra/launcher/sglang_server.py | 8 ++--- areal/infra/remote_inf_engine.py | 26 +++++++++------ areal/infra/rpc/rpc_server.py | 4 +-- 4 files changed, 47 insertions(+), 24 deletions(-) diff --git a/areal/infra/controller/rollout_controller.py b/areal/infra/controller/rollout_controller.py index 1de3b44e5a..d8ebc9083b 100644 --- a/areal/infra/controller/rollout_controller.py +++ b/areal/infra/controller/rollout_controller.py @@ -34,7 +34,7 @@ from areal.utils import logging, perf_tracer from areal.utils.data import concat_padded_tensors, cycle_dataloader from areal.utils.dynamic_import import import_from_string -from areal.utils.network import find_free_ports, gethostip +from areal.utils.network import find_free_ports, format_addr, gethostip from areal.utils.perf_tracer import trace_perf from ..staleness_manager import StalenessManager @@ -236,7 +236,7 @@ async def _async_initialize( engine_name=self._engine_name(rank), # args in `engine_api` engine_id=str(rank), - addr=f"{info.host}:{info.port}", + addr=format_addr(info.host, info.port), *args, **kwargs, ) @@ -246,9 +246,26 @@ async def _async_initialize( ] await asyncio.gather(*tasks) else: - self.server_infos = await self._collective_rpc_async( - "launch_server", server_args=server_args - ) + # Assign each worker a unique port range to avoid concurrent port conflicts. + # Workers on the same node calling find_free_ports simultaneously can pick + # the same port (TOCTOU race). Partitioning by global rank prevents this. + n_workers = len(self.workers) + ports_per_worker = 40000 // max(n_workers, 1) + server_info_tasks = [] + for rank, worker in enumerate(self.workers): + worker_server_args = dict(server_args or {}) + port_range_start = 10000 + rank * ports_per_worker + port_range_end = port_range_start + ports_per_worker - 1 + worker_server_args["port_range"] = (port_range_start, port_range_end) + server_info_tasks.append( + self.scheduler.async_call_engine( + worker_id=worker.id, + method="launch_server", + engine_name=self._engine_name(rank), + server_args=worker_server_args, + ) + ) + self.server_infos = list(await asyncio.gather(*server_info_tasks)) tasks = [ self.scheduler.async_call_engine( worker_id=worker.id, @@ -352,10 +369,10 @@ async def _async_start_proxy(self) -> None: worker_id=worker.id, method="initialize", engine_name=f"proxy/{rank}", - addr=f"{server_info.host}:{server_info.port}", + addr=format_addr(server_info.host, server_info.port), ) ) - self.proxy_addrs.append(f"http://{worker.ip}:{worker.worker_ports[0]}") + self.proxy_addrs.append(f"http://{format_addr(worker.ip, worker.worker_ports[0])}") await asyncio.gather(*init_tasks) logger.info(f"Proxy servers initialized. Addresses: {self.proxy_addrs}") @@ -492,7 +509,7 @@ def callback_addr(self) -> str: """Return callback server address as 'host:port'.""" if self._callback_host is None or self._callback_port is None: raise RuntimeError("Callback server not started") - return f"{self._callback_host}:{self._callback_port}" + return format_addr(self._callback_host, self._callback_port) def _resolve_task_future(self, task_id: int): """Resolve a pending future with the task result.""" diff --git a/areal/infra/launcher/sglang_server.py b/areal/infra/launcher/sglang_server.py index 7bd370eb1f..ba8c8a3e70 100644 --- a/areal/infra/launcher/sglang_server.py +++ b/areal/infra/launcher/sglang_server.py @@ -24,7 +24,7 @@ ) from areal.infra.utils.proc import kill_process_tree from areal.utils import logging, name_resolve, names -from areal.utils.network import find_free_ports, gethostip +from areal.utils.network import find_free_ports, format_addr, gethostip logger = logging.getLogger("SGLangWrapper") @@ -185,7 +185,7 @@ def run(self): node_rank=node_rank, ) launch_server_args.append((cmd, host_ip, server_port, node_rank)) - server_addresses.append(f"http://{host_ip}:{server_port}") + server_addresses.append(f"http://{format_addr(host_ip, server_port)}") with ThreadPoolExecutor(max_workers=n_servers_per_proc) as executor: server_iterator = executor.map( @@ -199,10 +199,10 @@ def run(self): def launch_one_server(self, cmd, host_ip, server_port, node_rank): server_process = launch_server_cmd(cmd) - wait_for_server(f"http://{host_ip}:{server_port}") + wait_for_server(f"http://{format_addr(host_ip, server_port)}") if node_rank == 0: name = names.gen_servers(self.experiment_name, self.trial_name) - name_resolve.add_subentry(name, f"{host_ip}:{server_port}") + name_resolve.add_subentry(name, format_addr(host_ip, server_port)) logger.info(f"SGLang server launched at: http://{host_ip}:{server_port}") return server_process diff --git a/areal/infra/remote_inf_engine.py b/areal/infra/remote_inf_engine.py index 3d071b1d55..cee5dc8ffe 100644 --- a/areal/infra/remote_inf_engine.py +++ b/areal/infra/remote_inf_engine.py @@ -43,7 +43,7 @@ from areal.utils import logging, name_resolve, names from areal.utils.data import concat_padded_tensors from areal.utils.dynamic_import import import_from_string -from areal.utils.network import find_free_ports, gethostip +from areal.utils.network import find_free_ports, format_addr, gethostip, is_ipv6_address from areal.utils.perf_tracer import trace_perf from .workflow_executor import WorkflowExecutor @@ -405,7 +405,7 @@ def initialize( self.addresses = addr if isinstance(addr, list) else [addr] self.logger.info("Get server addresses from the `addr` argument.") elif len(self.local_server_processes) > 0: - self.addresses = [f"{s.host}:{s.port}" for s in self.local_server_processes] + self.addresses = [format_addr(s.host, s.port) for s in self.local_server_processes] self.logger.info("Get server addresses from the local subprocess.") elif ( self.config.experiment_name is not None @@ -1166,13 +1166,19 @@ async def _fn(): def launch_server(self, server_args: dict[str, Any]) -> LocalInfServerInfo: """Launch a local inference server.""" - server_args["host"] = gethostip() - server_args["port"] = find_free_ports(1)[0] + host_ip = gethostip() + port_range = server_args.pop("port_range", (1024, 65535)) + port = find_free_ports(1, port_range=port_range)[0] + # Use wildcard bind address so uvicorn/SGLang can bind on any IP family, + # but keep the actual host IP for external address registration. + bind_host = "::" if is_ipv6_address(host_ip) else "0.0.0.0" + server_args["host"] = bind_host + server_args["port"] = port process = self.backend.launch_server(server_args) - address = f"{server_args['host']}:{server_args['port']}" + address = format_addr(host_ip, port) server_info = LocalInfServerInfo( - host=server_args["host"], - port=server_args["port"], + host=host_ip, + port=port, process=process, ) try: @@ -1181,8 +1187,8 @@ def launch_server(self, server_args: dict[str, Any]) -> LocalInfServerInfo: if ray.is_initialized(): # do not return with process for ray as it is not picklable return LocalInfServerInfo( - host=server_args["host"], - port=server_args["port"], + host=host_ip, + port=port, process=None, ) return server_info @@ -1194,7 +1200,7 @@ def launch_server(self, server_args: dict[str, Any]) -> LocalInfServerInfo: raise def _shutdown_one_server(self, server_info: LocalInfServerInfo): - addr = f"{server_info.host}:{server_info.port}" + addr = format_addr(server_info.host, server_info.port) if addr in self.addresses: self.addresses.remove(addr) if server_info.process.poll() is not None: diff --git a/areal/infra/rpc/rpc_server.py b/areal/infra/rpc/rpc_server.py index 999f87d59d..ee8a20f0ff 100644 --- a/areal/infra/rpc/rpc_server.py +++ b/areal/infra/rpc/rpc_server.py @@ -34,7 +34,7 @@ tensor_container_to, ) from areal.utils.dynamic_import import import_from_string -from areal.utils.network import find_free_ports, gethostip +from areal.utils.network import find_free_ports, format_addr, gethostip logger = logging.getLogger("SyncRPCServer") @@ -182,7 +182,7 @@ def _wait_for_worker_ready(host: str, port: int, timeout: float = 60) -> bool: Returns: True if the worker is ready, False if timeout is reached. """ - url = f"http://{host}:{port}/health" + url = f"http://{format_addr(host, port)}/health" deadline = time.time() + timeout while time.time() < deadline: From 0a6880f9c9003b9daf31ba1d7a3262b4b3a343f8 Mon Sep 17 00:00:00 2001 From: "luhongyu.4869" Date: Mon, 9 Mar 2026 21:53:17 -0700 Subject: [PATCH 15/18] fix: IPv6 NCCL tcp URL and cpu_group for cross-node Gloo collectives - Bracket IPv6 addresses in tcp://[addr]:port init_method so PyTorch dist.init_process_group can parse the URL correctly - Pass group=self.cpu_group (Gloo) to allocate_balanced_mbs_synced inside split_padded_tensor_dict_into_mb_list; without this it falls back to NCCL WORLD which hangs on nodes where IB has no routing Co-Authored-By: Claude Sonnet 4.6 --- areal/engine/fsdp_engine.py | 20 ++++++++++++++------ 1 file changed, 14 insertions(+), 6 deletions(-) diff --git a/areal/engine/fsdp_engine.py b/areal/engine/fsdp_engine.py index 618543a3fd..d8574c6d86 100644 --- a/areal/engine/fsdp_engine.py +++ b/areal/engine/fsdp_engine.py @@ -117,7 +117,7 @@ ) from areal.utils.functional import gather_logprobs, gather_logprobs_entropy from areal.utils.hf_utils import load_hf_processor_and_tokenizer, load_hf_tokenizer -from areal.utils.network import find_free_ports, gethostip +from areal.utils.network import find_free_ports, format_addr, gethostip, is_ipv6_address from areal.utils.offload import is_tms_enabled, torch_memory_saver from areal.utils.perf_tracer import trace_perf, trace_scope from areal.utils.save_load import get_state_dict_from_repo_id_or_path @@ -1048,7 +1048,10 @@ def _init_weight_update_from_distributed(self, meta: WeightUpdateMeta): assert meta.type == "xccl" # Reset weight weight meta with local info - meta.nccl_master_address = self.weight_update_master_addr = gethostip() + raw_addr = gethostip() + self.weight_update_master_addr = raw_addr + # Pre-bracket IPv6 so SGLang's f"tcp://{master_address}:{port}" constructs a valid URL + meta.nccl_master_address = f"[{raw_addr}]" if is_ipv6_address(raw_addr) else raw_addr meta.nccl_master_port = self.weight_update_master_port = find_free_ports(1)[0] meta.nccl_group_name = self.weight_update_group_name @@ -1060,15 +1063,16 @@ def _init_weight_update_from_distributed(self, meta: WeightUpdateMeta): fut = self.rollout_engine.init_weights_update_group(meta) + tcp_addr = format_addr(meta.nccl_master_address, meta.nccl_master_port) self.logger.info( f"Initializing weight update group: type={meta.type} " - f"init_method=tcp://{meta.nccl_master_address}:{meta.nccl_master_port} " + f"init_method=tcp://{tcp_addr} " f"group={meta.nccl_group_name}" ) self.weight_update_group = init_custom_process_group( backend=current_platform.communication_backend, world_size=meta.alloc_mode.gen.world_size + 1, - init_method=f"tcp://{meta.nccl_master_address}:{meta.nccl_master_port}", + init_method=f"tcp://{tcp_addr}", rank=0, group_name=meta.nccl_group_name, timeout=DIST_GROUP_DEFAULT_TIMEOUT, @@ -1081,7 +1085,9 @@ def _update_weights_from_distributed(self, meta: WeightUpdateMeta): """Broadcast parameters (chunked) from rank 0 (FSDP2 compatible).""" # Reset weight weight meta with local info - meta.nccl_master_address = self.weight_update_master_addr + # Pre-bracket IPv6 so SGLang's f"tcp://{master_address}:{port}" constructs a valid URL + raw_addr = self.weight_update_master_addr + meta.nccl_master_address = f"[{raw_addr}]" if is_ipv6_address(raw_addr) else raw_addr meta.nccl_master_port = self.weight_update_master_port meta.nccl_group_name = self.weight_update_group_name @@ -1301,7 +1307,9 @@ def _prepare_mb_list(self, input_: dict[str, Any]) -> MicroBatchList: else: input_ = amend_position_ids(input_) - mb_list = split_padded_tensor_dict_into_mb_list(input_, self.config.mb_spec) + mb_list = split_padded_tensor_dict_into_mb_list( + input_, self.config.mb_spec, group=self.cpu_group + ) mb_list.mbs = [pack_tensor_dict(mb) for mb in mb_list.mbs] mb_list = pad_mb_list( mb_list, From d2e358db0e0b24bc0a293cf50d574c3ccb48e6c4 Mon Sep 17 00:00:00 2001 From: "luhongyu.4869" Date: Mon, 9 Mar 2026 21:53:25 -0700 Subject: [PATCH 16/18] fix: eliminate data race in concurrent ScaffoldingLlm episodes When multiple arun_episode coroutines run concurrently they previously shared mutable state on PipelineTrajectoryMaker (task_data, prompt_str, input_tokens), causing later episodes to overwrite earlier ones before generation completed. - ScaffoldingLlm.generate_async: clone prototype_controller synchronously before any async handoff so each episode gets its own isolated clone - ScaffoldingWorkflow.arun_episode: pass per-episode data as kwargs to generate_async instead of mutating shared trajectory_maker fields - PipelineTrajectoryMaker.process: pop task_data/prompt_str/input_tokens from kwargs so each cloned controller sees its own episode data Co-Authored-By: Claude Sonnet 4.6 --- areal/experimental/scaffolding/controllers.py | 16 +++++++++++++--- .../scaffolding/core/scaffolding_llm.py | 10 +++++++--- areal/experimental/scaffolding/workflow.py | 17 +++++++++-------- 3 files changed, 29 insertions(+), 14 deletions(-) diff --git a/areal/experimental/scaffolding/controllers.py b/areal/experimental/scaffolding/controllers.py index 77db36a6b2..ce55baee59 100644 --- a/areal/experimental/scaffolding/controllers.py +++ b/areal/experimental/scaffolding/controllers.py @@ -501,18 +501,28 @@ def process(self, tasks: list[Task], **kwargs) -> Any: reward_tasks = [] interactions = {} + # Per-episode data may be passed via kwargs (from generate_async) to + # avoid race conditions when multiple episodes run concurrently. + effective_task_data = kwargs.pop("task_data", self.task_data) + effective_prompt_str = kwargs.pop("prompt_str", self.prompt_str) + effective_input_tokens = kwargs.pop("input_tokens", self.input_tokens) + for i, task in enumerate(tasks): if isinstance(task, GenerationTask): + # Update task input_tokens from per-episode data if not already set + if not task.input_tokens and effective_input_tokens: + task.input_tokens = effective_input_tokens + # Create interaction object interaction = self._create_interaction_from_task(task) task_id = f"task_{i}" interactions[task_id] = interaction - # Create reward task using constructor-provided task_data and prompt_str + # Create reward task using per-episode task_data and prompt_str reward_task = RLVRRewardTask.create_from_generation_task( gen_task=task, - prompt_str=self.prompt_str or task.input_str or "", - task_data=self.task_data, + prompt_str=effective_prompt_str or task.input_str or "", + task_data=effective_task_data, interaction=interaction, ) reward_tasks.append(reward_task) diff --git a/areal/experimental/scaffolding/core/scaffolding_llm.py b/areal/experimental/scaffolding/core/scaffolding_llm.py index 7e7ac751e2..068c8d3f12 100644 --- a/areal/experimental/scaffolding/core/scaffolding_llm.py +++ b/areal/experimental/scaffolding/core/scaffolding_llm.py @@ -167,16 +167,20 @@ def main_loop_thread(): self.main_loop_thread = threading.Thread(target=main_loop_thread, daemon=True) self.main_loop_thread.start() - def generate_async(self, prompt: str) -> ScaffoldingResult: + def generate_async(self, prompt: str, **kwargs) -> ScaffoldingResult: result = ScaffoldingResult() + # Clone synchronously here (before any async handoff) to avoid race + # conditions where concurrent callers mutate prototype_controller state + # between this call and when put_request actually runs on self.loop. + cloned_controller = self.prototype_controller.clone() async def put_request(): try: request = ScaffoldingRequest( prompt=prompt, - kwargs={}, + kwargs=kwargs, result=result, - controller=self.prototype_controller.clone(), + controller=cloned_controller, ) except Exception as e: await self.task_queue.put(None) diff --git a/areal/experimental/scaffolding/workflow.py b/areal/experimental/scaffolding/workflow.py index 343e57ca89..28c8719f79 100644 --- a/areal/experimental/scaffolding/workflow.py +++ b/areal/experimental/scaffolding/workflow.py @@ -188,14 +188,15 @@ async def arun_episode( ) prompt_str = self.tokenizer.decode(input_ids) - # Configure per-episode data on trajectory maker - # (clone() in scaffolding_llm will deep-copy these) - self.trajectory_maker.task_data = data - self.trajectory_maker.prompt_str = prompt_str - self.trajectory_maker.input_tokens = input_ids - - # Run full pipeline via scaffolding_llm - result = self.scaffolding_llm.generate_async(prompt_str) + # Pass per-episode data as kwargs so generate_async captures them + # in the synchronous clone, avoiding race conditions when multiple + # arun_episode coroutines run concurrently. + result = self.scaffolding_llm.generate_async( + prompt_str, + task_data=data, + prompt_str=prompt_str, + input_tokens=input_ids, + ) await result # Extract interaction and convert to tensor dict From 2e3f299cb870ac920839d5031271810137eb7c71 Mon Sep 17 00:00:00 2001 From: "luhongyu.4869" Date: Mon, 9 Mar 2026 21:53:32 -0700 Subject: [PATCH 17/18] fix: pass sampling_params correctly in GSM8KScaffoldingWorkflow GSM8KScaffoldingWorkflow.build_scaffolding_llm was constructing NativeGenerationController() with no arguments, leaving sampling_params empty. convert_task_params skips None fields, so SGLang never received max_tokens and defaulted to ~16 tokens, producing truncated completions and near-zero rewards (~0.3%). Fix: delegate to super().build_scaffolding_llm(engine) which correctly builds sampling_params from gconfig (max_tokens, temperature, stop). Also clean up debug instrumentation added during investigation. Co-Authored-By: Claude Sonnet 4.6 --- examples/scaffolding/gsm8k_rlvr_scaffolding.py | 10 +--------- 1 file changed, 1 insertion(+), 9 deletions(-) diff --git a/examples/scaffolding/gsm8k_rlvr_scaffolding.py b/examples/scaffolding/gsm8k_rlvr_scaffolding.py index 9cdbf8817d..2a96b16935 100644 --- a/examples/scaffolding/gsm8k_rlvr_scaffolding.py +++ b/examples/scaffolding/gsm8k_rlvr_scaffolding.py @@ -37,15 +37,7 @@ class GSM8KScaffoldingWorkflow(ScaffoldingWorkflow): """ def build_scaffolding_llm(self, engine: InferenceEngine) -> ScaffoldingLlm: - self.gen_controller = NativeGenerationController() - self.reward_controller = RLVRRewardController(self.reward_fn) - self.trajectory_maker = PipelineTrajectoryMaker( - self.gen_controller, self.reward_controller - ) - return ScaffoldingLlm( - self.trajectory_maker, - {NativeGenerationController.WorkerTag.GENERATION: self.worker}, - ) + return super().build_scaffolding_llm(engine) def main(args): From e60cfcdaa0ee9d5ec493efbfe243f9efb11d7043 Mon Sep 17 00:00:00 2001 From: "luhongyu.4869" Date: Mon, 9 Mar 2026 21:53:40 -0700 Subject: [PATCH 18/18] chore: increase timeouts, add attn options, add 2-node GSM8K config MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - cli_args: raise setup_timeout 300→1200s (SGLang model loading on large clusters takes >5 min); add sdpa/eager to attn_impl choices; use sys.executable instead of hardcoded python3 - ray.py: raise startup_timeout 30→600s for multi-node Ray worker join - gsm8k dataset: update prompt to request step-by-step reasoning - chat_scaffolding.yaml: update to 8-GPU xccl config - Add gsm8k_rlvr_scaffolding_2nodes.yaml for 2-node A100 Ray setup Co-Authored-By: Claude Sonnet 4.6 --- areal/api/cli_args.py | 7 +- areal/dataset/gsm8k.py | 2 +- areal/infra/scheduler/ray.py | 2 +- examples/scaffolding/chat_scaffolding.yaml | 6 +- .../gsm8k_rlvr_scaffolding_2nodes.yaml | 187 ++++++++++++++++++ 5 files changed, 196 insertions(+), 8 deletions(-) create mode 100644 examples/scaffolding/gsm8k_rlvr_scaffolding_2nodes.yaml diff --git a/areal/api/cli_args.py b/areal/api/cli_args.py index 7b1539ad7f..49d816fd1b 100644 --- a/areal/api/cli_args.py +++ b/areal/api/cli_args.py @@ -1,6 +1,7 @@ import argparse import json import os +import sys from dataclasses import MISSING as dataclass_missing from dataclasses import asdict, dataclass, field, fields from enum import Enum @@ -806,7 +807,7 @@ class TrainEngineConfig: default="flash_attention_2", metadata={ "help": "Attention implementation for huggingface transformers model.", - "choices": ["flash_attention_2"], + "choices": ["flash_attention_2", "sdpa", "eager"], }, ) init_from_scratch: bool = field( @@ -1092,7 +1093,7 @@ class PPOCriticConfig(TrainEngineConfig): def get_py_cmd(module: str, args: dict[str, Any]): # convert to flags - cmd = ["python3", "-m", module] + cmd = [sys.executable, "-m", module] for k, v in args.items(): if v is None or v is False or v == "" or (isinstance(v, list) and not v): continue @@ -1493,7 +1494,7 @@ class InferenceEngineConfig: metadata={"help": "Whether to dump the trajectories to files under fileroot."}, ) setup_timeout: float = field( - default=300.0, + default=1200.0, metadata={ "help": "Timeout in seconds of connecting to remote servers or launching local servers." }, diff --git a/areal/dataset/gsm8k.py b/areal/dataset/gsm8k.py index 779363b475..d4d6ddb1b3 100644 --- a/areal/dataset/gsm8k.py +++ b/areal/dataset/gsm8k.py @@ -39,7 +39,7 @@ def process(sample): { "role": "user", "content": sample["question"] - + "\nPlease put your final answer within \\boxed{}.", + + "\nPlease reason step by step and provide your final numeric answer.", } ] return {"messages": messages} diff --git a/areal/infra/scheduler/ray.py b/areal/infra/scheduler/ray.py index 04b890504e..9bc1bebff5 100644 --- a/areal/infra/scheduler/ray.py +++ b/areal/infra/scheduler/ray.py @@ -63,7 +63,7 @@ class RayWorkerInfo: class RayScheduler(Scheduler): def __init__( self, - startup_timeout: float = 30.0, + startup_timeout: float = 600.0, *, exp_config: BaseExperimentConfig | None = None, ): diff --git a/examples/scaffolding/chat_scaffolding.yaml b/examples/scaffolding/chat_scaffolding.yaml index 4f9a1b8b4e..7866224c41 100644 --- a/examples/scaffolding/chat_scaffolding.yaml +++ b/examples/scaffolding/chat_scaffolding.yaml @@ -12,13 +12,13 @@ tokenizer_path: ${actor.path} cluster: n_nodes: 1 - n_gpus_per_node: 1 + n_gpus_per_node: 8 fileroot: /tmp/areal/experiments name_resolve: type: nfs nfs_record_root: /tmp/areal/name_resolve -allocation_mode: sglang:d1+d1 +allocation_mode: sglang:d4p1t1+d4p1t1 rollout: experiment_name: ${experiment_name} @@ -76,7 +76,7 @@ actor: adv_norm: mean_level: batch std_level: batch - weight_update_mode: disk + weight_update_mode: xccl max_new_tokens: ${gconfig.max_new_tokens} scheduling_spec: - task_type: worker diff --git a/examples/scaffolding/gsm8k_rlvr_scaffolding_2nodes.yaml b/examples/scaffolding/gsm8k_rlvr_scaffolding_2nodes.yaml new file mode 100644 index 0000000000..2394371567 --- /dev/null +++ b/examples/scaffolding/gsm8k_rlvr_scaffolding_2nodes.yaml @@ -0,0 +1,187 @@ +# RLVR Scaffolding Example Configuration for GSM8K +# Multi-node setup: 2 nodes x 8 A100-80GB GPUs (Ray scheduler) +# allocation: 8 GPUs for SGLang inference, 8 GPUs for actor+ref training + +experiment_name: gsm8k-rlvr-scaffolding +trial_name: trial0 + +seed: 1 +enable_offload: false +total_train_epochs: 10 +tokenizer_path: ${actor.path} + +cluster: + n_nodes: 2 + n_gpus_per_node: 8 + fileroot: /tmp/areal/experiments + name_resolve: + type: ray + ray_actor_name: ray_kv_store + +allocation_mode: sglang:d8+d8 + +scheduler: + type: ray + +rollout: + experiment_name: ${experiment_name} + trial_name: ${trial_name} + max_concurrent_rollouts: 64 + queue_size: null + consumer_batch_size: ${train_dataset.batch_size} + max_head_offpolicyness: 2 + enable_rollout_tracing: false + scheduling_spec: ${actor.scheduling_spec} + fileroot: ${cluster.fileroot} + tokenizer_path: ${tokenizer_path} + dump_to_file: false + +gconfig: + n_samples: 8 + min_new_tokens: 0 + max_new_tokens: 1024 + greedy: false + temperature: 1.0 + +actor: + experiment_name: ${experiment_name} + trial_name: ${trial_name} + path: Qwen/Qwen2.5-3B-Instruct + init_from_scratch: false + attn_impl: sdpa + disable_dropout: true + gradient_checkpointing: true + dtype: bfloat16 + mb_spec: + max_tokens_per_mb: 10240 + optimizer: + type: adam + lr: 1.0e-6 + weight_decay: 0.01 + beta1: 0.9 + beta2: 0.999 + eps: 1e-8 + lr_scheduler_type: constant + gradient_clipping: 1.0 + warmup_steps_proportion: 0.001 + eps_clip: 0.2 + temperature: ${gconfig.temperature} + reward_scaling: 10.0 + reward_bias: -0.5 + kl_ctl: 0.0 + ppo_n_minibatches: 1 + recompute_logprob: true + use_decoupled_loss: true + behav_imp_weight_cap: 5.0 + reward_norm: + mean_level: group + std_level: group + group_size: ${gconfig.n_samples} + adv_norm: + mean_level: batch + std_level: batch + weight_update_mode: xccl + max_new_tokens: ${gconfig.max_new_tokens} + scheduling_spec: + - task_type: worker + port_count: 2 + gpu: 1 + mem: 32 + cmd: python3 -m areal.infra.rpc.rpc_server + env_vars: + NCCL_SOCKET_FAMILY: AF_INET6 + NCCL_SOCKET_IFNAME: eth0 + TORCH_NCCL_HEARTBEAT_TIMEOUT_SEC: "3600" + NCCL_IB_DISABLE: "1" + +ref: + experiment_name: ${experiment_name} + trial_name: ${trial_name} + path: ${actor.path} + init_from_scratch: false + disable_dropout: true + dtype: ${actor.dtype} + mb_spec: + max_tokens_per_mb: 10240 + optimizer: null + scheduling_strategy: + type: colocation + target: actor + scheduling_spec: ${actor.scheduling_spec} + +# SGLang +sglang: + model_path: ${actor.path} + random_seed: ${seed} + skip_tokenizer_init: false + dtype: ${actor.dtype} + max_running_requests: null + context_length: 2048 + mem_fraction_static: 0.5 + attention_backend: flashinfer + +vllm: + model: ${actor.path} + seed: ${seed} + skip_tokenizer_init: false + dtype: ${actor.dtype} + max_model_len: 4096 + gpu_memory_utilization: 0.8 + +# Datasets +train_dataset: + batch_size: 256 + shuffle: true + pin_memory: false + num_workers: 0 + path: openai/gsm8k + type: rl + max_length: 2048 + +valid_dataset: + batch_size: 256 + pin_memory: true + num_workers: 4 + path: openai/gsm8k + type: rl + +# Utilities +saver: + experiment_name: ${experiment_name} + trial_name: ${trial_name} + fileroot: ${cluster.fileroot} + freq_epochs: 1 + freq_steps: null + freq_secs: null + +recover: + mode: disabled + experiment_name: ${experiment_name} + trial_name: ${trial_name} + fileroot: ${cluster.fileroot} + freq_epochs: 1 + freq_steps: null + freq_secs: 3600 + +evaluator: + experiment_name: ${experiment_name} + trial_name: ${trial_name} + fileroot: ${cluster.fileroot} + freq_epochs: 1 + freq_steps: null + freq_secs: null + +stats_logger: + experiment_name: ${experiment_name} + trial_name: ${trial_name} + fileroot: ${cluster.fileroot} + wandb: + mode: disabled + +perf_tracer: + experiment_name: ${experiment_name} + trial_name: ${trial_name} + fileroot: ${cluster.fileroot} + enabled: false + session_tracer: + enabled: false