feat: add Arena single- and multi-stream rollouts

areal/api/cli_args.py

@@ -25,6 +25,7 @@
     is_valid_attn_impl,
 )
 from areal.utils import logging, name_resolve, pkg_version
+from areal.utils.config_utils import redact_sensitive_config
 from areal.utils.constants import (
     PROX_LOGP_METHOD_RECOMPUTE,
     PROX_LOGP_METHODS_ALL,
@@ -4200,7 +4201,7 @@ def save_config(cfg, log_dir):
     os.makedirs(log_dir, exist_ok=True)
     config_save_path = os.path.join(log_dir, "config.yaml")
     with open(config_save_path, "w") as f:
-        config_dict: dict = asdict(cfg)
+        config_dict: dict = redact_sensitive_config(asdict(cfg))
         yaml.dump(
             config_dict,
             f,

areal/engine/fsdp_engine.py

@@ -738,7 +738,23 @@ def forward_backward_batch(
         ],
         forward_only: bool = False,
     ) -> None:
-        for mb_item in mb_list:
+        mb_items = list(mb_list)
+        n_micro_batches = len(mb_items)
+        if not self.enable_tree_training:
+            counts = torch.tensor(
+                [n_micro_batches, int(not mb_items)], dtype=torch.int64, device="cpu"
+            )
+            dist.all_reduce(counts, op=dist.ReduceOp.MAX, group=self.cpu_group)
+            if counts[1]:
+                raise ValueError("FSDP requires at least one real micro-batch per rank")
+            n_micro_batches = int(counts[0])
+        # Keep real packing and output metadata intact. Short ranks still need
+        # to participate in every FSDP forward/backward collective. Reusing the
+        # shortest valid input also preserves multimodal and SP/TP input layout.
+        dummy_index = min(range(len(mb_items)), key=mb_list.group_lens.__getitem__)
+        for mb_index in range(n_micro_batches):
+            is_dummy = mb_index >= len(mb_items)
+            mb_item = mb_items[dummy_index if is_dummy else mb_index]
             inputs, ctx = self._prepare_mb_inputs(mb_item)
 
             # Lazily create tree attention metadata just before forward.
@@ -764,8 +780,13 @@ def forward_backward_batch(
             for key in tree_attn_keys:
                 del inputs[key]
 
-            ctx_dict = ctx.to_dict()
-            loss = process_output_fn(logits, ctx_dict)
+            if is_dummy:
+                # No output callback: dummy rows must not affect returned
+                # logprobs, loss normalization, or training statistics.
+                loss = None if forward_only else logits.mean() * 0.0
+            else:
+                ctx_dict = ctx.to_dict()
+                loss = process_output_fn(logits, ctx_dict)
 
             if not forward_only and loss is not None:
                 with trace_scope("fsdp_engine.backward"):
@@ -1943,7 +1964,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, sync_mbs=False
+        )
         mb_list.mbs = [pack_tensor_dict(mb) for mb in mb_list.mbs]
         mb_list = pad_mb_list(
             mb_list,

areal/experimental/openai/client.py

@@ -82,6 +82,10 @@ async def agenerate(self, req: ModelRequest) -> ModelResponse:
 logger = logging.getLogger("OpenAIClient")
 
 
+class ContextLengthExceededError(ValueError):
+    """Raised before generation when the prompt exhausts the context window."""
+
+
 @dataclass
 class _PreparedPrompt:
     input_ids: list[int]
@@ -1145,7 +1149,7 @@ async def create(
                 ):
                     # Remove the interaction from cache on failure
                     del cache[completion_id]
-                raise ValueError(
+                raise ContextLengthExceededError(
                     f"len of prompt tokens {len(prompt_token_ids)} exceeds max_total_tokens {max_total_tokens_final}"
                 )
         if not is_omitted(max_completion_tokens):
@@ -1579,7 +1583,7 @@ async def create(
                 if interaction is not None and cache is not None and resp_id in cache:
                     # Remove the interaction from cache on failure
                     del cache[resp_id]
-                raise ValueError(
+                raise ContextLengthExceededError(
                     f"len of prompt tokens {len(prompt_token_ids)} exceeds engine_max_tokens {self.engine_max_tokens}"
                 )
         if not is_omitted(max_output_tokens):

areal/experimental/openai/proxy/client_session.py

@@ -110,6 +110,11 @@ def __init__(
             self._shared_tensor_resolver = SharedTensorResolver()
         self.session_id: str | None = None
         self._session_api_key: str | None = None
+        self.interaction_count = 0
+        self.context_overflow = False
+        self.context_overflow_message = ""
+        self.system_error = False
+        self.system_error_message = ""
 
     @property
     def session_api_key(self) -> str:
@@ -276,11 +281,18 @@ async def __aexit__(
 
         # Always try to end the session, even on exception
         try:
-            await post_json_with_retry(
+            data = await post_json_with_retry(
                 self._session,
                 url=f"{self.base_url}{RL_END_SESSION_PATHNAME}",
                 headers=self._session_auth_headers(),
             )
+            self.interaction_count = int(data.get("interaction_count", 0))
+            self.context_overflow = bool(data.get("context_overflow", False))
+            self.context_overflow_message = str(
+                data.get("context_overflow_message", "")
+            )
+            self.system_error = bool(data.get("system_error", False))
+            self.system_error_message = str(data.get("system_error_message", ""))
         except Exception as e:
             # Raised errors will be properly handled by OpenAIProxyWorkflow
             logger.warning(f"Failed to end session {self.session_id}: {e}")