Scaffolding rollout workspace multi worker

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."
         },

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}

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,

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))

areal/experimental/scaffolding/__init__.py

@@ -0,0 +1,46 @@
+"""
+Scaffolding Framework Integration for AReaL.
+
+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
+- 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
+"""
+
+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__ = [
+    "ScaffoldingWorkflow",
+    "RLVRRewardTask",
+    "RLVRRewardController",
+    "PipelineTrajectoryMaker",
+    "ChatTracer",
+    "TraceTrajectoryMaker",
+    "TraceGenerationTask",
+    "ChatRewardTask",
+    "CreateWorkerFromEngine",
+    "SGLangWorker",
+]

areal/experimental/scaffolding/_compat.py

@@ -0,0 +1,61 @@
+"""Scaffolding framework primitives vendored from TensorRT-LLM.
+
+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 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__ = [
+    "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",
+]