diff --git a/areal/api/cli_args.py b/areal/api/cli_args.py index 1f91268545..0a3b965a1c 100644 --- a/areal/api/cli_args.py +++ b/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, diff --git a/areal/engine/fsdp_engine.py b/areal/engine/fsdp_engine.py index 8549aefba3..ed70ce8597 100644 --- a/areal/engine/fsdp_engine.py +++ b/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, diff --git a/areal/experimental/openai/client.py b/areal/experimental/openai/client.py index f16a5e8f7a..453433c4ad 100644 --- a/areal/experimental/openai/client.py +++ b/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): diff --git a/areal/experimental/openai/proxy/client_session.py b/areal/experimental/openai/proxy/client_session.py index 2d1f8f1c4d..648f64f801 100644 --- a/areal/experimental/openai/proxy/client_session.py +++ b/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}") diff --git a/areal/experimental/openai/proxy/proxy_rollout_server.py b/areal/experimental/openai/proxy/proxy_rollout_server.py index ac9b16d4d6..8e8c5a4a5b 100644 --- a/areal/experimental/openai/proxy/proxy_rollout_server.py +++ b/areal/experimental/openai/proxy/proxy_rollout_server.py @@ -6,11 +6,13 @@ import asyncio import hmac import inspect +import json import os import secrets import threading import time -from collections.abc import AsyncGenerator, Mapping +import uuid +from collections.abc import AsyncGenerator, Iterable, Mapping from typing import TYPE_CHECKING, Any import uvicorn @@ -30,7 +32,7 @@ translate_anthropic_response, translate_anthropic_stream, ) -from areal.experimental.openai.client import ArealOpenAI +from areal.experimental.openai.client import ArealOpenAI, ContextLengthExceededError from areal.experimental.openai.types import InteractionWithTokenLogpReward from areal.infra.processor_cache import ProcessorCacheRegistry from areal.infra.rpc.serialization import deserialize_value, serialize_value @@ -47,6 +49,7 @@ DEFAULT_ADMIN_API_KEY, EXPORT_TRAJECTORIES_PATHNAME, GRANT_CAPACITY_PATHNAME, + OPENAI_CHAT_COMPLETIONS_PATHNAME, RESPONSES_PATHNAME, RL_END_PROCESSOR_CACHE_GROUP_PATHNAME, RL_END_SESSION_PATHNAME, @@ -62,6 +65,8 @@ SetRewardRequest, StartSessionRequest, StartSessionResponse, + derive_session_gateway_api_key, + derive_session_gateway_token, serialize_interactions, ) from .tensor_reference import GroupTensorStoreRegistry @@ -73,6 +78,8 @@ logger = getLogger("ProxyRolloutServer") +_STREAM_HEARTBEAT_INTERVAL_SECONDS = 30.0 + # ============================================================================= # Warning Deduplication @@ -199,6 +206,29 @@ async def validate_json_request(raw_request: Request): # ============================================================================= +def _extract_api_tokens(request: Request) -> tuple[str, ...]: + """Extract all supported API credentials without assuming proxy precedence.""" + + tokens: list[str] = [] + auth_header = request.headers.get("authorization", "") + if auth_header.lower().startswith("bearer "): + token = auth_header[7:].strip() + if token: + tokens.append(token) + x_api_key = request.headers.get("x-api-key", "").strip() + if x_api_key and x_api_key not in tokens: + tokens.append(x_api_key) + if not tokens: + raise HTTPException( + status_code=401, + detail=( + "Missing or malformed Authorization header. Expected " + "'Bearer ' or 'x-api-key: '." + ), + ) + return tuple(tokens) + + def _extract_bearer_token(request: Request) -> str: """Extract API token from Authorization header or x-api-key header. @@ -206,17 +236,7 @@ def _extract_bearer_token(request: Request) -> str: per RFC 6750) and 'x-api-key: ' (Anthropic SDK) for cross-SDK compatibility. """ - auth_header = request.headers.get("authorization", "") - if auth_header.lower().startswith("bearer "): - return auth_header[7:].strip() # len("Bearer ") == 7 - # Fallback to x-api-key header (used by Anthropic SDK) - x_api_key = request.headers.get("x-api-key", "") - if x_api_key: - return x_api_key - raise HTTPException( - status_code=401, - detail="Missing or malformed Authorization header. Expected 'Bearer ' or 'x-api-key: '.", - ) + return _extract_api_tokens(request)[0] def _require_admin_key(request: Request) -> str: @@ -231,10 +251,50 @@ def _require_admin_key(request: Request) -> str: def _require_session_key(request: Request) -> str: - """Resolve session_id from the session API key in the Authorization header.""" - token = _extract_bearer_token(request) + """Resolve an active session from a session key or trusted gateway header. + + Normal clients authenticate with their least-privilege per-session API key. + A trusted upstream gateway may instead authenticate with a generation-only + derived key and present a session-bound capability. This lets a public + gateway keep one stable upstream registration per proxy worker without + exposing the admin credential or allowing one Harness to select another + concurrent RL session. + """ + tokens = _extract_api_tokens(request) + gateway_key = derive_session_gateway_api_key(_admin_api_key) with _lock: - session_id = _api_key_to_session.get(token) + if any(hmac.compare_digest(token, gateway_key) for token in tokens): + if request.url.path not in { + f"/{CHAT_COMPLETIONS_PATHNAME}", + f"/{OPENAI_CHAT_COMPLETIONS_PATHNAME}", + f"/{RESPONSES_PATHNAME}", + f"/{ANTHROPIC_MESSAGES_PATHNAME}", + }: + raise HTTPException( + status_code=403, detail="Gateway keys are generation-only." + ) + session_id = request.headers.get("x-session-id", "").strip() + session_token = request.headers.get("x-session-token", "").strip() + expected_token = derive_session_gateway_token(_admin_api_key, session_id) + if ( + session_id + and session_id in _session_to_api_key + and session_token + and hmac.compare_digest(session_token, expected_token) + ): + return session_id + raise HTTPException( + status_code=401, + detail="Missing, invalid, or expired session routing capability.", + ) + session_id = next( + ( + resolved + for token in tokens + if (resolved := _api_key_to_session.get(token)) is not None + ), + None, + ) if session_id is None: raise HTTPException( status_code=401, detail="Invalid or expired session API key." @@ -601,7 +661,14 @@ def end_session(session_id: str = Depends(_require_session_key)): cache_group_id = session.take_processor_cache_group_id() if cache_group_id is not None: _processor_cache_registry.release(cache_group_id) - return {"message": "success", "interaction_count": interaction_count} + return { + "message": "success", + "interaction_count": interaction_count, + "context_overflow": session.context_overflow, + "context_overflow_message": session.context_overflow_message, + "system_error": session.system_error, + "system_error_message": session.system_error_message, + } @app.post( @@ -655,10 +722,19 @@ def set_reward( raise HTTPException(status_code=400, detail="No interactions in session") interaction_id = completions.last_interaction_id elif interaction_id not in completions: - logger.error(f"Interaction {interaction_id} not found in session {session_id}") - raise HTTPException( - status_code=400, detail=f"Interaction {interaction_id} not found" + requested_interaction_id = interaction_id + interaction_id = session_data.stream_completion_aliases.get( + interaction_id, interaction_id ) + if interaction_id not in completions: + logger.error( + f"Interaction {requested_interaction_id} not found in session " + f"{session_id}" + ) + raise HTTPException( + status_code=400, + detail=f"Interaction {requested_interaction_id} not found", + ) session_data.completions.set_reward(interaction_id, reward) return {"message": "success"} @@ -668,6 +744,92 @@ def set_reward( # ============================================================================= +def _contains_image_content(value: Any) -> bool: + """Return whether a nested content value contains an image block.""" + if isinstance(value, (list, tuple)): + return any(_contains_image_content(item) for item in value) + if not isinstance(value, dict): + return False + if value.get("type") in {"image", "image_url", "input_image"}: + return True + return any(_contains_image_content(item) for item in value.values()) + + +def _content_block_text(value: Any) -> str: + """Extract text from nested Claude/OpenAI content blocks.""" + if isinstance(value, str): + return value + if isinstance(value, (list, tuple)): + return "\n".join(filter(None, (_content_block_text(item) for item in value))) + if not isinstance(value, dict): + return str(value) + + text = value.get("text") + if isinstance(text, str): + return text + content = value.get("content") + if isinstance(content, (str, list, tuple, dict)): + return _content_block_text(content) + return json.dumps(value, ensure_ascii=False, sort_keys=True) + + +def _flatten_text_content_lists(messages: list[dict]) -> None: + """Flatten text-only content block lists to strings in-place. + + LiteLLM can forward Claude requests through the OpenAI chat-completions + endpoint while retaining Anthropic-style text block lists. Preserve + multimodal or otherwise structured content for the AReaL client. + """ + for msg in messages: + content = msg.get("content") + if not isinstance(content, list): + continue + + if msg.get("role") == "system" or not _contains_image_content(content): + msg["content"] = _content_block_text(content) + + +def _preprocess_messages(messages: list[dict]) -> list[dict]: + """Normalize messages shared by OpenAI and Anthropic proxy endpoints.""" + _flatten_text_content_lists(messages) + for preprocessor in _message_preprocessors: + messages = preprocessor(messages) + return messages + + +def _materialize_iterables(value: Any) -> Any: + """Recursively convert Pydantic validator iterators to plain containers.""" + if isinstance(value, BaseModel): + return _materialize_iterables(value.model_dump()) + if isinstance(value, Mapping): + return {key: _materialize_iterables(item) for key, item in value.items()} + if isinstance(value, (str, bytes)): + return value + if isinstance(value, Iterable): + return [_materialize_iterables(item) for item in value] + return value + + +def _prepare_request_messages(messages: Any) -> Any: + """Convert request message iterables to mutable dicts and preprocess them.""" + if isinstance(messages, (str, bytes, Mapping)) or not isinstance( + messages, Iterable + ): + return messages + + message_list = list(messages) + normalized: list[dict] = [] + for message in message_list: + if isinstance(message, BaseModel): + message = message.model_dump() + elif isinstance(message, Mapping): + message = dict(message) + else: + return message_list + normalized.append(_materialize_iterables(message)) + return _preprocess_messages(normalized) + + async def _call_client_create( create_fn, request: dict[str, Any] | BaseModel, @@ -712,6 +874,17 @@ async def _call_client_create( ) kwargs = request.model_dump() if isinstance(request, BaseModel) else dict(request) + messages = kwargs.get("messages") + if messages is not None: + try: + prepared_messages = _prepare_request_messages(messages) + except Exception as e: + message = f"{type(e).__name__}: {e}" + session_data.mark_system_error(message) + logger.exception("Failed to normalize proxy request messages") + raise HTTPException(status_code=500, detail=message) from e + kwargs["messages"] = prepared_messages + dropped_args = [] for k, v in kwargs.items(): if k not in areal_client_allowed_args: @@ -778,9 +951,19 @@ def _is_default_value(k: str, v: Any) -> bool: if supports_processor_cache: client_kwargs["processor_cache"] = session_data.processor_cache return await create_fn(**client_kwargs) + except ContextLengthExceededError as e: + session_data.mark_context_overflow(str(e)) + logger.warning("Session %s exceeded its context window: %s", session_id, e) + raise HTTPException( + status_code=400, + detail={"type": "context_length_exceeded", "message": str(e)}, + ) except ValueError as e: + session_data.mark_system_error(f"{type(e).__name__}: {e}") raise HTTPException(status_code=500, detail=str(e)) except Exception as e: + session_data.mark_system_error(f"{type(e).__name__}: {e}") + logger.exception("AReaL client request failed") raise HTTPException(status_code=500, detail=f"{type(e).__name__}: {e}") @@ -789,12 +972,21 @@ def _is_default_value(k: str, v: Any) -> bool: dependencies=[Depends(validate_json_request)], response_model=None, ) +@app.post( + f"/{OPENAI_CHAT_COMPLETIONS_PATHNAME}", + dependencies=[Depends(validate_json_request)], + response_model=None, +) async def chat_completions( - request: CompletionCreateParams, session_id: str = Depends(_require_session_key) + request: dict[str, Any], session_id: str = Depends(_require_session_key) ) -> ChatCompletion | StreamingResponse: """OpenAI-compatible chat completions endpoint. Supports both streaming (stream=True) and non-streaming requests. + The wire body stays an untyped mapping here because OpenAI-compatible + gateways may place multimodal content in roles that the SDK TypedDicts + validate as text-only. The AReaL client performs the actual normalization + after session routing. For streaming requests, returns a StreamingResponse with Server-Sent Events in the OpenAI streaming format (data: {json}\\n\\n ... data: [DONE]\\n\\n). """ @@ -808,39 +1000,18 @@ async def chat_completions( is_streaming = request.get("stream") is True if is_streaming: - openai_stream = None - try: - openai_stream = await _call_client_create( - create_fn=_openai_client.chat.completions.create, - request=request, - session_id=session_id, - stream=True, - ) - - # Convert ChatCompletionChunk objects to OpenAI SSE format - async def _openai_sse_generator( - chunk_stream: AsyncGenerator[ChatCompletionChunk, None], - ) -> AsyncGenerator[str, None]: - async for chunk in chunk_stream: - yield f"data: {chunk.model_dump_json()}\n\n" - yield "data: [DONE]\n\n" - - safe_stream = _safe_stream_wrapper(_openai_sse_generator(openai_stream)) - - return StreamingResponse( - safe_stream, - media_type="text/event-stream", - headers={ - "Cache-Control": "no-cache", - "Connection": "keep-alive", - "X-Accel-Buffering": "no", - }, - ) - except Exception as e: - if openai_stream is not None and hasattr(openai_stream, "aclose"): - await openai_stream.aclose() - logger.error(f"Error setting up streaming response: {e}") - raise HTTPException(status_code=500, detail=f"Streaming setup failed: {e}") + safe_stream = _safe_stream_wrapper( + _deferred_openai_sse_stream(request, session_id) + ) + return StreamingResponse( + safe_stream, + media_type="text/event-stream", + headers={ + "Cache-Control": "no-cache", + "Connection": "keep-alive", + "X-Accel-Buffering": "no", + }, + ) return await _call_client_create( create_fn=_openai_client.chat.completions.create, @@ -849,6 +1020,95 @@ async def _openai_sse_generator( ) +async def _deferred_openai_sse_stream( + request: CompletionCreateParams, + session_id: str, +) -> AsyncGenerator[str, None]: + """Start an SSE response before the non-streaming engine finishes. + + ``ArealOpenAI`` currently waits for ``agenerate()`` to finish before it + returns its simulated chunk stream. Emit a protocol-valid empty role + chunk first so gateways and OpenAI-compatible clients can establish the + stream, then send SSE comments while waiting. Comments are ignored by + compliant clients and do not become assistant content. + """ + stream_id = f"chatcmpl-{uuid.uuid4().hex[:29]}" + stream_model = str(request.get("model") or "") + initial_chunk = { + "id": stream_id, + "choices": [ + { + "delta": {"role": "assistant", "content": ""}, + "index": 0, + "finish_reason": None, + } + ], + "created": int(time.time()), + "model": stream_model, + "object": "chat.completion.chunk", + } + yield f"data: {json.dumps(initial_chunk, separators=(',', ':'))}\n\n" + + create_task = asyncio.create_task( + _call_client_create( + create_fn=_openai_client.chat.completions.create, + request=request, + session_id=session_id, + stream=True, + ) + ) + openai_stream: AsyncGenerator[ChatCompletionChunk, None] | None = None + real_completion_id: str | None = None + try: + while not create_task.done(): + try: + openai_stream = await asyncio.wait_for( + asyncio.shield(create_task), + timeout=_STREAM_HEARTBEAT_INTERVAL_SECONDS, + ) + except TimeoutError: + yield ": ping\n\n" + + if openai_stream is None: + openai_stream = await create_task + + async for chunk in openai_stream: + if real_completion_id is None and chunk.id: + real_completion_id = chunk.id + if real_completion_id != stream_id: + with _lock: + session = _session_cache.get(session_id) + if session is not None: + session.stream_completion_aliases[stream_id] = ( + real_completion_id + ) + stable_chunk = chunk.model_copy( + update={"id": stream_id, "model": stream_model} + ) + yield f"data: {stable_chunk.model_dump_json()}\n\n" + yield "data: [DONE]\n\n" + except asyncio.CancelledError: + raise + except Exception: + logger.exception("Error producing streaming response") + error = { + "error": { + "message": "AReaL stream generation failed", + "type": "server_error", + } + } + yield f"data: {json.dumps(error, separators=(',', ':'))}\n\n" + finally: + if not create_task.done(): + create_task.cancel() + try: + await create_task + except asyncio.CancelledError: + pass + if openai_stream is not None: + await openai_stream.aclose() + + @app.post( f"/{RESPONSES_PATHNAME}", dependencies=[Depends(validate_json_request)], @@ -889,10 +1149,8 @@ def _anthropic_tool_error_ids(anthropic_request: dict[str, Any]) -> set[str]: def _translate_anthropic_to_openai_request(anthropic_request: dict[str, Any]) -> dict: """Translate an Anthropic Messages API request to OpenAI format.""" - openai_request = translate_anthropic_request( - anthropic_request, - message_preprocessors=_message_preprocessors, - ) + # _call_client_create applies preprocessors once for both API formats. + openai_request = translate_anthropic_request(anthropic_request) if openai_request is None: raise ValueError("Failed to translate request") openai_request = dict(openai_request) @@ -1017,6 +1275,8 @@ async def anthropic_messages( "X-Accel-Buffering": "no", }, ) + except HTTPException: + raise except Exception as e: # Clean up stream on error during setup if openai_stream is not None and hasattr(openai_stream, "aclose"): diff --git a/areal/experimental/openai/proxy/server.py b/areal/experimental/openai/proxy/server.py index 381f48df7a..fec0e4d1e4 100644 --- a/areal/experimental/openai/proxy/server.py +++ b/areal/experimental/openai/proxy/server.py @@ -3,6 +3,8 @@ from __future__ import annotations import asyncio +import hashlib +import hmac import threading import time from typing import TYPE_CHECKING, Any @@ -20,6 +22,32 @@ # Session timeout for cleanup (1 hour) SESSION_TIMEOUT_SECONDS = 3600 +_SESSION_GATEWAY_KEY_CONTEXT = b"areal-session-gateway-key-v1" +_SESSION_GATEWAY_TOKEN_CONTEXT = b"areal-session-gateway-token-v1:" + + +def derive_session_gateway_api_key(admin_api_key: str) -> str: + """Derive a generation-only gateway credential from the admin key.""" + + digest = hmac.new( + admin_api_key.encode(), + _SESSION_GATEWAY_KEY_CONTEXT, + hashlib.sha256, + ).hexdigest() + return f"areal-gateway-{digest}" + + +def derive_session_gateway_token(admin_api_key: str, session_id: str) -> str: + """Bind a public-gateway request to exactly one proxy session.""" + + gateway_key = derive_session_gateway_api_key(admin_api_key) + digest = hmac.new( + gateway_key.encode(), + _SESSION_GATEWAY_TOKEN_CONTEXT + session_id.encode(), + hashlib.sha256, + ).hexdigest() + return f"areal-session-{digest}" + # ============================================================================= # Request/Response Models @@ -117,6 +145,11 @@ def __init__( self._last_access_time = time.time() self._end_time = None self._lock = threading.Lock() + self.stream_completion_aliases: dict[str, str] = {} + self.context_overflow = False + self.context_overflow_message = "" + self.system_error = False + self.system_error_message = "" self._next_sampling_request_index = 0 self._processor_cache_released = False @@ -151,6 +184,19 @@ def finish(self): self._end_time = time.time() self._completed_event.set() + def mark_context_overflow(self, message: str) -> None: + """Mark the session as recoverable after a context-length failure.""" + with self._lock: + self.context_overflow = True + self.context_overflow_message = message + + def mark_system_error(self, message: str) -> None: + """Record an internal proxy failure for downstream rejection policy.""" + with self._lock: + self.system_error = True + if not self.system_error_message: + self.system_error_message = message + @property def is_completed(self) -> bool: """Whether this session has been completed via ``finish()``.""" @@ -247,6 +293,7 @@ def deserialize_interactions( RL_FETCH_SHARED_TENSORS_PATHNAME = "rl/fetch_shared_tensors" RL_SET_REWARD_PATHNAME = "rl/set_reward" CHAT_COMPLETIONS_PATHNAME = "chat/completions" +OPENAI_CHAT_COMPLETIONS_PATHNAME = "v1/chat/completions" RESPONSES_PATHNAME = "responses" ANTHROPIC_MESSAGES_PATHNAME = "v1/messages" GRANT_CAPACITY_PATHNAME = "grant_capacity" diff --git a/areal/experimental/openai/proxy/workflow.py b/areal/experimental/openai/proxy/workflow.py index b0f3d48210..f86c865686 100644 --- a/areal/experimental/openai/proxy/workflow.py +++ b/areal/experimental/openai/proxy/workflow.py @@ -4,10 +4,11 @@ import asyncio import atexit +import inspect import os import threading from concurrent.futures import ProcessPoolExecutor -from typing import TYPE_CHECKING, Any +from typing import TYPE_CHECKING, Any, Literal import aiohttp @@ -22,6 +23,8 @@ GRANT_CAPACITY_PATHNAME, RL_END_PROCESSOR_CACHE_GROUP_PATHNAME, ProcessorCacheGroupRequest, + derive_session_gateway_api_key, + derive_session_gateway_token, ) from .tensor_reference import SharedTensorResolver @@ -33,6 +36,13 @@ logger = logging.getLogger("OpenAIProxyWorkflow") +AgentFailureDisposition = Literal[ + "model_failure_zero", + "system_failure_reject", + "unknown_failure_reject", +] + + # Lazy-initialized process pool for running agent tasks _executor: ProcessPoolExecutor | None = None _executor_lock = threading.Lock() @@ -135,14 +145,33 @@ def __init__( self._shared_tensor_resolver = SharedTensorResolver() @trace_session("run_agent") - async def _run_agent(self, session_api_key: str, data: dict): + async def _run_agent( + self, + session_api_key: str, + data: dict, + *, + session_id: str | None = None, + worker_runtime: Any | None = None, + ): if self.mode == "inline": http_client = await workflow_context.get_httpx_client() extra_kwargs = { "base_url": self.proxy_addr, "http_client": http_client, "api_key": session_api_key, + # The public gateway receives only a generation-scoped key, + # never the control-plane admin credential. + "proxy_gateway_api_key": derive_session_gateway_api_key( + self._admin_api_key + ), } + if session_id is not None: + extra_kwargs["session_id"] = session_id + extra_kwargs["proxy_session_token"] = derive_session_gateway_token( + self._admin_api_key, session_id + ) + if worker_runtime is not None: + extra_kwargs["worker_runtime"] = worker_runtime return await self.agent.run(data, **extra_kwargs) if self.mode == "subproc": extra_envs = { @@ -169,6 +198,78 @@ async def _run_agent(self, session_api_key: str, data: dict): return await self.agent.run(data, **extra_kwargs) raise ValueError(f"Unsupported mode: {self.mode}") + def record_group_metrics( + self, + data: dict[str, Any], + rewards: list[float | None], + group_size: int, + ) -> None: + """Delegate optional group-level metric recording to the inline agent.""" + recorder = getattr(self.agent, "record_group_metrics", None) + if callable(recorder): + recorder(data, rewards, group_size) + + def record_episode_metrics( + self, + data: dict[str, Any], + reward: float, + ) -> None: + """Delegate optional episode metrics after interactions are exported.""" + recorder = getattr(self.agent, "record_episode_metrics", None) + if callable(recorder): + recorder(data, reward) + + async def _call_agent_hook(self, name: str, *args: Any, **kwargs: Any) -> None: + """Call an optional agent lifecycle hook without blocking the event loop.""" + + hook = getattr(self.agent, name, None) + if not callable(hook): + return + result = hook(*args, **kwargs) + if inspect.isawaitable(result): + await result + + def _classify_agent_failure( + self, + error: Exception, + *, + context_overflow: bool, + interaction_count: int, + system_error: bool = False, + ) -> AgentFailureDisposition: + """Classify whether an agent error is a trainable model failure. + + Existing agents retain the historical behavior: a typed proxy context + overflow is recoverable as reward zero. Agents backed by authoritative + external graders may provide ``classify_proxy_failure`` to reject system + and unknown failures even when an earlier model request overflowed. + """ + + if system_error: + return "system_failure_reject" + + classifier = getattr(self.agent, "classify_proxy_failure", None) + if classifier is None: + return ( + "model_failure_zero" if context_overflow else "unknown_failure_reject" + ) + disposition = classifier( + error, + context_overflow=context_overflow, + interaction_count=interaction_count, + ) + valid_dispositions = { + "model_failure_zero", + "system_failure_reject", + "unknown_failure_reject", + } + if disposition not in valid_dispositions: + raise ValueError( + "classify_proxy_failure returned an invalid disposition: " + f"{disposition!r}" + ) + return disposition + async def _grant_capacity(self, session: aiohttp.ClientSession) -> None: """Grant capacity via HTTP.""" url = f"{self.proxy_addr}/{GRANT_CAPACITY_PATHNAME}" @@ -277,27 +378,83 @@ async def arun_episode( processor_cache_group_size=context.group_size, shared_tensor_resolver=self._shared_tensor_resolver, ) + agent_error: Exception | None = None async with proxy_client: # Run the user code. try: - rewards = await self._run_agent(proxy_client.session_api_key, data) - except Exception as exc: - logger.warning( - "Agent task failed (%s: %s). This trajectory will be rejected.", - type(exc).__name__, - exc, - exc_info=True, + rewards = await self._run_agent( + proxy_client.session_api_key, + data, + session_id=proxy_client.session_id, + worker_runtime=engine, ) - raise - - # Assign rewards back according to user code output - if isinstance(rewards, dict): - for completion_id, reward in rewards.items(): - await proxy_client.set_reward(completion_id, reward) - elif isinstance(rewards, float): - await proxy_client.set_last_reward(rewards) + except Exception as exc: + agent_error = exc else: - raise ValueError(f"Invalid reward type: {type(rewards)}") + # Assign rewards back according to user code output + if isinstance(rewards, dict): + for completion_id, reward in rewards.items(): + await proxy_client.set_reward(completion_id, reward) + elif isinstance(rewards, float): + await proxy_client.set_last_reward(rewards) + else: + raise ValueError(f"Invalid reward type: {type(rewards)}") + + failure_disposition: AgentFailureDisposition | None = None + if agent_error is not None: + failure_disposition = self._classify_agent_failure( + agent_error, + context_overflow=proxy_client.context_overflow, + interaction_count=proxy_client.interaction_count, + system_error=proxy_client.system_error, + ) + await self._call_agent_hook( + "record_failure_disposition", + data, + agent_error, + failure_disposition, + ) + + if agent_error is not None and failure_disposition != "model_failure_zero": + logger.warning( + "Agent task failed with disposition %s (%s: %s). This " + "trajectory will be rejected.", + failure_disposition, + type(agent_error).__name__, + agent_error, + exc_info=agent_error, + ) + stats_tracker.get(workflow_context.stat_scope()).scalar( + context_overflow=float(proxy_client.context_overflow), + proxy_system_error=float(proxy_client.system_error), + ) + raise agent_error + + if proxy_client.context_overflow or failure_disposition == "model_failure_zero": + logger.warning( + "Recovering model failure with reward 0: context_overflow=%s, " + "interactions=%d, agent_error=%s", + proxy_client.context_overflow, + proxy_client.interaction_count, + agent_error, + ) + stats_tracker.get(workflow_context.stat_scope()).scalar( + context_overflow=float(proxy_client.context_overflow), + proxy_system_error=float(proxy_client.system_error), + ) + if proxy_client.interaction_count == 0: + await self._call_agent_hook( + "persist_episode_result", + data, + None, + ) + return None + await proxy_client.set_last_reward(0.0) + else: + stats_tracker.get(workflow_context.stat_scope()).scalar( + context_overflow=0.0, + proxy_system_error=float(proxy_client.system_error), + ) # Apply turn-level discount and export interactions interactions = await proxy_client.export_interactions( @@ -309,14 +466,27 @@ async def arun_episode( if not interactions: logger.warning( - "Proxy export returned no interactions; rejecting trajectory" + "Task %s exported no usable interactions; trajectory will be rejected.", + task_id, + ) + await self._call_agent_hook( + "persist_episode_result", + data, + None, ) return None # Record stats - last_id = list(interactions.keys())[-1] if interactions else None - if last_id and interactions: + last_id = list(interactions.keys())[-1] + if last_id: last_reward = interactions[last_id].reward - stats_tracker.get(workflow_context.stat_scope()).scalar(reward=last_reward) + await self._call_agent_hook( + "persist_episode_result", + data, + last_reward, + ) + tracker = stats_tracker.get(workflow_context.stat_scope()) + tracker.scalar(reward=last_reward) + self.record_episode_metrics(data, last_reward) return interactions diff --git a/areal/infra/remote_inf_engine.py b/areal/infra/remote_inf_engine.py index b3f5446a1d..5b70ea2905 100644 --- a/areal/infra/remote_inf_engine.py +++ b/areal/infra/remote_inf_engine.py @@ -101,6 +101,9 @@ async def arun_episode( processor_cache=shared_processor_cache, ) + completed_results: list[Any] = [None] * self.group_size + group_error: Exception | None = None + async def run_sample(sample_idx: int) -> tuple[int, Any]: workflow_context.set( WorkflowContext( @@ -112,6 +115,7 @@ async def run_sample(sample_idx: int) -> tuple[int, Any]: ) ) result = await self.workflow.arun_episode(engine, data) + completed_results[sample_idx] = result return sample_idx, result group_tasks = [ @@ -120,12 +124,14 @@ async def run_sample(sample_idx: int) -> tuple[int, Any]: ] try: indexed_results = await asyncio.gather(*group_tasks) - except BaseException: + except BaseException as exc: # gather does not cancel siblings when a child fails or is cancelled. for task in group_tasks: if not task.done(): task.cancel() - raise + if not isinstance(exc, Exception): + raise + group_error = exc finally: try: # Drain cancellation handlers before releasing shared resources. @@ -143,16 +149,33 @@ async def run_sample(sample_idx: int) -> tuple[int, Any]: type(exc).__name__, exc, ) - indexed_results.sort(key=lambda item: item[0]) - sample_indices = [sample_idx for sample_idx, _ in indexed_results] - if sample_indices != list(range(self.group_size)): - raise RuntimeError( - "Grouped rollout returned invalid sample indices: " - f"expected {list(range(self.group_size))}, got {sample_indices}" - ) - results = [result for _, result in indexed_results] - + if group_error is None: + indexed_results.sort(key=lambda item: item[0]) + sample_indices = [sample_idx for sample_idx, _ in indexed_results] + if sample_indices != list(range(self.group_size)): + raise RuntimeError( + "Grouped rollout returned invalid sample indices: " + f"expected {list(range(self.group_size))}, got {sample_indices}" + ) + results = [result if result else None for result in completed_results] valid_results = [r for r in results if r is not None] + interaction_group = not valid_results or all( + isinstance(result, dict) + and all( + isinstance(v, InteractionWithTokenLogpReward) for v in result.values() + ) + for result in valid_results + ) + if interaction_group: + record_group_metrics = getattr(self.workflow, "record_group_metrics", None) + if callable(record_group_metrics): + rewards = [ + result[next(reversed(result))].reward if result else None + for result in results + ] + record_group_metrics(data, rewards, self.group_size) + if group_error is not None: + raise group_error # All results None -> return None if not valid_results: @@ -481,6 +504,8 @@ def __init__( self._workflow_executor: WorkflowExecutor | None = None self._initialized = False + self._destroy_callbacks: dict[str, Callable[[], None]] = {} + self._destroy_callbacks_lock = Lock() self._proxy_gateway_addr: str | None = None self.local_server_processes: list[LocalInfServerInfo] = [] @@ -613,10 +638,38 @@ def initialize( def destroy(self): """Destroy the engine and clean up resources.""" self._initialized = False - if self._workflow_executor is not None: - self._workflow_executor.destroy() - if len(self.local_server_processes) > 0: - self.teardown_server() + try: + if self._workflow_executor is not None: + self._workflow_executor.destroy() + finally: + self._run_destroy_callbacks() + if len(self.local_server_processes) > 0: + self.teardown_server() + + def register_destroy_callback(self, key: str, callback: Callable[[], None]) -> None: + """Register one idempotent worker-lifetime cleanup callback. + + Agent workflows are reconstructed for every submitted prompt, while a + ``RemoteInfEngine`` instance lives for the whole rollout worker. This + hook lets agents retain bounded worker-scoped resources without leaking + them when the worker is destroyed. + """ + + if not key: + raise ValueError("destroy callback key must not be empty") + with self._destroy_callbacks_lock: + self._destroy_callbacks.setdefault(key, callback) + + def _run_destroy_callbacks(self) -> None: + with self._destroy_callbacks_lock: + callbacks = list(self._destroy_callbacks.items()) + self._destroy_callbacks.clear() + for key, callback in reversed(callbacks): + try: + callback() + except Exception: + callback_logger = getattr(self, "logger", logger) + callback_logger.exception("Destroy callback %s failed", key) @property def workflow_executor(self) -> WorkflowExecutor: diff --git a/areal/utils/config_utils.py b/areal/utils/config_utils.py new file mode 100644 index 0000000000..8b2a153978 --- /dev/null +++ b/areal/utils/config_utils.py @@ -0,0 +1,46 @@ +# SPDX-License-Identifier: Apache-2.0 + +"""Utilities for safely serializing application configuration.""" + +from __future__ import annotations + +from typing import Any + +REDACTED_VALUE = "" + + +def _is_sensitive_key(key: object) -> bool: + if not isinstance(key, str): + return False + normalized = key.lower() + return ( + normalized + in { + "authorization", + "password", + "secret", + "token", + # The encoded Stream payload may contain per-task environment + # credentials whose nested keys are invisible to recursive redaction. + "arena_streams_yaml_b64", + } + or "api_key" in normalized + or normalized.endswith(("_password", "_secret", "_token", "_credential")) + or "private_key" in normalized + ) + + +def redact_sensitive_config(value: Any) -> Any: + """Return a copy with credentials redacted while preserving ordinary fields.""" + if isinstance(value, dict): + return { + key: REDACTED_VALUE + if _is_sensitive_key(key) and item not in (None, "") + else redact_sensitive_config(item) + for key, item in value.items() + } + if isinstance(value, list): + return [redact_sensitive_config(item) for item in value] + if isinstance(value, tuple): + return tuple(redact_sensitive_config(item) for item in value) + return value diff --git a/areal/utils/data.py b/areal/utils/data.py index 29408a7495..d17bb5d118 100644 --- a/areal/utils/data.py +++ b/areal/utils/data.py @@ -737,6 +737,8 @@ def split_padded_tensor_dict_into_mb_list( data: dict[str, Any], mb_spec: MicroBatchSpec, group: dist.ProcessGroup | None = None, + *, + sync_mbs: bool = True, ) -> MicroBatchList: """Split a padded dict of tensors into micro-batches based on the attention mask. @@ -744,6 +746,8 @@ def split_padded_tensor_dict_into_mb_list( data (Dict): Dictionary containing padded tensors. mb_spec (MicroBatchSpec): Specification for micro-batch splitting. group (Optional[dist.ProcessGroup]): Process group for distributed synchronization. + sync_mbs: Synchronize micro-batch counts across ranks. Engines that pad + execution with zero-contribution forwards can disable this. Returns: MicroBatchList: A structure containing the split micro-batches and metadata. @@ -788,7 +792,11 @@ def split_padded_tensor_dict_into_mb_list( not_to_split[key] = value # split - group_indices = allocate_balanced_mbs_synced(mb_spec, input_lens, group=group) + group_indices = ( + allocate_balanced_mbs_synced(mb_spec, input_lens, group=group) + if sync_mbs + else allocate_balanced_mbs(mb_spec, input_lens) + ) group_indices = [ seqpack.flat2d( [list(range(i * granularity, (i + 1) * granularity)) for i in group_index] diff --git a/areal/utils/logging.py b/areal/utils/logging.py index 5f79181e64..6f228719e4 100644 --- a/areal/utils/logging.py +++ b/areal/utils/logging.py @@ -116,6 +116,7 @@ "PRMRunner": "light_purple", "ProxyUtils": "light_purple", "AReaL-SWEAgent": "light_purple", + "ArenaStreamAgent": "light_purple", "SWETrain": "light_green", # Agent Service - purple "AgentGateway": "light_purple", diff --git a/areal/utils/network.py b/areal/utils/network.py index 8cf9f5955e..c469ee6982 100644 --- a/areal/utils/network.py +++ b/areal/utils/network.py @@ -152,10 +152,13 @@ def find_free_ports( max_attempts = count * 10 # Reasonable limit to avoid infinite loops attempts = 0 + # Network resources must not follow or advance the training RNG: jobs with + # the same training seed otherwise compete for the same port sequence. + port_rng = random.SystemRandom() while len(free_ports) < count and attempts < max_attempts: # Generate random port within range - port = random.randint(min_port, max_port) + port = port_rng.randint(min_port, max_port) # Skip if port already attempted or excluded if port in attempted_ports or port in exclude_ports: diff --git a/areal/utils/stats_logger.py b/areal/utils/stats_logger.py index 220e6d3920..411ea0402c 100644 --- a/areal/utils/stats_logger.py +++ b/areal/utils/stats_logger.py @@ -14,6 +14,7 @@ from areal.api import FinetuneSpec from areal.api.cli_args import BaseExperimentConfig, StatsLoggerConfig from areal.utils import logging +from areal.utils.config_utils import redact_sensitive_config from areal.utils.printing import tabulate_stats from areal.version import version_info @@ -52,7 +53,7 @@ def init(self): if suffix == "timestamp": suffix = time.strftime("%Y_%m_%d_%H_%M_%S") - exp_config_dict = asdict(self.exp_config) + exp_config_dict = redact_sensitive_config(asdict(self.exp_config)) exp_config_dict["version_info"] = { "commit_id": version_info.commit, "branch": version_info.branch, diff --git a/examples/__init__.py b/examples/__init__.py new file mode 100644 index 0000000000..bbc5468a5c --- /dev/null +++ b/examples/__init__.py @@ -0,0 +1 @@ +"""Example entrypoints and training recipes for AReaL.""" diff --git a/examples/swe/README.md b/examples/swe/README.md index bcacc14a80..c6e6cca548 100644 --- a/examples/swe/README.md +++ b/examples/swe/README.md @@ -1,5 +1,8 @@ # SWE-bench RL training with AReaL-SWEAgent +For Arena-hosted single- or multi-Stream training, see +[Arena integration](README_arena.md). + This example runs SWE-bench coding-agent RL (GRPO) in AReaL. The actual agent loop, sandboxing and reward computation live in a **separate repository**, [AReaL-SWEAgent](https://github.com/areal-project/AReaL-SWEAgent): for each SWE-bench diff --git a/examples/swe/README_arena.md b/examples/swe/README_arena.md new file mode 100644 index 0000000000..f37e003a67 --- /dev/null +++ b/examples/swe/README_arena.md @@ -0,0 +1,185 @@ +# Arena single- and multi-Stream training + +Arena supplies datasets, runs the external Harness, and returns terminal rewards. AReaL +serves the policy through its rollout proxy and trains on the recorded tokens. This +integration does not require AReaL-SWEAgent, PRM, or RewardSystem. + +## Configuration + +`arena_grpo.yaml` is a two-node FSDP/SGLang smoke profile: eight training GPUs, eight +inference GPUs, two samples per prompt, and one training step. The request budget is +32767 tokens against a 32768-token SGLang context, leaving its required one-token +margin. Set these environment variables to your deployment's values: + +```bash +export AREAL_DIR=/path/to/AReaL +export AREAL_IMAGE=/path/to/compatible-image.sif +export AREAL_PYTHON=/path/to/python/in/container +export AREAL_FILEROOT=/shared/path/to/experiments +export MODEL_PATH=/shared/path/to/Qwen3-4B-Instruct +export ARENA_CREDENTIALS_FILE=/shared/private/arena.env +``` + +Keep only credentials and the Arena API base in the private credentials file. It must +export `ARENA_OPENAPI_BASE`, `ARENA_OPENAPI_TOKEN`, `ARENA_LLM_API_KEY`, and +`SWE_RL_ADMIN_API_KEY`. Load it in the controller environment as well. Workers load the +same file; credentials are not embedded in worker commands or committed YAML. + +`ARENA_LLM_API_KEY` must authenticate to Arena's model gateway. It is not an upstream +model-provider key. If your deployment uses the same Arena bearer token for both APIs, +set `export ARENA_LLM_API_KEY="$ARENA_OPENAPI_TOKEN"` in the credentials file. +Control-plane preflight success alone does not validate model-gateway authentication. + +The example assumes a compatible container with Slurm available to the controller, +shared storage, and network reachability from Arena to rollout proxies. Use the included +submission script after setting the site-specific mounts: + +```bash +# Comma-separated Apptainer binds supplied by your cluster setup. +# Workers need the shared checkout, model, credentials and output paths. +export AREAL_WORKER_MOUNTS="$SITE_SHARED_MOUNTS" +# Controller additionally needs working Slurm commands and authentication. +export AREAL_CONTROLLER_MOUNTS="$SITE_SHARED_MOUNTS,$SITE_SLURM_MOUNTS" +export AREAL_CONTAINER_BIN=apptainer # or singularity, available on compute nodes +export SBATCH_PARTITION="$SITE_PARTITION" +export SBATCH_TIMELIMIT=01:00:00 # worker limit +export ARENA_CONTROLLER_TIME=01:30:00 +export ARENA_STREAM_ID=your-stream + +bash examples/swe/submit_arena.sh --check-arena +bash examples/swe/submit_arena.sh +``` + +The script prints the submitted Slurm job ID. It allocates a **CPU-only controller** (4 +CPUs, 16 GB); config validation and Arena preflight run inside the container before +AReaL submits any GPU workers. `--check-arena` stops after preflight. Check +`$AREAL_FILEROOT/submissions/$TRIAL_NAME/controller-.log` and `preflight.json`; +set `TRIAL_NAME` explicitly to identify a run, otherwise a unique name is generated and +used as the controller job name. Use `squeue` to find the job and its name. + +`CONFIG_PATH` defaults to `examples/swe/arena_grpo.yaml`. For multi-stream training: + +```bash +export CONFIG_PATH=examples/swe/arena_multi_stream.yaml +# Fill in the Stream/Harness/reward entries before submitting. +bash examples/swe/submit_arena.sh +``` + +Use a custom YAML to change the GPU topology or training settings. Both preflight and +training resolve that same file, including its Hydra defaults. The supplied profile uses +two GPU nodes; the controller does not reserve those nodes itself. Optional Slurm site +settings use the usual `SBATCH_ACCOUNT`, `SBATCH_RESERVATION`, and related environment +variables. The worker scheduler retains its existing default shared-storage bind; +`AREAL_WORKER_MOUNTS` adds the mounts required by your site. + +Keep the checkout and configuration unchanged while the job is queued or running; use a +dedicated checkout for each revision. Scripts do not copy or snapshot code. All paths +must be visible at the same locations inside the containers. The credentials file is +sourced only in the controller container and workers, never printed or embedded in +generated commands. In the controller, explicit launch paths, Stream ID and API base +take precedence over older credential files that also contain deployment settings. +Workers source the file again, so remove unrelated deployment settings from it. + +For one Stream, set `ARENA_STREAM_ID`. For multiple Streams, edit the literal +`econfig.arena_streams` entries in `arena_multi_stream.yaml` and select that config. Pin +each Stream's actual Harness and reward key/version. Every group of `n_samples` uses one +Stream. The default epoch interleaves all source rows without replacement; a positive +`arena_mixture_epoch_size` selects a weighted subset, must not exceed the union, and +must be divisible by the training batch size. Rejection can change the mixture that +actually reaches training. + +Before submission, extract the Stream list and validate it without launching tasks: + +```bash +python -m examples.swe.arena_stream_config_projection \ + examples/swe/arena_multi_stream.yaml | base64 --decode > /tmp/arena-streams.yaml +python -m examples.swe.arena_preflight \ + --streams-file /tmp/arena-streams.yaml --base-url "$ARENA_OPENAPI_BASE" +``` + +The preflight checks Stream/Harness state, reward version, dataset availability, and +recent task health. A new Stream with no history can use `--min-terminal-tasks 0`; this +does not bypass the other checks. + +## Routing and rewards + +`session_gateway` reuses a registration for each rollout worker and binds every request +to an individual proxy session. The registered key only permits generation; it cannot +end sessions, assign rewards, or export trajectories. Registrations are probed while +tasks are active and cleaned up when the worker is destroyed. `gateway` retains the +older per-rollout registration mode; `direct` requires private network access from the +Harness to the proxy. + +Only the result envelope's top-level `score` is used as reward. Per-Stream thresholds +and optional `reward_transform_fn` hooks can transform it. Heterogeneous `raw` payloads +are retained for audit, not parsed for training rewards. Model-attributed failures can +retain interactions with zero reward; system and ambiguous failures are rejected. Result +shards are written under `arena_result_dump_dir` when configured. + +## Port provenance and scope + +This port applies selected changes from `swe-dev` commits `da1da65c7`, `1d0fc7ba0`, +`714f733a0`, `586d7cc06`, `84f7bb341`, `ddb7d7f3c`, `2b994c10c`, and `9ecbde97c`. It +preserves main's processor cache, sample identity, and cancellation cleanup. Training +engine APIs, mean-only reward normalization, internal Astra launch scripts, AWEX/Qwen3.8 +changes, and PRM/RewardSystem are outside this port. + +## Flash V3 Megatron profile + +`arena_flash_v3.yaml` and `submit_arena_flash_v3.sh` adapt the model, resource and +sampling settings from `swe-dev@5478d6384`'s `submit_arena_astra_multi_stream.sh`. The +default is eight shared training/inference nodes, 128K context, batch 32, 12 samples per +prompt, 1500 concurrent rollouts and 500 training steps. Generation leaves one token of +context headroom. `N_NODES` selects the following 8-GPU-node profiles; the CPU +controller is additional: + +| Nodes | Megatron actor | SGLang rollout | +| ----- | ----------------------------- | -------------- | +| 2 | `(attn:d1p2t4c2\|ffn:d1p2e8)` | `d4t4p1` | +| 4 | `(attn:d2p2t4c2\|ffn:d2p2e8)` | `d8t4p1` | +| 8 | `(attn:d2p2t2c8\|ffn:d4p2e8)` | `d16t4p1` | +| 16 | `(attn:d4p2t2c8\|ffn:d8p2e8)` | `d32t4p1` | + +Use a checkpoint with architecture `BailingMoeV3ForCausalLM`. The profile uses the +existing `mbridge` bridge, CP-capable KDA, Adam at `3e-6`, FP32 LM head, and AWEX weight +updates. MTP is disabled and virtual pipeline size is one. Supply the common submission +environment above, then set deployment-specific paths: + +```bash +export MODEL_PATH="$FLASH_V3_CHECKPOINT" +export ROLLOUT_IMAGE="$FLASH_V3_SGLANG_IMAGE" +export AWEX_ROOT="$FLASH_V3_AWEX_SOURCE" +export FLASH_LINEAR_ATTENTION_ROOT="$FLASH_V3_FLA_SOURCE" +# Optional colon-separated module directories, visible inside worker containers: +export ACTOR_RUNTIME_PYTHONPATH="$FLASH_V3_ACTOR_MODULES" +export ROLLOUT_RUNTIME_PYTHONPATH="$FLASH_V3_ROLLOUT_MODULES" + +N_NODES=2 bash examples/swe/submit_arena_flash_v3.sh --print-profile +N_NODES=2 bash examples/swe/submit_arena_flash_v3.sh --check-arena +N_NODES=2 TRAIN_BATCH_SIZE=4 N_SAMPLES=2 TOTAL_TRAIN_STEPS=1 \ + MAX_CONCURRENT_ROLLOUTS=4 bash examples/swe/submit_arena_flash_v3.sh +``` + +The small trial retains the 128K context. Set `CONTEXT_LENGTH` to change it together +with the default generation and microbatch budgets. Adjust worker/controller time limits +for long runs. This is a resource profile, not a guarantee that the selected checkpoint +and sequence lengths fit available GPU memory. + +The default dataset file pins the two source streams, their Harnesses and reward +versions at equal sampling weights. Set `DATASET_CONFIG` to another shared streams YAML +when needed. Submission runs the existing Arena preflight before creating GPU workers; +`--print-profile` neither contacts Arena nor submits a job. + +The runtime must supply Ling3 tool/reasoning parsers and Flash V3 inference support. The +current AWEX plugin accepts SGLang `0.5.9` or `0.5.10.post1`; the external AWEX checkout +must register the Flash V3 converter. FLA must support the checkpoint's KDA `safe_gate` +and `lower_bound` arguments. The script does not install dependencies. GPU workers +explicitly disable expandable-segment allocation for AWEX compatibility. + +This profile keeps main's outcome-only GRPO: rollout group normalization includes +standard deviation, and incomplete groups are dropped. It does not reproduce `swe-dev`'s +PRM/RewardSystem rewards, mean-only group normalization, minimum-eight partial groups, +or adaptive thinking controls. Chat rendering uses the checkpoint's template defaults +and any Harness request overrides. These differences affect training semantics; the +profile is intended for Arena integration validation, not a claim of full recipe +equivalence. diff --git a/examples/swe/arena_agent.py b/examples/swe/arena_agent.py new file mode 100644 index 0000000000..a814efc491 --- /dev/null +++ b/examples/swe/arena_agent.py @@ -0,0 +1,1283 @@ +"""Arena Stream agent workflow using AReaL's OpenAI-compatible proxy.""" + +from __future__ import annotations + +import asyncio +import hmac +import json +import math +import os +import re +import socket +import stat +import uuid +from contextvars import ContextVar +from dataclasses import dataclass +from numbers import Real +from pathlib import Path +from typing import Any +from urllib.parse import urlparse + +import httpx + +from examples.swe.arena_client import ( + ArenaAPIError, + ArenaOpenAPIClient, + ArenaTaskFailedError, + ArenaTaskResult, + LLMProtocol, +) +from examples.swe.arena_config import load_arena_stream_configs +from examples.swe.arena_types import ArenaStreamConfig + +from areal.infra import workflow_context +from areal.utils import logging, stats_tracker +from areal.utils.dynamic_import import import_from_string + +logger = logging.getLogger("ArenaStreamAgent") + +_ARENA_SESSION_ID_PATTERN = re.compile(r"^[A-Za-z0-9][A-Za-z0-9_.:-]{0,199}$") +_WORKER_GATEWAY_REGISTRY_ATTR = "_arena_session_gateway_registry_v1" +_WORKER_GATEWAY_CLEANUP_KEY = "arena-session-gateway-registrations" + + +def _record_arena_metrics(**metrics: float) -> None: + stats_tracker.get(workflow_context.stat_scope()).scalar(**metrics) + + +def _record_arena_domain_reward(reward: float, arena_task_type: str) -> None: + _record_arena_metrics(**{f"{arena_task_type}/reward": float(reward)}) + + +@dataclass +class _WorkerGatewayTarget: + target: tuple[str, str] + deployment_id: str + client: ArenaOpenAPIClient + http_client: httpx.AsyncClient + proxy_gateway_api_key: str + registration_timeout: float + probe_interval: float + last_probe_at: float + known_missing: bool = False + users: int = 0 + probe_task: asyncio.Task[None] | None = None + + +class _WorkerSessionGatewayRegistry: + """Own Arena registrations for one rollout worker lifetime.""" + + def __init__(self, *, persistent: bool) -> None: + self.persistent = persistent + self._lock = asyncio.Lock() + self._targets: dict[tuple[str, str], _WorkerGatewayTarget] = {} + self._closed = False + + async def acquire( + self, + *, + arena_client: ArenaOpenAPIClient, + proxy_base_url: str, + proxy_gateway_api_key: str, + http_client: httpx.AsyncClient, + registration_timeout: float, + probe_interval: float, + ) -> tuple[str, str]: + key = (arena_client.base_url, proxy_base_url) + async with self._lock: + if self._closed: + raise RuntimeError("Arena worker session gateway registry is closed") + entry = self._targets.get(key) + if entry is None: + suffix = uuid.uuid4().hex[:16] + model_name = f"stream-areal-session-{suffix}" + deployment_id = f"session-{suffix}" + try: + target = await self._register_target( + arena_client=arena_client, + model_name=model_name, + deployment_id=deployment_id, + proxy_base_url=proxy_base_url, + proxy_gateway_api_key=proxy_gateway_api_key, + http_client=http_client, + registration_timeout=registration_timeout, + ) + except (Exception, asyncio.CancelledError): + _record_arena_metrics(**{"arena/registration_success": 0.0}) + _record_arena_metrics(**{"arena/call_success": 0.0}) + # A POST may have reached Arena even when its response was + # lost or the local task was cancelled. The generated name + # is exact and DELETE treats 404 as success, so remove any + # uncertain registration before propagating the failure. + try: + await asyncio.wait_for( + arena_client.delete_llm_proxy_async( + model_name, + client=http_client, + timeout=registration_timeout, + ), + timeout=registration_timeout, + ) + _record_arena_metrics(**{"arena/cleanup_success": 1.0}) + except Exception: + logger.exception( + "Failed to delete uncertain worker-scoped Arena " + "session gateway: model_id=%s", + model_name, + ) + _record_arena_metrics(**{"arena/cleanup_success": 0.0}) + raise + entry = _WorkerGatewayTarget( + target=target, + deployment_id=deployment_id, + client=arena_client, + http_client=http_client, + proxy_gateway_api_key=proxy_gateway_api_key, + registration_timeout=registration_timeout, + probe_interval=probe_interval, + last_probe_at=asyncio.get_running_loop().time(), + ) + self._targets[key] = entry + _record_arena_metrics(**{"arena/registration_success": 1.0}) + logger.info( + "Registered worker-scoped Arena session gateway: " + "model_id=%s, proxy_base_url=%s", + target[1], + proxy_base_url, + ) + elif not hmac.compare_digest( + entry.proxy_gateway_api_key, proxy_gateway_api_key + ): + raise RuntimeError( + "Arena proxy gateway key changed within one rollout worker" + ) + else: + entry.client = arena_client + entry.http_client = http_client + entry.registration_timeout = registration_timeout + entry.probe_interval = min(entry.probe_interval, probe_interval) + await self._ensure_target_locked(key, entry) + entry.users += 1 + if entry.probe_task is None: + entry.probe_task = asyncio.create_task( + self._probe_while_active(key, entry), + name=f"arena-model-probe-{entry.target[1]}", + ) + return entry.target + + @staticmethod + async def _register_target( + *, + arena_client: ArenaOpenAPIClient, + model_name: str, + deployment_id: str, + proxy_base_url: str, + proxy_gateway_api_key: str, + http_client: httpx.AsyncClient, + registration_timeout: float, + ) -> tuple[str, str]: + return await asyncio.wait_for( + arena_client.register_llm_proxy_async( + model_name=model_name, + upstream_base_url=proxy_base_url, + upstream_api_key=proxy_gateway_api_key, + deployment_id=deployment_id, + protocol="chat_completions", + client=http_client, + timeout=registration_timeout, + ), + timeout=registration_timeout, + ) + + async def _ensure_target_locked( + self, + key: tuple[str, str], + entry: _WorkerGatewayTarget, + ) -> None: + now = asyncio.get_running_loop().time() + if now - entry.last_probe_at < entry.probe_interval: + if entry.known_missing: + raise ArenaAPIError( + "Arena session gateway registration is unavailable while " + f"restore retries are throttled: model_id={entry.target[1]}" + ) + return + + model_name = entry.target[1] + try: + exists = await asyncio.wait_for( + entry.client.llm_proxy_exists_async( + model_name, + client=entry.http_client, + timeout=entry.registration_timeout, + ), + timeout=entry.registration_timeout, + ) + except asyncio.CancelledError: + raise + except Exception: + # The probe is advisory unless Arena explicitly confirms a 404. + # Throttle failures too, otherwise task concurrency can become a + # control-plane request storm during an Arena outage. + entry.last_probe_at = asyncio.get_running_loop().time() + _record_arena_metrics(**{"arena/registration_probe_success": 0.0}) + logger.warning( + "Arena session gateway probe failed; retaining cached route: " + "model_id=%s", + model_name, + exc_info=True, + ) + if entry.known_missing: + raise ArenaAPIError( + "Arena session gateway registration remains unavailable after " + f"a failed liveness probe: model_id={model_name}" + ) + return + + entry.last_probe_at = asyncio.get_running_loop().time() + _record_arena_metrics(**{"arena/registration_probe_success": 1.0}) + _record_arena_metrics(**{"arena/registration_present": float(exists)}) + if exists: + entry.known_missing = False + return + + entry.known_missing = True + logger.warning( + "Arena session gateway registration is missing (possibly garbage-" + "collected); restoring the original model_id=%s", + model_name, + ) + try: + restored_target = await self._register_target( + arena_client=entry.client, + model_name=model_name, + deployment_id=entry.deployment_id, + proxy_base_url=key[1], + proxy_gateway_api_key=entry.proxy_gateway_api_key, + http_client=entry.http_client, + registration_timeout=entry.registration_timeout, + ) + except asyncio.CancelledError: + entry.last_probe_at = asyncio.get_running_loop().time() + _record_arena_metrics(**{"arena/registration_restore_success": 0.0}) + raise + except Exception: + # Reusing the original name is required by already-queued tasks. + # A concurrent restore or a lost POST response can therefore look + # like a failed create even though the exact route now exists. + try: + restored = await asyncio.wait_for( + entry.client.llm_proxy_exists_async( + model_name, + client=entry.http_client, + timeout=entry.registration_timeout, + ), + timeout=entry.registration_timeout, + ) + except Exception: + restored = False + if not restored: + entry.last_probe_at = asyncio.get_running_loop().time() + _record_arena_metrics(**{"arena/registration_restore_success": 0.0}) + raise + restored_target = entry.target + entry.target = restored_target + entry.known_missing = False + entry.last_probe_at = asyncio.get_running_loop().time() + _record_arena_metrics(**{"arena/registration_restore_success": 1.0}) + logger.info( + "Restored missing Arena session gateway: model_id=%s, proxy_base_url=%s", + model_name, + key[1], + ) + + async def _probe_while_active( + self, + key: tuple[str, str], + expected_entry: _WorkerGatewayTarget, + ) -> None: + try: + while True: + await asyncio.sleep(expected_entry.probe_interval) + async with self._lock: + entry = self._targets.get(key) + if self._closed or entry is not expected_entry or entry.users <= 0: + return + try: + await self._ensure_target_locked(key, entry) + except asyncio.CancelledError: + raise + except Exception: + logger.warning( + "Failed to restore Arena session gateway during " + "background probe: model_id=%s", + entry.target[1], + exc_info=True, + ) + except asyncio.CancelledError: + return + + async def release( + self, + *, + arena_base_url: str, + proxy_base_url: str, + http_client: httpx.AsyncClient, + registration_timeout: float, + ) -> None: + key = (arena_base_url, proxy_base_url) + probe_task: asyncio.Task[None] | None = None + delete_entry: _WorkerGatewayTarget | None = None + async with self._lock: + entry = self._targets.get(key) + if entry is None or entry.users <= 0: + logger.error( + "Arena session gateway reference underflow: proxy_base_url=%s", + proxy_base_url, + ) + _record_arena_metrics(**{"arena/cleanup_success": 0.0}) + return + entry.users -= 1 + if entry.users == 0: + probe_task = entry.probe_task + entry.probe_task = None + if not self.persistent and entry.users == 0: + delete_entry = self._targets.pop(key) + + if probe_task is not None: + probe_task.cancel() + try: + await probe_task + except asyncio.CancelledError: + pass + + if delete_entry is not None: + try: + await asyncio.wait_for( + delete_entry.client.delete_llm_proxy_async( + delete_entry.target[1], + client=http_client, + timeout=registration_timeout, + ), + timeout=registration_timeout, + ) + logger.info( + "Deleted prompt-scoped Arena session gateway: model_id=%s", + delete_entry.target[1], + ) + _record_arena_metrics(**{"arena/cleanup_success": 1.0}) + except Exception: + logger.exception( + "Failed to delete prompt-scoped Arena session gateway: model_id=%s", + delete_entry.target[1], + ) + _record_arena_metrics(**{"arena/cleanup_success": 0.0}) + + def close(self) -> None: + """Synchronously remove persistent registrations during worker destroy.""" + + if self._closed: + return + self._closed = True + entries = list(self._targets.values()) + self._targets.clear() + for entry in entries: + if entry.probe_task is not None: + loop = entry.probe_task.get_loop() + try: + if loop.is_running(): + loop.call_soon_threadsafe(entry.probe_task.cancel) + elif not loop.is_closed(): + entry.probe_task.cancel() + except RuntimeError: + logger.warning( + "Failed to cancel Arena session gateway probe during " + "worker shutdown: model_id=%s", + entry.target[1], + exc_info=True, + ) + try: + entry.client.delete_llm_proxy(entry.target[1]) + logger.info( + "Deleted worker-scoped Arena session gateway: model_id=%s", + entry.target[1], + ) + except Exception: + logger.exception( + "Failed to delete worker-scoped Arena session gateway: model_id=%s", + entry.target[1], + ) + + +def _open_directory_no_symlinks(path: Path) -> int: + """Create/open an absolute directory path without traversing symlinks.""" + + absolute_path = Path(os.path.abspath(path)) + directory_flags = os.O_RDONLY | os.O_DIRECTORY | os.O_NOFOLLOW + directory_fd = os.open(absolute_path.anchor, directory_flags) + try: + for part in absolute_path.parts[1:]: + try: + os.mkdir(part, mode=0o700, dir_fd=directory_fd) + except FileExistsError: + pass + next_fd = os.open(part, directory_flags, dir_fd=directory_fd) + os.close(directory_fd) + directory_fd = next_fd + return directory_fd + except BaseException: + os.close(directory_fd) + raise + + +def _arena_task_type(data: dict[str, Any]) -> str: + arena_task_type = data.get("arena_task_type", "unknown") + if not isinstance(arena_task_type, str) or not arena_task_type.strip(): + return "unknown" + return arena_task_type.strip() + + +class ArenaStreamAgentWorkflow: + """Launch an Arena online task and use its returned reward for RL.""" + + _MODEL_FAILURE_STATUSES_WITH_INTERACTIONS = {"NO_OUTPUT", "TIMEOUT"} + _SYSTEM_FAILURE_STATUSES = ArenaOpenAPIClient.FAILED_TASK_STATUSES - { + "HARNESS_FAILED", + *_MODEL_FAILURE_STATUSES_WITH_INTERACTIONS, + } + + def __init__( + self, + econfig: dict[str, Any] | None = None, + gen_args: dict[str, Any] | None = None, + timeout: float = 3600.0, + ) -> None: + self.econfig = econfig or {} + self.gen_args = gen_args or {} + self.timeout = float(self.econfig.get("timeout", timeout)) + self.registration_timeout = float( + self.econfig.get("arena_registration_timeout", 180.0) + ) + self.registration_probe_interval = float( + self.econfig.get("arena_registration_probe_interval", 60.0) + ) + if not math.isfinite(self.registration_probe_interval) or ( + self.registration_probe_interval <= 0 + ): + raise ValueError( + "arena_registration_probe_interval must be finite and positive" + ) + self.request_timeout = float(self.econfig.get("arena_request_timeout", 60.0)) + if self.request_timeout <= 0: + raise ValueError("arena_request_timeout must be positive") + self.llm_route_mode = str( + self.econfig.get("arena_llm_route_mode", "gateway") or "gateway" + ).strip() + if self.llm_route_mode not in {"direct", "gateway", "session_gateway"}: + raise ValueError( + "arena_llm_route_mode must be 'direct', 'gateway', or 'session_gateway'" + ) + stream_configs = load_arena_stream_configs(self.econfig) + self.stream_configs = {config.name: config for config in stream_configs} + self.default_stream_name = ( + stream_configs[0].name if len(stream_configs) == 1 else "" + ) + self.reward_transforms = { + config.name: ( + import_from_string(config.reward_transform_fn) + if config.reward_transform_fn + else None + ) + for config in stream_configs + } + self._task_result: ContextVar[ArenaTaskResult | None] = ContextVar( + "arena_task_result", default=None + ) + self._task_result_dumped: ContextVar[bool] = ContextVar( + "arena_task_result_dumped", default=False + ) + self.result_dump_dir = str( + self.econfig.get("arena_result_dump_dir", "") or "" + ).strip() + self.result_dump_max_bytes = int( + self.econfig.get("arena_result_dump_max_bytes", 1_000_000) + ) + if self.result_dump_max_bytes <= 0: + raise ValueError("arena_result_dump_max_bytes must be positive") + self._result_dump_lock = asyncio.Lock() + # Unit tests and callers that invoke the agent directly do not provide a + # rollout-worker runtime. Keep their registration scoped to this agent. + # Normal inline RL attaches a persistent registry to RemoteInfEngine. + self._local_session_gateway_registry = _WorkerSessionGatewayRegistry( + persistent=False + ) + self.client = ArenaOpenAPIClient( + base_url=str(self.econfig.get("arena_base_url", "")), + timeout=self.timeout, + poll_interval=float(self.econfig.get("arena_poll_interval", 5.0)), + request_retries=int(self.econfig.get("arena_request_retries", 3)), + ) + + def record_group_metrics( + self, + data: dict[str, Any], + rewards: list[float | None], + group_size: int, + ) -> None: + """Record whether any of the k rollouts solved the Arena task.""" + if len(rewards) != group_size: + raise ValueError( + f"Expected {group_size} Arena rewards for pass@k, got {len(rewards)}" + ) + arena_task_type = _arena_task_type(data) + stream_config = self._stream_config_for_data(data) + pass_count = sum( + reward is not None and self._is_solved_reward(reward, stream_config) + for reward in rewards + ) + group_pass_distribution = { + f"group_pass_{count}": float(count == pass_count) + for count in range(group_size + 1) + } + _record_arena_metrics( + **{ + "pass@k": float(pass_count > 0), + f"{arena_task_type}/pass@k": float(pass_count > 0), + f"stream/{stream_config.name}/pass@k": float(pass_count > 0), + **group_pass_distribution, + } + ) + + @classmethod + def _is_model_attributed_harness_failure(cls, error: ArenaTaskFailedError) -> bool: + """Recognize the Harness' explicit Claude agent-phase failure envelope.""" + + result = error.result + raw = result.raw if result is not None else None + if not isinstance(raw, dict): + return False + detail = raw.get("error") + if not isinstance(detail, str): + return False + normalized = detail.lower() + return ( + "harness: harness agent phase exited with code" in normalized + and "harness: agent phase error: claude reported error" in normalized + ) + + @classmethod + def classify_proxy_failure( + cls, + error: Exception, + *, + context_overflow: bool, + interaction_count: int, + ) -> str: + """Classify Arena failures without parsing grader-owned ``raw`` text. + + A typed local context overflow or the Harness' explicit Claude + agent-phase error envelope attributes the failure to the model-facing + agent lifecycle. Everything ambiguous is rejected until Arena exposes + versioned failure attribution and semantic codes in its stable result + envelope. + """ + + if not isinstance(error, ArenaTaskFailedError): + return "system_failure_reject" + if error.status in cls._SYSTEM_FAILURE_STATUSES: + return "system_failure_reject" + if ( + error.status in cls._MODEL_FAILURE_STATUSES_WITH_INTERACTIONS + and interaction_count > 0 + ): + return "model_failure_zero" + if ( + error.status == "HARNESS_FAILED" + and interaction_count > 0 + and (context_overflow or cls._is_model_attributed_harness_failure(error)) + ): + return "model_failure_zero" + return "unknown_failure_reject" + + def record_episode_metrics( + self, + data: dict[str, Any], + reward: float, + ) -> None: + """Record terminal result metrics after a valid interaction export.""" + result = self._task_result.get() + stream_config = self._stream_config_for_data(data) + arena_task_type = _arena_task_type(data) + metrics = { + "training_score": float(reward), + f"stream/{stream_config.name}/training_score": float(reward), + } + if result is not None and result.score is not None: + metrics.update( + { + "arena_score": result.score, + # Compatibility alias: this has always meant the scalar score, + # never the heterogeneous Arena ``raw`` object. + "raw_reward": result.score, + f"{arena_task_type}/raw_reward": result.score, + "arena_raw_present": float(result.raw is not None), + "arena_trace_present": float(result.trace_id is not None), + "arena_artifacts_present": float(result.artifacts_uri is not None), + f"stream/{stream_config.name}/arena_score": result.score, + } + ) + try: + _record_arena_metrics(**metrics) + if result is not None and result.score is not None: + data_id = str(data.get("data_id") or "").strip() + else: + data_id = "" + if data_id: + stats_tracker.get("debug").scalar( + **{f"task/{data_id}/raw_reward": result.score} + ) + _record_arena_domain_reward(reward, arena_task_type) + finally: + self._task_result.set(None) + self._task_result_dumped.set(False) + + async def persist_episode_result( + self, + data: dict[str, Any], + reward: float | None, + ) -> None: + """Persist the final reward after proxy-level failure recovery.""" + + result = self._task_result.get() + if result is None or self._task_result_dumped.get(): + return + await self._dump_task_result( + data, + self._stream_config_for_data(data), + result, + training_score=None if reward is None else float(reward), + ) + self._task_result_dumped.set(True) + + async def record_failure_disposition( + self, + data: dict[str, Any], + error: Exception, + disposition: str, + ) -> None: + """Audit rejected results now; recovered failures use the final reward hook.""" + + if isinstance(error, ArenaTaskFailedError) and error.result is not None: + self._task_result.set(error.result) + if disposition == "model_failure_zero": + return + await self.persist_episode_result(data, None) + + async def run( + self, + data: dict[str, Any], + **extra_kwargs: Any, + ) -> float: + """Launch the row's task with the current rollout proxy session.""" + self._task_result.set(None) + self._task_result_dumped.set(False) + stream_config = self._stream_config_for_data(data) + stream_id = str(data.get("stream_id") or stream_config.stream_id or "") + data_id = str(data.get("data_id") or "") + llm_protocol: LLMProtocol = data.get( + "llm_protocol", stream_config.llm_protocol or "chat_completions" + ) + proxy_base_url = extra_kwargs.get("base_url") + proxy_api_key = extra_kwargs.get("api_key") + proxy_session_id = str(extra_kwargs.get("session_id") or "").strip() + arena_http_client: httpx.AsyncClient | None = extra_kwargs.get( + "arena_http_client" + ) or extra_kwargs.get("http_client") + + if not stream_id: + raise ValueError("stream_id is required for ArenaStreamAgentWorkflow") + if not data_id: + raise ValueError("data_id is required for ArenaStreamAgentWorkflow") + if not proxy_base_url: + raise ValueError("base_url is required for ArenaStreamAgentWorkflow") + if not proxy_api_key: + raise ValueError("api_key is required for ArenaStreamAgentWorkflow") + if self.llm_route_mode in {"direct", "session_gateway"} and not ( + proxy_session_id and _ARENA_SESSION_ID_PATTERN.fullmatch(proxy_session_id) + ): + raise ValueError( + "a valid session_id is required when arena_llm_route_mode is " + "'direct' or 'session_gateway'" + ) + proxy_gateway_api_key = str( + extra_kwargs.get("proxy_gateway_api_key") or "" + ).strip() + proxy_session_token = str(extra_kwargs.get("proxy_session_token") or "").strip() + worker_runtime = extra_kwargs.get("worker_runtime") + if self.llm_route_mode == "session_gateway" and not ( + proxy_gateway_api_key and proxy_session_token + ): + raise ValueError( + "proxy_gateway_api_key and proxy_session_token are required when " + "arena_llm_route_mode='session_gateway'" + ) + self._validate_row_reward_ref(data, stream_config) + self._validate_row_llm_protocol(llm_protocol, stream_config) + if arena_http_client is None: + async with httpx.AsyncClient(timeout=self.request_timeout) as check_client: + await self._validate_live_reward_ref( + stream_id, stream_config, check_client + ) + else: + await self._validate_live_reward_ref( + stream_id, stream_config, arena_http_client + ) + + suffix = uuid.uuid4().hex[:12] + model_name = f"stream-areal-{suffix}" + deployment_id = str(uuid.uuid4()) + registered_model_id = model_name + registered_url = str(proxy_base_url).rstrip("/") + launch_api_key = str(proxy_api_key) + launch_task_envs = dict(stream_config.task_envs) + registration_created = False + session_gateway_acquired = False + session_gateway_registry: _WorkerSessionGatewayRegistry | None = None + session_gateway_key = str(proxy_base_url).rstrip("/") + owns_client = arena_http_client is None + client = arena_http_client or httpx.AsyncClient(timeout=self.timeout) + launch_recorded = False + launched_task_id: str | None = None + + def record_launch_success(task_id: str | None) -> None: + nonlocal launch_recorded, launched_task_id + launched_task_id = task_id + if not launch_recorded: + _record_arena_metrics(**{"arena/launch_success": 1.0}) + launch_recorded = True + + async def audit_incomplete_task(status: str) -> None: + if not launched_task_id: + return + # TODO(agent): call Arena task cancellation here once a supported + # endpoint is present in the published OpenAPI schema. + incomplete_result = ArenaTaskResult( + task_id=launched_task_id, + status=status, + score=None, + ) + self._task_result.set(incomplete_result) + await self._dump_task_result( + data, + stream_config, + incomplete_result, + training_score=None, + ) + self._task_result_dumped.set(True) + + try: + if self.llm_route_mode == "gateway": + # Registration POST retries can have an uncertain-success + # outcome, so always attempt exact-name cleanup afterward. + registration_created = True + try: + ( + registered_url, + registered_model_id, + ) = await asyncio.wait_for( + self.client.register_llm_proxy_async( + model_name=model_name, + upstream_base_url=str(proxy_base_url), + upstream_api_key=str(proxy_api_key), + deployment_id=deployment_id, + protocol=llm_protocol, + client=client, + timeout=self.registration_timeout, + ), + timeout=self.registration_timeout, + ) + except Exception: + _record_arena_metrics(**{"arena/registration_success": 0.0}) + _record_arena_metrics(**{"arena/call_success": 0.0}) + raise + launch_api_key = self.client.llm_gateway_api_key + _record_arena_metrics(**{"arena/registration_success": 1.0}) + logger.info( + "Registered Arena LLM proxy: model_name=%s, model_id=%s, " + "registered_url=%s, protocol=%s", + model_name, + registered_model_id, + registered_url, + llm_protocol, + ) + elif self.llm_route_mode == "direct": + # The launch API injects these values into the Harness. Route + # directly with the existing least-privilege per-session key; + # this avoids a control-plane model registration for every leaf. + proxy_host = urlparse(registered_url).hostname + if not proxy_host: + raise ValueError( + f"Direct Arena LLM route has no hostname: {registered_url!r}" + ) + no_proxy_hosts = [ + host.strip() + for host in str( + launch_task_envs.get("NO_PROXY") + or launch_task_envs.get("no_proxy") + or "" + ).split(",") + if host.strip() + ] + if proxy_host not in no_proxy_hosts: + no_proxy_hosts.append(proxy_host) + no_proxy = ",".join(no_proxy_hosts) + launch_task_envs.update( + { + "ANTHROPIC_BASE_URL": registered_url, + "ANTHROPIC_API_KEY": launch_api_key, + "OPENAI_BASE_URL": registered_url, + "OPENAI_API_KEY": launch_api_key, + # ARCA installs an egress gateway in Harness sandboxes. + # The rollout proxy is a cluster-private address and + # must bypass that gateway to remain reachable. + "NO_PROXY": no_proxy, + "no_proxy": no_proxy, + } + ) + _record_arena_metrics(**{"arena/direct_route": 1.0}) + logger.info( + "Using direct Arena LLM route: model_name=%s, protocol=%s, " + "session_id=%s", + model_name, + llm_protocol, + proxy_session_id, + ) + else: + session_gateway_registry = self._session_gateway_registry( + worker_runtime + ) + ( + registered_url, + registered_model_id, + ) = await session_gateway_registry.acquire( + arena_client=self.client, + proxy_base_url=session_gateway_key, + proxy_gateway_api_key=proxy_gateway_api_key, + http_client=client, + registration_timeout=self.registration_timeout, + probe_interval=self.registration_probe_interval, + ) + session_gateway_acquired = True + launch_api_key = self.client.llm_gateway_api_key + custom_headers = str( + launch_task_envs.get("ANTHROPIC_CUSTOM_HEADERS") or "" + ).strip() + session_headers = ( + f"X-Session-Id: {proxy_session_id}\n" + f"X-Session-Token: {proxy_session_token}" + ) + launch_task_envs["ANTHROPIC_CUSTOM_HEADERS"] = ( + f"{custom_headers}\n{session_headers}" + if custom_headers + else session_headers + ) + _record_arena_metrics(**{"arena/session_gateway_route": 1.0}) + logger.info( + "Using Arena session gateway route: model_id=%s, " + "protocol=%s, session_id=%s", + registered_model_id, + llm_protocol, + proxy_session_id, + ) + logger.info( + f"Launching Arena task: stream_id={stream_id}, data_id={data_id}" + ) + try: + task_result_or_score = await asyncio.wait_for( + self.client.launch_one_task_result( + stream_id=stream_id, + data_id=data_id, + model_name=registered_model_id, + proxy_base_url=registered_url, + proxy_api_key=launch_api_key, + session_id=proxy_session_id, + harness=stream_config.harness, + task_envs=launch_task_envs, + client=client, + on_launch_success=record_launch_success, + ), + timeout=self.timeout, + ) + except ArenaTaskFailedError as exc: + _record_arena_metrics(**{"arena/terminal_success": 0.0}) + _record_arena_metrics(**{"arena/call_success": 0.0}) + if exc.result is not None: + self._task_result.set(exc.result) + raise + except asyncio.CancelledError: + await audit_incomplete_task("LOCAL_WAIT_CANCELLED") + if launch_recorded: + _record_arena_metrics(**{"arena/terminal_success": 0.0}) + _record_arena_metrics(**{"arena/call_success": 0.0}) + logger.warning( + "Arena task wait cancelled locally: task_id=%s data_id=%s", + launched_task_id or "unknown", + data_id, + ) + raise + except (TimeoutError, asyncio.TimeoutError) as exc: # noqa: UP041 + await audit_incomplete_task("LOCAL_WAIT_TIMEOUT") + logger.warning( + "Arena task wait timed out: task_id=%s data_id=%s " + "timeout_seconds=%s", + launched_task_id or "unknown", + data_id, + self.timeout, + ) + if launch_recorded: + _record_arena_metrics(**{"arena/terminal_success": 0.0}) + else: + _record_arena_metrics(**{"arena/launch_success": 0.0}) + _record_arena_metrics(**{"arena/call_success": 0.0}) + raise TimeoutError from exc + except Exception: + if launch_recorded: + _record_arena_metrics(**{"arena/terminal_success": 0.0}) + else: + _record_arena_metrics(**{"arena/launch_success": 0.0}) + _record_arena_metrics(**{"arena/call_success": 0.0}) + raise + if isinstance(task_result_or_score, ArenaTaskResult): + task_result = task_result_or_score + elif isinstance(task_result_or_score, Real) and not isinstance( + task_result_or_score, bool + ): + # Compatibility for existing custom fakes and older clients. + task_result = ArenaTaskResult( + task_id=launched_task_id or "legacy-result", + status="DONE", + score=float(task_result_or_score), + ) + else: + raise ArenaAPIError( + "Arena launch returned neither ArenaTaskResult nor numeric score" + ) + self._task_result.set(task_result) + try: + if task_result.score is None: + raise ArenaAPIError("Successful Arena task result is missing score") + if not math.isfinite(task_result.score): + raise ArenaAPIError("Successful Arena task score must be finite") + reward = self._transform_reward(task_result.score, data, stream_config) + except Exception: + await self._dump_task_result( + data, + stream_config, + task_result, + training_score=None, + ) + self._task_result_dumped.set(True) + _record_arena_metrics(**{"arena/terminal_success": 1.0}) + _record_arena_metrics(**{"arena/result_processing_success": 0.0}) + _record_arena_metrics(**{"arena/call_success": 0.0}) + raise + _record_arena_metrics(**{"arena/result_processing_success": 1.0}) + _record_arena_metrics(**{"arena/terminal_success": 1.0}) + _record_arena_metrics(**{"arena/call_success": 1.0}) + finally: + if session_gateway_acquired and session_gateway_registry is not None: + await session_gateway_registry.release( + arena_base_url=self.client.base_url, + proxy_base_url=session_gateway_key, + http_client=client, + registration_timeout=self.registration_timeout, + ) + elif registration_created: + try: + await asyncio.wait_for( + self.client.delete_llm_proxy_async( + registered_model_id, + client=client, + timeout=self.registration_timeout, + ), + timeout=self.registration_timeout, + ) + logger.info( + "Deleted Arena LLM proxy registration: model_id=%s", + registered_model_id, + ) + _record_arena_metrics(**{"arena/cleanup_success": 1.0}) + except Exception: + logger.exception( + "Failed to delete Arena LLM proxy registration: model_id=%s", + registered_model_id, + ) + _record_arena_metrics(**{"arena/cleanup_success": 0.0}) + if owns_client: + await client.aclose() + logger.info( + f"Finished Arena task: stream_id={stream_id}, data_id={data_id}, " + f"reward={reward}" + ) + self._task_result.set(task_result) + return reward + + def _session_gateway_registry( + self, worker_runtime: Any | None + ) -> _WorkerSessionGatewayRegistry: + """Return the registry owned by this rollout worker. + + ``RemoteInfEngine`` survives across prompt submissions even though each + submission reconstructs this agent. Keeping the registry on that + runtime makes registration cardinality equal the number of proxy + workers, not the number of prompts or samples. + """ + + if worker_runtime is None: + return self._local_session_gateway_registry + + registry = getattr(worker_runtime, _WORKER_GATEWAY_REGISTRY_ATTR, None) + if registry is None: + registry = _WorkerSessionGatewayRegistry(persistent=True) + setattr(worker_runtime, _WORKER_GATEWAY_REGISTRY_ATTR, registry) + if not isinstance(registry, _WorkerSessionGatewayRegistry): + raise TypeError( + f"{_WORKER_GATEWAY_REGISTRY_ATTR} is not an Arena gateway registry" + ) + + register_cleanup = getattr(worker_runtime, "register_destroy_callback", None) + if not callable(register_cleanup): + raise TypeError( + "worker_runtime must provide register_destroy_callback for " + "persistent Arena session gateway cleanup" + ) + register_cleanup(_WORKER_GATEWAY_CLEANUP_KEY, registry.close) + return registry + + def _stream_config_for_data(self, data: dict[str, Any]) -> ArenaStreamConfig: + stream_name = str(data.get("arena_stream_name") or self.default_stream_name) + if not stream_name: + raise ValueError("arena_stream_name is required for a multi-Stream run") + try: + stream_config = self.stream_configs[stream_name] + except KeyError as exc: + raise ValueError(f"Unknown Arena Stream name: {stream_name!r}") from exc + row_stream_id = str(data.get("stream_id") or "") + if ( + row_stream_id + and stream_config.stream_id + and row_stream_id != stream_config.stream_id + ): + raise ValueError( + f"Arena row Stream mismatch for {stream_name!r}: " + f"{row_stream_id!r} != {stream_config.stream_id!r}" + ) + return stream_config + + @staticmethod + def _validate_row_reward_ref( + data: dict[str, Any], stream_config: ArenaStreamConfig + ) -> None: + expected = stream_config.expected_reward_ref + row_key = str(data.get("reward_ref_key") or "") + row_version = str(data.get("reward_ref_version") or "") + if expected.key and (row_key, row_version) != ( + expected.key, + expected.version, + ): + raise ValueError( + f"Arena row reward_ref mismatch for {stream_config.name!r}: " + f"{row_key}@{row_version} != {expected.key}@{expected.version}" + ) + + @staticmethod + def _validate_row_llm_protocol( + row_protocol: str, stream_config: ArenaStreamConfig + ) -> None: + if row_protocol not in ("anthropic", "responses", "chat_completions"): + raise ValueError(f"Unsupported Arena row LLM protocol: {row_protocol!r}") + if stream_config.llm_protocol and row_protocol != stream_config.llm_protocol: + raise ValueError( + f"Arena row LLM protocol mismatch for {stream_config.name!r}: " + f"{row_protocol!r} != {stream_config.llm_protocol!r}" + ) + + async def _validate_live_reward_ref( + self, + stream_id: str, + stream_config: ArenaStreamConfig, + client: httpx.AsyncClient, + ) -> None: + """Fail before launch when a Stream's default grader has drifted.""" + + expected = stream_config.expected_reward_ref + if not expected.key: + return + try: + stream = await asyncio.wait_for( + self.client.resolve_stream_async( + stream_id, + client=client, + timeout=self.request_timeout, + ), + timeout=self.request_timeout, + ) + actual = stream.get("default_reward_ref") + if not isinstance(actual, dict): + raise ArenaAPIError( + f"Arena Stream {stream_config.name!r} has no default_reward_ref" + ) + actual_key = actual.get("key") + actual_version = actual.get("version") + if (actual_key, actual_version) != (expected.key, expected.version): + raise ArenaAPIError( + f"Arena Stream {stream_config.name!r} reward_ref drifted before " + f"launch: expected {expected.key}@{expected.version}, got " + f"{actual_key}@{actual_version}" + ) + except Exception: + _record_arena_metrics(**{"arena/reward_ref_validation_success": 0.0}) + raise + _record_arena_metrics(**{"arena/reward_ref_validation_success": 1.0}) + + def _transform_reward( + self, + reward: float, + data: dict[str, Any], + stream_config: ArenaStreamConfig | None = None, + ) -> float: + """Apply the configured Arena reward mapping at the RL boundary.""" + if stream_config is None: + stream_config = self._stream_config_for_data(data) + reward_transform = self.reward_transforms[stream_config.name] + if reward_transform is not None: + transform_kwargs = ( + {"reward_threshold": stream_config.reward_threshold} + if stream_config.reward_threshold is not None + else {} + ) + transformed = float(reward_transform(reward, data, **transform_kwargs)) + if not math.isfinite(transformed): + raise ValueError( + f"Arena reward transform for {stream_config.name!r} returned " + f"a non-finite value: {transformed}" + ) + return transformed + if stream_config.reward_threshold is not None: + return float(reward >= stream_config.reward_threshold) + return reward + + async def _dump_task_result( + self, + data: dict[str, Any], + stream_config: ArenaStreamConfig, + result: ArenaTaskResult, + *, + training_score: float | None, + ) -> None: + """Optionally append one bounded result envelope to a private audit shard.""" + + if not self.result_dump_dir: + return + expected_ref = stream_config.expected_reward_ref + + def encode_record() -> bytes: + raw_json = json.dumps(result.raw, ensure_ascii=False, separators=(",", ":")) + raw_size_bytes = len(raw_json.encode("utf-8")) + raw_truncated = raw_size_bytes > self.result_dump_max_bytes + record = { + "arena_stream_name": stream_config.name, + "stream_id": str(data.get("stream_id") or stream_config.stream_id), + "data_id": str(data.get("data_id") or ""), + "task_id": result.task_id, + "status": result.status, + "expected_reward_ref": { + "key": expected_ref.key, + "version": expected_ref.version, + }, + "arena_score": result.score, + "training_score": training_score, + "raw": None if raw_truncated else result.raw, + "raw_present": result.raw is not None, + "raw_size_bytes": raw_size_bytes, + "raw_truncated": raw_truncated, + "artifacts_uri": result.artifacts_uri, + "trace_id": result.trace_id, + "computed_at": result.computed_at, + } + return ( + json.dumps(record, ensure_ascii=False, separators=(",", ":")) + "\n" + ).encode("utf-8") + + try: + encoded = await asyncio.to_thread(encode_record) + except (TypeError, ValueError): + logger.exception( + "Failed to encode Arena result audit record: task_id=%s", + result.task_id, + ) + _record_arena_metrics(**{"arena/result_dump_success": 0.0}) + return + dump_dir = Path(self.result_dump_dir) + dump_path = ( + dump_dir / f"arena_results_{socket.gethostname()}_{os.getpid()}.jsonl" + ) + + def append_record() -> None: + directory_fd = _open_directory_no_symlinks(dump_dir) + try: + directory_stat = os.fstat(directory_fd) + if not stat.S_ISDIR(directory_stat.st_mode): + raise NotADirectoryError(str(dump_dir)) + if directory_stat.st_uid != os.geteuid(): + raise PermissionError( + f"Arena result directory must be owned by uid " + f"{os.geteuid()}: {dump_dir}" + ) + os.fchmod(directory_fd, 0o700) + file_descriptor = os.open( + dump_path.name, + os.O_APPEND | os.O_CREAT | os.O_WRONLY | os.O_NOFOLLOW, + 0o600, + dir_fd=directory_fd, + ) + try: + file_stat = os.fstat(file_descriptor) + if not stat.S_ISREG(file_stat.st_mode): + raise OSError( + f"Arena result shard must be a regular file: {dump_path}" + ) + if file_stat.st_uid != os.geteuid(): + raise PermissionError( + f"Arena result shard must be owned by uid " + f"{os.geteuid()}: {dump_path}" + ) + os.fchmod(file_descriptor, 0o600) + with os.fdopen(file_descriptor, "ab", closefd=False) as output: + output.write(encoded) + finally: + os.close(file_descriptor) + finally: + os.close(directory_fd) + + try: + async with self._result_dump_lock: + await asyncio.to_thread(append_record) + except Exception: + logger.exception( + "Failed to persist Arena result audit record: task_id=%s", + result.task_id, + ) + _record_arena_metrics(**{"arena/result_dump_success": 0.0}) + else: + _record_arena_metrics(**{"arena/result_dump_success": 1.0}) + + def _is_solved_reward( + self, reward: float, stream_config: ArenaStreamConfig + ) -> bool: + if self.reward_transforms[stream_config.name] is None: + # Threshold-only mapping has already converted the score to 0/1. + return reward > 0.0 + if stream_config.reward_threshold is not None: + return reward >= stream_config.reward_threshold + return reward > 0.0 diff --git a/examples/swe/arena_client.py b/examples/swe/arena_client.py new file mode 100644 index 0000000000..3945aa2fea --- /dev/null +++ b/examples/swe/arena_client.py @@ -0,0 +1,1003 @@ +"""Client helpers for the Arena online Stream OpenAPI.""" + +from __future__ import annotations + +import asyncio +import math +import os +import time +import uuid +from collections.abc import Callable, Mapping +from dataclasses import dataclass +from numbers import Real +from typing import Any, Literal +from urllib.parse import quote + +import httpx + +LLMProtocol = Literal["anthropic", "responses", "chat_completions"] + + +class ArenaAPIError(RuntimeError): + """Raised when the Arena OpenAPI returns an invalid or failed response.""" + + +@dataclass(frozen=True) +class ArenaTaskResult: + """Stable Arena task-result envelope. + + ``raw`` is deliberately opaque. Its schema is owned by the configured Arena + reward implementation and may also be ``None`` for successful tasks. + """ + + task_id: str + status: str + score: float | None + raw: Any = None + artifacts_uri: str | None = None + trace_id: str | None = None + computed_at: str | None = None + + +def _arena_task_type_from_tags(tags: Any) -> str: + """Extract a task type from tags, falling back when metadata is unusable.""" + if tags is None: + return "unknown" + if not isinstance(tags, list) or not all(isinstance(tag, str) for tag in tags): + return "unknown" + + task_types = { + tag.removeprefix("domain:").strip() + for tag in tags + if tag.startswith("domain:") and tag.removeprefix("domain:").strip() + } + if len(task_types) > 1: + return "unknown" + return next(iter(task_types), "unknown") + + +class ArenaTaskFailedError(ArenaAPIError): + """Raised when an Arena task reaches a failed terminal status.""" + + def __init__( + self, + task_id: str, + status: str, + payload: Mapping[str, Any] | None = None, + result: ArenaTaskResult | None = None, + ) -> None: + self.task_id = task_id + self.status = status + self.payload = dict(payload or {}) + self.result = result + super().__init__(f"Arena task {task_id!r} failed with {status}") + + +def _is_retryable_status(status_code: int) -> bool: + return status_code == 429 or status_code >= 500 + + +def _is_retryable_response(response: httpx.Response) -> bool: + if _is_retryable_status(response.status_code): + return True + if response.status_code != 403: + return False + return "spanner-http-ant-group-watch-all" in response.text[:1000] + + +def resolve_arena_credentials( + base_url: str = "", + api_token: str = "", +) -> tuple[str, str]: + """Resolve Arena connection settings without embedding credentials in configs.""" + resolved_base_url = base_url or os.getenv("ARENA_OPENAPI_BASE", "") + resolved_api_token = api_token or os.getenv("ARENA_OPENAPI_TOKEN", "") + if not resolved_base_url: + raise ValueError( + "Arena OpenAPI base URL is required; set econfig.arena_base_url or " + "ARENA_OPENAPI_BASE" + ) + if not resolved_api_token: + raise ValueError("ARENA_OPENAPI_TOKEN is required") + return resolved_base_url.rstrip("/"), resolved_api_token + + +def _response_json(response: httpx.Response) -> Any: + try: + response.raise_for_status() + except httpx.HTTPStatusError as exc: + body = response.text[:500] + raise ArenaAPIError( + f"Arena OpenAPI returned HTTP {response.status_code}: {body}" + ) from exc + try: + return response.json() + except ValueError as exc: + raise ArenaAPIError("Arena OpenAPI returned a non-JSON response") from exc + + +def _stream_from_payload(payload: Any) -> dict[str, Any]: + """Extract one Stream from direct-get response envelope variants.""" + + if not isinstance(payload, Mapping): + raise ArenaAPIError("Arena Stream response must be a JSON object") + data = payload.get("data", payload) + if not isinstance(data, Mapping): + raise ArenaAPIError("Arena Stream response is missing the Stream object") + stream = data.get("stream", data) + if not isinstance(stream, Mapping): + raise ArenaAPIError("Arena Stream response is missing the Stream object") + return dict(stream) + + +def _optional_string(payload: Mapping[str, Any], key: str) -> str | None: + value = payload.get(key) + if value is None: + return None + if not isinstance(value, str): + raise ArenaAPIError(f"Arena task result field {key!r} must be a string") + return value + + +def _parse_task_result( + payload: Any, + *, + fallback_task_id: str = "", +) -> ArenaTaskResult: + """Parse only the stable top-level result envelope. + + In particular, no recursive reward discovery is performed: heterogeneous + grader data under ``raw`` must never become an RL reward accidentally. + """ + + # Preserve the legacy launch API's explicit control-plane envelope support + # without recursively searching opaque grader-owned ``raw`` data. + result_payload = payload + control_fields = { + "task_id", + "status", + "score", + "raw", + "artifacts_uri", + "trace_id", + "computed_at", + } + for _ in range(4): + if isinstance(result_payload, Real) and not isinstance(result_payload, bool): + score = float(result_payload) + if not math.isfinite(score): + raise ArenaAPIError("Arena task result 'score' must be finite") + return ArenaTaskResult( + task_id=fallback_task_id or "launch_one_task", + status="DONE", + score=score, + ) + if not isinstance(result_payload, Mapping): + raise ArenaAPIError("Arena task result must be a JSON object or number") + if control_fields.intersection(result_payload): + break + nested = next( + ( + result_payload[key] + for key in ("result", "output", "task", "data") + if key in result_payload + ), + None, + ) + if nested is None: + break + result_payload = nested + if not isinstance(result_payload, Mapping): + raise ArenaAPIError("Arena task result must be a JSON object or number") + + task_id = result_payload.get("task_id", fallback_task_id) + if not isinstance(task_id, str) or not task_id: + raise ArenaAPIError("Arena task result is missing a valid 'task_id'") + + status_value = result_payload.get("status") + if status_value is None and "score" in result_payload: + status = "DONE" + elif status_value is None and isinstance(result_payload.get("task_id"), str): + status = "PENDING" + elif isinstance(status_value, str) and status_value.strip(): + status = status_value.strip().upper() + else: + raise ArenaAPIError("Arena task result is missing a valid 'status'") + + score_value = result_payload.get("score") + if score_value is None: + score = None + elif isinstance(score_value, Real) and not isinstance(score_value, bool): + score = float(score_value) + if not math.isfinite(score): + raise ArenaAPIError("Arena task result 'score' must be finite") + else: + raise ArenaAPIError("Arena task result 'score' must be numeric or null") + + return ArenaTaskResult( + task_id=task_id, + status=status, + score=score, + raw=result_payload.get("raw"), + artifacts_uri=_optional_string(result_payload, "artifacts_uri"), + trace_id=_optional_string(result_payload, "trace_id"), + computed_at=_optional_string(result_payload, "computed_at"), + ) + + +def infer_llm_protocol_from_harness(harness_key: str) -> LLMProtocol: + """Infer the Harness client protocol from a Harness key or reference.""" + + harness_key = harness_key.lower() + if "claude" in harness_key: + return "anthropic" + if "codex" in harness_key: + return "responses" + return "chat_completions" + + +def infer_llm_protocol(stream: Mapping[str, Any]) -> LLMProtocol: + """Select the client protocol used by a Stream's default Harness.""" + harness_ref = stream.get("default_harness_ref") + harness_key = "" + if isinstance(harness_ref, Mapping): + key = harness_ref.get("key") + if isinstance(key, str): + harness_key = key + return infer_llm_protocol_from_harness(harness_key) + + +def resolve_llm_protocol( + stream: Mapping[str, Any], + override: str = "", +) -> LLMProtocol: + """Return an explicit protocol override or infer one from the Harness.""" + if not override: + return infer_llm_protocol(stream) + if override not in ("anthropic", "responses", "chat_completions"): + raise ValueError(f"Unsupported Arena LLM protocol override: {override!r}") + return override + + +def _llm_registration_payload( + model_name: str, + upstream_base_url: str, + upstream_api_key: str, + deployment_id: str, + protocol: LLMProtocol, +) -> dict[str, Any]: + if protocol not in ("anthropic", "responses", "chat_completions"): + raise ValueError(f"Unsupported Arena LLM protocol: {protocol!r}") + + # ``protocol`` describes the Arena Harness client. The AReaL rollout + # proxy itself exposes an OpenAI-compatible Chat Completions endpoint, so + # advertise only that native upstream capability. Arena converts Messages + # or Responses requests to Chat before forwarding them to this endpoint. + return { + "model_name": model_name, + "endpoints": [ + { + "endpoint_id": deployment_id, + "upstream_model": model_name, + "base_url": upstream_base_url.rstrip("/"), + "api_key": upstream_api_key, + "inbound_protos": ["chat"], + "enabled": True, + } + ], + "enabled": True, + "metadata": {"deployment_id": deployment_id}, + } + + +class ArenaOpenAPIClient: + """Small sync/async client for Stream discovery, datasets, and task launch.""" + + MAX_DATASET_LIMIT = 1000 + FAILED_TASK_STATUSES = { + "CANCELLED", + "COLLECT_FAILED", + "EVAL_FAILED", + "FAILED", + "HARNESS_FAILED", + "NO_OUTPUT", + "SETUP_FAILED", + "TIMEOUT", + } + + def __init__( + self, + base_url: str = "", + api_token: str = "", + llm_api_key: str = "", + timeout: float = 60.0, + poll_interval: float = 5.0, + request_retries: int = 3, + ) -> None: + self.base_url, self.api_token = resolve_arena_credentials( + base_url=base_url, + api_token=api_token, + ) + self.llm_api_key = llm_api_key or os.getenv("ARENA_LLM_API_KEY", "") + self.timeout = timeout + self.poll_interval = poll_interval + self.request_retries = request_retries + + @property + def _headers(self) -> dict[str, str]: + return {"Authorization": f"Bearer {self.api_token}"} + + @property + def llm_gateway_api_key(self) -> str: + """Return the Arena gateway key required by launched Harness tasks.""" + if not self.llm_api_key: + raise ValueError("ARENA_LLM_API_KEY is required for LLM gateway traffic") + return self.llm_api_key + + def list_streams( + self, + status: str | None = "ACTIVE", + *, + client: httpx.Client | None = None, + ) -> list[dict[str, Any]]: + """Return online Streams, optionally filtered by status.""" + params = {"status": status} if status else None + response = self._sync_request( + "GET", + f"{self.base_url}/openapi/v1/streams", + client=client, + params=params, + headers=self._headers, + ) + payload = _response_json(response) + items = payload.get("items") if isinstance(payload, Mapping) else None + if not isinstance(items, list): + raise ArenaAPIError("Stream list response is missing an 'items' array") + return [dict(item) for item in items if isinstance(item, Mapping)] + + async def list_streams_async( + self, + status: str | None = "ACTIVE", + *, + client: httpx.AsyncClient, + timeout: float, + ) -> list[dict[str, Any]]: + """Asynchronously return Streams, optionally filtered by status.""" + + params = {"status": status} if status else None + response = await self._async_request( + client, + "GET", + f"{self.base_url}/openapi/v1/streams", + params=params, + headers=self._headers, + timeout=timeout, + ) + payload = _response_json(response) + items = payload.get("items") if isinstance(payload, Mapping) else None + if not isinstance(items, list): + raise ArenaAPIError("Stream list response is missing an 'items' array") + return [dict(item) for item in items if isinstance(item, Mapping)] + + def resolve_stream( + self, + stream_id: str = "", + *, + client: httpx.Client | None = None, + ) -> dict[str, Any]: + """Resolve a Stream and retain metadata needed by the rollout Harness.""" + if stream_id: + encoded_stream_id = quote(stream_id, safe="") + response = self._sync_request( + "GET", + f"{self.base_url}/openapi/v1/streams/{encoded_stream_id}", + client=client, + headers=self._headers, + ) + stream = _stream_from_payload(_response_json(response)) + if stream.get("stream_id") != stream_id: + raise ArenaAPIError( + f"Arena Stream get returned id {stream.get('stream_id')!r}, " + f"expected {stream_id!r}" + ) + status = stream.get("status") + if status is not None and status != "ACTIVE": + raise ArenaAPIError( + f"Arena Stream {stream_id!r} is not ACTIVE; status={status!r}" + ) + return stream + + streams = self.list_streams(client=client) + if not streams: + raise ArenaAPIError("Arena OpenAPI returned no active Streams") + return streams[0] + + async def resolve_stream_async( + self, + stream_id: str = "", + *, + client: httpx.AsyncClient, + timeout: float, + ) -> dict[str, Any]: + """Asynchronously resolve one active Stream.""" + if stream_id: + encoded_stream_id = quote(stream_id, safe="") + response = await self._async_request( + client, + "GET", + f"{self.base_url}/openapi/v1/streams/{encoded_stream_id}", + headers=self._headers, + timeout=timeout, + ) + stream = _stream_from_payload(_response_json(response)) + if stream.get("stream_id") != stream_id: + raise ArenaAPIError( + f"Arena Stream get returned id {stream.get('stream_id')!r}, " + f"expected {stream_id!r}" + ) + status = stream.get("status") + if status is not None and status != "ACTIVE": + raise ArenaAPIError( + f"Arena Stream {stream_id!r} is not ACTIVE; status={status!r}" + ) + return stream + + streams = await self.list_streams_async(client=client, timeout=timeout) + if not streams: + raise ArenaAPIError("Arena OpenAPI returned no active Streams") + return streams[0] + + def resolve_stream_id( + self, + stream_id: str = "", + *, + client: httpx.Client | None = None, + ) -> str: + """Use an explicit Stream id or fall back to the first active Stream.""" + stream = self.resolve_stream(stream_id, client=client) + selected = stream.get("stream_id") + if not isinstance(selected, str) or not selected: + raise ArenaAPIError("The selected Stream is missing 'stream_id'") + return selected + + def _select_dataset_page( + self, + stream_id: str, + limit: int, + offset: int = 0, + *, + client: httpx.Client | None = None, + ) -> dict[str, Any]: + encoded_stream_id = quote(stream_id, safe="") + url = f"{self.base_url}/openapi/v1/streams/{encoded_stream_id}/dataset" + kwargs = { + "params": {"limit": limit, "offset": offset}, + "headers": self._headers, + } + response = self._sync_request("POST", url, client=client, **kwargs) + payload = _response_json(response) + if not isinstance(payload, Mapping): + raise ArenaAPIError("Dataset response must be a JSON object") + return dict(payload) + + def get_all_dataset_rows( + self, + stream_id: str, + llm_protocol: LLMProtocol = "chat_completions", + *, + client: httpx.Client | None = None, + ) -> list[dict[str, str]]: + """Load one Stream's complete dataset with stable, validated pagination.""" + first_page = self._select_dataset_page( + stream_id, + limit=1, + client=client, + ) + total = first_page.get("total") + if not isinstance(total, int) or isinstance(total, bool) or total < 0: + raise ArenaAPIError("Dataset response has an invalid 'total'") + if total == 0: + raise ArenaAPIError(f"Stream {stream_id!r} contains no dataset rows") + + if total == 1: + page_specs = iter([(0, 1, first_page)]) + else: + page_specs = ( + ( + offset, + min(self.MAX_DATASET_LIMIT, total - offset), + None, + ) + for offset in range(0, total, self.MAX_DATASET_LIMIT) + ) + + data_ids: list[str] = [] + item_by_data_id: dict[str, Mapping[str, Any]] = {} + for expected_offset, expected_limit, page in page_specs: + if page is None: + page = self._select_dataset_page( + stream_id, + limit=expected_limit, + offset=expected_offset, + client=client, + ) + page_total = page.get("total") + if page_total != total: + raise ArenaAPIError( + f"Stream {stream_id!r} dataset total changed during pagination: " + f"expected {total}, got {page_total!r}" + ) + page_offset = page.get("offset") + if page_offset != expected_offset: + raise ArenaAPIError( + f"Stream {stream_id!r} dataset pagination returned offset " + f"{page_offset!r}, expected {expected_offset}" + ) + page_limit = page.get("limit") + if page_limit != expected_limit: + raise ArenaAPIError( + f"Stream {stream_id!r} dataset pagination returned limit " + f"{page_limit!r}, expected {expected_limit}" + ) + page_data_ids = page.get("data_ids") + if not isinstance(page_data_ids, list) or not all( + isinstance(data_id, str) and data_id for data_id in page_data_ids + ): + raise ArenaAPIError("Dataset response has an invalid 'data_ids' array") + page_count = page.get("count") + if page_count != len(page_data_ids): + raise ArenaAPIError( + f"Stream {stream_id!r} dataset pagination returned count " + f"{page_count!r}, expected {len(page_data_ids)}" + ) + duplicate_ids = set(data_ids).intersection(page_data_ids) + if duplicate_ids or len(set(page_data_ids)) != len(page_data_ids): + raise ArenaAPIError( + f"Stream {stream_id!r} dataset pagination returned duplicate ids" + ) + data_ids.extend(page_data_ids) + + items = page.get("items") + if items is None: + continue + if not isinstance(items, list) or len(items) != len(page_data_ids): + raise ArenaAPIError("Dataset response has an invalid 'items' array") + page_item_ids: set[str] = set() + for item in items: + if not isinstance(item, Mapping): + raise ArenaAPIError( + "Dataset response contains an invalid dataset item" + ) + item_data_id = item.get("data_id") + if not isinstance(item_data_id, str) or not item_data_id: + raise ArenaAPIError("Dataset item is missing a valid 'data_id'") + if item_data_id in item_by_data_id or item_data_id in page_item_ids: + raise ArenaAPIError( + f"Stream {stream_id!r} dataset items contain duplicate ids" + ) + item_by_data_id[item_data_id] = item + page_item_ids.add(item_data_id) + if page_item_ids != set(page_data_ids): + raise ArenaAPIError( + "Dataset items do not match the returned 'data_ids'" + ) + + if len(data_ids) != total: + raise ArenaAPIError( + f"Dataset response returned {len(data_ids)} rows, expected {total}" + ) + + rows: list[dict[str, str]] = [] + for data_id in data_ids: + item = item_by_data_id.get(data_id) + rows.append( + { + "data_id": data_id, + "stream_id": stream_id, + "llm_protocol": llm_protocol, + "arena_task_type": _arena_task_type_from_tags( + item.get("tags") if item is not None else None + ), + } + ) + return rows + + def register_llm_proxy( + self, + model_name: str, + upstream_base_url: str, + upstream_api_key: str, + *, + deployment_id: str | None = None, + protocol: LLMProtocol = "chat_completions", + client: httpx.Client | None = None, + ) -> tuple[str, str]: + """Register an AReaL OpenAI proxy and return its URL and model alias.""" + if not model_name.startswith("stream-areal-"): + raise ValueError("Arena model_name must start with 'stream-areal-'") + if not upstream_base_url: + raise ValueError("upstream_base_url is required") + if not upstream_api_key: + raise ValueError("upstream_api_key is required") + + resolved_deployment_id = deployment_id or str(uuid.uuid4()) + payload = _llm_registration_payload( + model_name=model_name, + upstream_base_url=upstream_base_url, + upstream_api_key=upstream_api_key, + deployment_id=resolved_deployment_id, + protocol=protocol, + ) + response = self._sync_request( + "POST", + f"{self.base_url}/openapi/v1/llm/models", + client=client, + headers=self._headers, + json=payload, + ) + registration = _response_json(response) + registered_url, registered_model_id = self._registered_llm_target( + registration, + model_name, + ) + return registered_url, registered_model_id + + def _registered_llm_target( + self, + registration: Any, + model_name: str, + ) -> tuple[str, str]: + """Resolve the LLM gateway API base and model alias from registration.""" + if not isinstance(registration, Mapping): + raise ArenaAPIError("LLM registration response must be a JSON object") + registered_model_id = registration.get("model_name") + if not isinstance(registered_model_id, str) or not registered_model_id: + raise ArenaAPIError("LLM registration response is missing model_name") + if registered_model_id != model_name: + raise ArenaAPIError( + "LLM registration response returned an unexpected model_name: " + f"{registered_model_id!r}" + ) + return f"{self.base_url}/api", registered_model_id + + async def register_llm_proxy_async( + self, + model_name: str, + upstream_base_url: str, + upstream_api_key: str, + deployment_id: str, + *, + protocol: LLMProtocol = "chat_completions", + client: httpx.AsyncClient, + timeout: float = 180.0, + ) -> tuple[str, str]: + """Asynchronously register one rollout proxy session.""" + if not model_name.startswith("stream-areal-"): + raise ValueError("Arena model_name must start with 'stream-areal-'") + if not upstream_base_url: + raise ValueError("upstream_base_url is required") + if not upstream_api_key: + raise ValueError("upstream_api_key is required") + payload = _llm_registration_payload( + model_name=model_name, + upstream_base_url=upstream_base_url, + upstream_api_key=upstream_api_key, + deployment_id=deployment_id, + protocol=protocol, + ) + response = await self._async_request( + client, + "POST", + f"{self.base_url}/openapi/v1/llm/models", + headers=self._headers, + json=payload, + timeout=timeout, + ) + registration = _response_json(response) + return self._registered_llm_target(registration, model_name) + + async def llm_proxy_exists_async( + self, + model_name: str, + *, + client: httpx.AsyncClient, + timeout: float = 180.0, + ) -> bool: + """Return whether one dynamic LLM route still exists in Arena.""" + if not model_name.startswith("stream-areal-"): + raise ValueError("Arena model_name must start with 'stream-areal-'") + response = await self._async_request( + client, + "GET", + f"{self.base_url}/openapi/v1/llm/models/{quote(model_name, safe='')}", + headers=self._headers, + timeout=timeout, + ) + if response.status_code == 404: + return False + payload = _response_json(response) + if not isinstance(payload, Mapping): + raise ArenaAPIError("LLM model response must be a JSON object") + model = payload.get("data", payload) + if not isinstance(model, Mapping): + raise ArenaAPIError("LLM model response is missing the model object") + returned_name = model.get("model_name") + if returned_name != model_name: + raise ArenaAPIError( + "LLM model response returned an unexpected model_name: " + f"{returned_name!r}" + ) + return True + + def delete_llm_proxy( + self, + model_name: str, + *, + client: httpx.Client | None = None, + ) -> None: + """Delete one registered LLM model; missing models are clean.""" + response = self._sync_request( + "DELETE", + f"{self.base_url}/openapi/v1/llm/models/{quote(model_name, safe='')}", + client=client, + headers=self._headers, + ) + if response.status_code == 404: + return + if 200 <= response.status_code < 300: + return + _response_json(response) + + async def delete_llm_proxy_async( + self, + model_name: str, + *, + client: httpx.AsyncClient, + timeout: float = 180.0, + ) -> None: + """Asynchronously delete one registered LLM model.""" + response = await self._async_request( + client, + "DELETE", + f"{self.base_url}/openapi/v1/llm/models/{quote(model_name, safe='')}", + headers=self._headers, + timeout=timeout, + ) + if response.status_code == 404: + return + if 200 <= response.status_code < 300: + return + _response_json(response) + + def _sync_request( + self, + method: str, + url: str, + *, + client: httpx.Client | None = None, + **kwargs: Any, + ) -> httpx.Response: + request_client = ( + client if client is not None else httpx.Client(timeout=self.timeout) + ) + try: + for attempt in range(self.request_retries + 1): + try: + response = request_client.request(method, url, **kwargs) + except httpx.RequestError as exc: + if attempt == self.request_retries: + raise ArenaAPIError( + "Arena OpenAPI request failed after " + f"{attempt + 1} attempts: {type(exc).__name__}" + ) from exc + time.sleep(min(2**attempt, 10)) + continue + if not _is_retryable_response(response): + return response + if attempt == self.request_retries: + return response + response.close() + time.sleep(min(2**attempt, 10)) + finally: + if client is None: + request_client.close() + raise AssertionError("unreachable") + + async def _async_request( + self, + client: httpx.AsyncClient, + method: str, + url: str, + **kwargs: Any, + ) -> httpx.Response: + for attempt in range(self.request_retries + 1): + try: + response = await client.request(method, url, **kwargs) + except httpx.RequestError as exc: + if attempt == self.request_retries: + raise ArenaAPIError( + f"Arena OpenAPI request failed after {attempt + 1} attempts: " + f"{type(exc).__name__}" + ) from exc + await asyncio.sleep(min(2**attempt, 10)) + continue + if not _is_retryable_response(response): + return response + if attempt == self.request_retries: + return response + await response.aclose() + await asyncio.sleep(min(2**attempt, 10)) + raise AssertionError("unreachable") + + async def launch_one_task( + self, + stream_id: str, + data_id: str, + model_name: str, + proxy_base_url: str, + proxy_api_key: str, + *, + session_id: str = "", + harness: str = "", + task_envs: Mapping[str, str] | None = None, + client: httpx.AsyncClient | None = None, + on_launch_success: Callable[[str | None], None] | None = None, + ) -> float: + """Launch one Arena task and return its terminal score. + + This compatibility wrapper retains the original scalar API. New callers + that need ``raw`` and provenance metadata should use + :meth:`launch_one_task_result`. + """ + + result = await self.launch_one_task_result( + stream_id=stream_id, + data_id=data_id, + model_name=model_name, + proxy_base_url=proxy_base_url, + proxy_api_key=proxy_api_key, + session_id=session_id, + harness=harness, + task_envs=task_envs, + client=client, + on_launch_success=on_launch_success, + ) + if result.score is None: + raise ArenaAPIError("Successful Arena task result is missing score") + return result.score + + async def launch_one_task_result( + self, + stream_id: str, + data_id: str, + model_name: str, + proxy_base_url: str, + proxy_api_key: str, + *, + session_id: str = "", + harness: str = "", + task_envs: Mapping[str, str] | None = None, + client: httpx.AsyncClient | None = None, + on_launch_success: Callable[[str | None], None] | None = None, + ) -> ArenaTaskResult: + """Launch one Arena task and return its structured terminal result.""" + encoded_stream_id = quote(stream_id, safe="") + url = f"{self.base_url}/openapi/v1/streams/{encoded_stream_id}/launch_one_task" + envs = dict(task_envs or {}) + if not all( + isinstance(key, str) and key and isinstance(value, str) + for key, value in envs.items() + ): + raise ValueError("Arena task environment variables must be strings") + envs.update( + { + "MODEL_NAME": model_name, + "BASE_URL": proxy_base_url, + "API_KEY": proxy_api_key, + } + ) + request = { + "data_id": data_id, + "model_name": model_name, + "base_url": proxy_base_url, + "api_key": proxy_api_key, + "envs": envs, + } + if session_id: + request["session_id"] = session_id + if harness: + request["harness"] = harness + if client is not None: + return await self._launch_and_wait( + client, + url, + request, + on_launch_success=on_launch_success, + ) + async with httpx.AsyncClient(timeout=self.timeout) as owned_client: + return await self._launch_and_wait( + owned_client, + url, + request, + on_launch_success=on_launch_success, + ) + + async def _launch_and_wait( + self, + client: httpx.AsyncClient, + url: str, + request: dict[str, Any], + *, + on_launch_success: Callable[[str | None], None] | None = None, + ) -> ArenaTaskResult: + response = await self._async_request( + client, + "POST", + url, + json=request, + headers=self._headers, + ) + payload = _response_json(response) + result = _parse_task_result(payload, fallback_task_id="launch_one_task") + status = result.status + task_id = result.task_id + has_task_id = task_id != "launch_one_task" + if status in self.FAILED_TASK_STATUSES: + raise ArenaTaskFailedError( + result.task_id, + status, + payload if isinstance(payload, Mapping) else None, + result, + ) + if status and status not in {"DONE", "OK"}: + if has_task_id: + if on_launch_success is not None: + on_launch_success(task_id) + return await self._poll_task_result(client, task_id) + raise ArenaAPIError( + f"launch_one_task response has status {status} but no task_id to poll" + ) + + if result.score is None: + raise ArenaAPIError( + "launch_one_task terminal response is missing top-level 'score'" + ) + if on_launch_success is not None: + on_launch_success(task_id if has_task_id else None) + return result + + async def _poll_task_result( + self, + client: httpx.AsyncClient, + task_id: str, + ) -> ArenaTaskResult: + encoded_task_id = quote(task_id, safe="") + url = f"{self.base_url}/openapi/v1/tasks/{encoded_task_id}/result" + while True: + response = await self._async_request( + client, + "GET", + url, + headers=self._headers, + ) + payload = _response_json(response) + result = _parse_task_result(payload, fallback_task_id=task_id) + if result.task_id != task_id: + raise ArenaAPIError( + "Arena task result returned an unexpected task_id: " + f"{result.task_id!r}" + ) + if result.status in self.FAILED_TASK_STATUSES: + raise ArenaTaskFailedError( + task_id, + result.status, + payload if isinstance(payload, Mapping) else None, + result, + ) + if result.status in {"DONE", "OK"}: + if result.score is None: + raise ArenaAPIError( + "Arena terminal task result is missing top-level 'score'" + ) + return result + await asyncio.sleep(self.poll_interval) diff --git a/examples/swe/arena_config.py b/examples/swe/arena_config.py new file mode 100644 index 0000000000..7f48125fa1 --- /dev/null +++ b/examples/swe/arena_config.py @@ -0,0 +1,345 @@ +"""Configuration helpers for prompt-level Arena Stream mixtures.""" + +from __future__ import annotations + +import base64 +import binascii +import math +import re +from collections.abc import Mapping +from dataclasses import fields, replace +from pathlib import Path +from typing import Any + +import yaml + +from examples.swe.arena_types import ArenaRewardRefConfig, ArenaStreamConfig + +_STREAM_NAME_PATTERN = re.compile(r"^[A-Za-z0-9_.-]+$") + + +def _config_value(config: Any, name: str, default: Any = None) -> Any: + if isinstance(config, Mapping): + return config.get(name, default) + return getattr(config, name, default) + + +def _parse_reward_ref(value: Any, *, stream_name: str) -> ArenaRewardRefConfig: + if value in (None, ""): + return ArenaRewardRefConfig() + if isinstance(value, ArenaRewardRefConfig): + reward_ref = value + elif isinstance(value, Mapping): + unknown = set(value) - {"key", "version"} + if unknown: + raise ValueError( + f"Arena Stream {stream_name!r} expected_reward_ref has unknown " + f"fields: {sorted(unknown)}" + ) + reward_ref = ArenaRewardRefConfig( + key=str(value.get("key") or "").strip(), + version=str(value.get("version") or "").strip(), + ) + else: + raise ValueError( + f"Arena Stream {stream_name!r} expected_reward_ref must be an object" + ) + if bool(reward_ref.key) != bool(reward_ref.version): + raise ValueError( + f"Arena Stream {stream_name!r} expected_reward_ref requires both key " + "and version" + ) + return reward_ref + + +def parse_arena_stream_config(value: Any) -> ArenaStreamConfig: + """Validate one inline or file-backed Stream definition.""" + + if isinstance(value, ArenaStreamConfig): + raw = {field.name: getattr(value, field.name) for field in fields(value)} + elif isinstance(value, Mapping): + raw = dict(value) + else: + raise ValueError("Each Arena Stream definition must be an object") + + known_fields = {field.name for field in fields(ArenaStreamConfig)} + unknown = set(raw) - known_fields + if unknown: + raise ValueError( + f"Arena Stream definition has unknown fields: {sorted(unknown)}" + ) + + name = str(raw.get("name") or "").strip() + stream_id = str(raw.get("stream_id") or "").strip() + if not name: + raise ValueError("Arena Stream definition requires a non-empty name") + if not _STREAM_NAME_PATTERN.fullmatch(name): + raise ValueError( + f"Arena Stream name {name!r} may only contain letters, digits, '.', " + "'_', and '-'" + ) + if not stream_id: + raise ValueError(f"Arena Stream {name!r} requires a non-empty stream_id") + + weight = raw.get("sampling_weight", 1.0) + if isinstance(weight, bool) or not isinstance(weight, (int, float)): + raise ValueError(f"Arena Stream {name!r} sampling_weight must be numeric") + weight = float(weight) + if not math.isfinite(weight) or weight <= 0: + raise ValueError( + f"Arena Stream {name!r} sampling_weight must be finite and positive" + ) + + llm_protocol = str(raw.get("llm_protocol") or "").strip() + if llm_protocol not in ("", "anthropic", "responses", "chat_completions"): + raise ValueError( + f"Arena Stream {name!r} has unsupported llm_protocol {llm_protocol!r}" + ) + task_envs = raw.get("task_envs") or {} + if not isinstance(task_envs, Mapping) or not all( + isinstance(key, str) and key and isinstance(item, str) + for key, item in task_envs.items() + ): + raise ValueError(f"Arena Stream {name!r} task_envs must map strings to strings") + + threshold = raw.get("reward_threshold") + if threshold is not None: + if isinstance(threshold, bool) or not isinstance(threshold, (int, float)): + raise ValueError( + f"Arena Stream {name!r} reward_threshold must be numeric or null" + ) + threshold = float(threshold) + if not math.isfinite(threshold): + raise ValueError(f"Arena Stream {name!r} reward_threshold must be finite") + + return ArenaStreamConfig( + name=name, + stream_id=stream_id, + sampling_weight=weight, + harness=str(raw.get("harness") or "").strip(), + llm_protocol=llm_protocol, + task_envs=dict(task_envs), + expected_reward_ref=_parse_reward_ref( + raw.get("expected_reward_ref"), stream_name=name + ), + reward_threshold=threshold, + reward_transform_fn=str(raw.get("reward_transform_fn") or "").strip(), + ) + + +def load_arena_stream_configs(econfig: Any) -> list[ArenaStreamConfig]: + """Load multi-Stream definitions or translate the legacy single Stream.""" + + inline = list(_config_value(econfig, "arena_streams", []) or []) + streams_yaml_b64 = str( + _config_value(econfig, "arena_streams_yaml_b64", "") or "" + ).strip() + streams_file = str(_config_value(econfig, "arena_streams_file", "") or "").strip() + if sum((bool(inline), bool(streams_yaml_b64), bool(streams_file))) > 1: + raise ValueError( + "arena_streams, arena_streams_yaml_b64, and arena_streams_file are " + "mutually exclusive" + ) + + values: list[Any] + if streams_yaml_b64 or streams_file: + if streams_yaml_b64: + source = "Encoded inline Arena Streams YAML" + try: + streams_yaml = base64.b64decode(streams_yaml_b64, validate=True).decode( + "utf-8" + ) + except (binascii.Error, UnicodeDecodeError) as exc: + raise ValueError( + "arena_streams_yaml_b64 is not valid base64 UTF-8" + ) from exc + else: + source = "Arena Streams file" + path = Path(streams_file).expanduser() + if not path.is_file(): + raise FileNotFoundError(f"Arena Streams file not found: {path}") + streams_yaml = path.read_text(encoding="utf-8") + payload = yaml.safe_load(streams_yaml) + if isinstance(payload, Mapping): + unknown = set(payload) - {"streams"} + if unknown: + raise ValueError( + f"{source} has unknown top-level fields: {sorted(unknown)}" + ) + values = payload.get("streams") + else: + values = payload + if not isinstance(values, list) or not values: + raise ValueError(f"{source} must contain a non-empty 'streams' list") + elif inline: + values = inline + else: + stream_id = str(_config_value(econfig, "stream_id", "") or "").strip() + legacy_task_envs = _config_value(econfig, "arena_task_envs", {}) or {} + if not isinstance(legacy_task_envs, Mapping) or not all( + isinstance(key, str) and key and isinstance(item, str) + for key, item in legacy_task_envs.items() + ): + raise ValueError("arena_task_envs must map strings to strings") + legacy_threshold = _config_value(econfig, "arena_reward_threshold", None) + if legacy_threshold is not None: + try: + legacy_threshold = float(legacy_threshold) + except (TypeError, ValueError) as exc: + raise ValueError( + "arena_reward_threshold must be numeric or null" + ) from exc + if not math.isfinite(legacy_threshold): + raise ValueError("arena_reward_threshold must be finite") + return [ + ArenaStreamConfig( + name="default", + stream_id=stream_id, + harness=str(_config_value(econfig, "arena_harness", "") or "").strip(), + llm_protocol=str( + _config_value(econfig, "arena_llm_protocol", "") or "" + ).strip(), + task_envs=dict(legacy_task_envs), + reward_threshold=legacy_threshold, + reward_transform_fn=str( + _config_value(econfig, "arena_reward_transform_fn", "") or "" + ).strip(), + ) + ] + + streams = [parse_arena_stream_config(value) for value in values] + common_task_envs_value = _config_value(econfig, "arena_task_envs", {}) or {} + if not isinstance(common_task_envs_value, Mapping) or not all( + isinstance(key, str) and key and isinstance(item, str) + for key, item in common_task_envs_value.items() + ): + raise ValueError("arena_task_envs must map strings to strings") + common_task_envs = dict(common_task_envs_value) + common_harness = str(_config_value(econfig, "arena_harness", "") or "").strip() + streams = [ + replace( + stream, + harness=stream.harness or common_harness, + task_envs={**common_task_envs, **stream.task_envs}, + ) + for stream in streams + ] + names = [stream.name for stream in streams] + if len(names) != len(set(names)): + raise ValueError("Arena Stream names must be unique") + stream_ids = [stream.stream_id for stream in streams] + if len(stream_ids) != len(set(stream_ids)): + raise ValueError("Arena Stream ids must be unique") + return streams + + +def build_weighted_arena_rows( + rows_by_stream: Mapping[str, list[dict[str, str]]], + streams: list[ArenaStreamConfig], + epoch_size: int = 0, + size_multiple: int = 1, +) -> list[dict[str, str]]: + """Build a deterministic Stream-weighted epoch covering every source row.""" + + if epoch_size < 0: + raise ValueError("arena_mixture_epoch_size must be non-negative") + if ( + isinstance(size_multiple, bool) + or not isinstance(size_multiple, int) + or size_multiple < 1 + ): + raise ValueError("Arena mixture size_multiple must be a positive integer") + for stream in streams: + if not rows_by_stream.get(stream.name): + raise ValueError(f"Arena Stream {stream.name!r} contains no dataset rows") + + # ``sampling_weight`` weights each source dataset rather than replacing its + # cardinality. The default epoch is the raw union, so every source appears + # exactly once regardless of weight. An explicit smaller epoch uses the + # weights to choose a deterministic subset without replacement. + capacities = [len(rows_by_stream[stream.name]) for stream in streams] + stream_masses = [ + capacity * stream.sampling_weight + for capacity, stream in zip(capacities, streams, strict=True) + ] + total_mass = sum(stream_masses) + if not all(math.isfinite(mass) and mass > 0 for mass in stream_masses) or not ( + math.isfinite(total_mass) and total_mass > 0 + ): + raise ValueError("Arena Stream weights overflowed the mixture mass") + total_source_rows = sum(capacities) + if epoch_size: + resolved_size = epoch_size + if resolved_size > total_source_rows: + raise ValueError( + f"arena_mixture_epoch_size={resolved_size} cannot exceed the " + f"{total_source_rows} unique source rows" + ) + if resolved_size % size_multiple: + raise ValueError( + f"arena_mixture_epoch_size={resolved_size} must be divisible by " + f"the training batch size {size_multiple}" + ) + else: + # Do not pad the raw union to a batch multiple: padding would repeat rows. + # The training sampler may drop a final incomplete batch instead. + resolved_size = total_source_rows + if resolved_size <= 0: + raise ValueError("Arena mixture epoch must contain at least one prompt row") + + if resolved_size == total_source_rows: + counts = capacities + else: + counts = [0] * len(streams) + remaining = resolved_size + active = set(range(len(streams))) + while remaining and active: + active_mass = sum(stream_masses[index] for index in active) + capped = [ + index + for index in active + if remaining * stream_masses[index] / active_mass >= capacities[index] + ] + if capped: + for index in capped: + counts[index] = capacities[index] + remaining -= counts[index] + active.remove(index) + continue + + exact = { + index: remaining * stream_masses[index] / active_mass + for index in active + } + for index, value in exact.items(): + counts[index] = int(value) + remainder = resolved_size - sum(counts) + order = sorted( + active, + key=lambda index: ( + -(exact[index] - counts[index]), + streams[index].name, + ), + ) + for index in order[:remainder]: + counts[index] += 1 + remaining = 0 + + mixed_rows: list[dict[str, str]] = [] + emitted = [0] * len(streams) + while len(mixed_rows) < resolved_size: + candidates = [ + index for index, count in enumerate(counts) if emitted[index] < count + ] + selected = min( + candidates, + key=lambda index: ( + (emitted[index] + 1) / streams[index].sampling_weight, + streams[index].name, + ), + ) + stream = streams[selected] + source_rows = rows_by_stream[stream.name] + mixed_rows.append(dict(source_rows[emitted[selected]])) + emitted[selected] += 1 + return mixed_rows diff --git a/examples/swe/arena_flash_v3.yaml b/examples/swe/arena_flash_v3.yaml new file mode 100644 index 0000000000..ef68f713a7 --- /dev/null +++ b/examples/swe/arena_flash_v3.yaml @@ -0,0 +1,202 @@ +# Flash V3 multi-stream resource/model settings from swe-dev@5478d6384. +# Uses main GRPO group handling; see README_arena.md for recipe differences. +defaults: +- arena_grpo +- _self_ +experiment_name: ${oc.env:EXPERIMENT_NAME,"arena-flash-v3"} +trial_name: ${oc.env:TRIAL_NAME,"flash-v3-128k"} +total_train_epochs: ${oc.env:TOTAL_TRAIN_EPOCHS,1000} +total_train_steps: ${oc.env:TOTAL_TRAIN_STEPS,500} +dynamic_bs: false +gconfig: + min_new_tokens: 0 + greedy: false + top_p: 1.0 + top_k: 1000000 + temperature: 1.0 + n_samples: ${oc.env:N_SAMPLES,12} + max_new_tokens: ${oc.env:MAX_NEW_TOKENS,131071} + max_tokens: ${oc.env:MAX_TOKENS,131071} + reward_normalization: true + drop_incomplete_group: true +eval_gconfig: + n_samples: ${gconfig.n_samples} + max_new_tokens: ${gconfig.max_new_tokens} + max_tokens: ${gconfig.max_tokens} + reward_normalization: false + drop_incomplete_group: ${gconfig.drop_incomplete_group} +rollout: + backend: ${oc.env:ROLLOUT_BACKEND,"sglang:d16t4p1"} + max_concurrent_rollouts: ${oc.env:MAX_CONCURRENT_ROLLOUTS,1500} + consumer_batch_size: ${train_dataset.batch_size} + scheduling_spec: + - task_type: ${actor.scheduling_spec.0.task_type} + port_count: ${actor.scheduling_spec.0.port_count} + gpu: ${actor.scheduling_spec.0.gpu} + cpu: ${actor.scheduling_spec.0.cpu} + mem: ${actor.scheduling_spec.0.mem} + image: ${oc.env:ROLLOUT_IMAGE} + cmd: ${actor.scheduling_spec.0.cmd} + env_vars: + PYTHONPATH: ${oc.env:AREAL_DIR}:${oc.env:ROLLOUT_RUNTIME_PYTHONPATH,""}:${oc.env:FLASH_LINEAR_ATTENTION_ROOT}:${oc.env:AWEX_ROOT} + AWEX_ROOT: ${oc.env:AWEX_ROOT} + AWEX_COLOCATE_TIMEOUT_S: ${oc.env:AWEX_COLOCATE_TIMEOUT_S,"7200"} + AREAL_AWEX_PROCESS_QUEUE_WHEN_IDLE: '1' + CUDA_DEVICE_MAX_CONNECTIONS: '1' + PYTORCH_CUDA_ALLOC_CONF: '' + additional_bash_cmds: ${actor.scheduling_spec.0.additional_bash_cmds} + scheduling_strategy: + type: colocation + target: actor + fork: true + agent: + tool_call_parser: ling3 + reasoning_parser: ling3 + engine_max_tokens: ${gconfig.max_tokens} + message_preprocessors: + - examples.swe.preprocessors.StripAnthropicBillingHeader + - examples.swe.preprocessors.StripAllSystemReminders + - examples.swe.preprocessors.StripAnthropicCacheFields + - examples.swe.preprocessors.NormalizeToolCallArguments + prefix_matcher: examples.swe.prefix_matchers.swe_prefix_matcher + drop_retry_orphans: true + session_timeout_seconds: ${oc.env:ARENA_PROXY_IDLE_TIMEOUT,7200} + max_head_offpolicyness: 1 + setup_timeout: 3600 + request_timeout: ${oc.env:ARENA_TASK_TIMEOUT,20000} +actor: + gradient_checkpointing: true + dtype: bfloat16 + grad_reduce_dtype: float32 + weight_update_mode: awex + eps_clip: 0.2 + eps_clip_higher: 0.28 + kl_ctl: 0.0 + ppo_n_minibatches: 1 + use_decoupled_loss: true + rejection_sampling: + level: token + action: mask + metric: ratio + upper: 5.0 + lower: 0.5 + discount: 1.0 + gae_timestep_unit: turn + gae_lambda: areal.trainer.ppo.lambda_fn.vapo_length_adaptive_gae + gae_lambda_kwargs: + alpha: 1.5 + reward_scaling: 1.0 + reward_bias: 0.0 + reward_clip: 5.0 + reward_norm: null + adv_norm: + mean_level: null + std_level: null + backend: ${oc.env:ACTOR_BACKEND,"megatron:(attn:d2p2t2c8|ffn:d4p2e8)"} + path: ${oc.env:MODEL_PATH} + disable_dropout: true + mb_spec: + n_mbs: ${oc.env:ACTOR_N_MBS,1} + granularity: 1 + max_tokens_per_mb: ${oc.env:ACTOR_MAX_TOKENS_PER_MB,131072} + n_mbs_divisor: ${oc.env:ACTOR_N_MBS_DIVISOR,1} + optimizer: + type: adam + lr: ${oc.env:ACTOR_LR,3.0e-6} + weight_decay: 0.01 + beta1: 0.9 + beta2: 0.999 + eps: 1.0e-08 + lr_scheduler_type: constant + gradient_clipping: 1.0 + warmup_steps: 5 + megatron: + wrap_with_ddp: true + ddp: + grad_reduce_in_fp32: true + overlap_grad_reduce: true + overlap_param_gather: false + use_distributed_optimizer: true + use_deterministic_algorithms: false + enable_fp32_lm_head: true + cross_entropy_loss_fusion: false + recompute_granularity: full + recompute_method: uniform + recompute_num_layers: 1 + moe_router_dtype: fp32 + moe_token_dispatcher_type: alltoall + moe_router_bias_update_rate: 0.0 + bridge_type: mbridge + virtual_pipeline_parallel_size: 1 + enable_mtp: false + recompute_logprob: false + prox_logp_method: reuse_train_logp + scheduling_spec: + - task_type: worker + port_count: 8 + gpu: 1 + cpu: 4 + mem: 64 + image: ${oc.env:AREAL_IMAGE} + cmd: ${oc.env:AREAL_PYTHON} -m areal.infra.rpc.rpc_server + env_vars: + PYTHONPATH: ${oc.env:AREAL_DIR}:${oc.env:ACTOR_RUNTIME_PYTHONPATH,""}:${oc.env:FLASH_LINEAR_ATTENTION_ROOT}:${oc.env:AWEX_ROOT} + AWEX_ROOT: ${oc.env:AWEX_ROOT} + AWEX_COLOCATE_TIMEOUT_S: ${oc.env:AWEX_COLOCATE_TIMEOUT_S,"7200"} + AREAL_AWEX_PROCESS_QUEUE_WHEN_IDLE: '1' + CUDA_DEVICE_MAX_CONNECTIONS: '1' + PYTORCH_CUDA_ALLOC_CONF: '' + additional_bash_cmds: + - set -a; source "$ARENA_CREDENTIALS_FILE"; set +a + - unset ALL_PROXY all_proxy HTTP_PROXY HTTPS_PROXY http_proxy https_proxy +ref: + disable_dropout: true + mb_spec: + n_mbs: ${actor.mb_spec.n_mbs} + granularity: ${actor.mb_spec.granularity} + max_tokens_per_mb: ${actor.mb_spec.max_tokens_per_mb} + n_mbs_divisor: ${actor.mb_spec.n_mbs_divisor} +econfig: + stream_id: '' + arena_streams_yaml_b64: '' + arena_streams_file: ${oc.env:DATASET_CONFIG} + arena_harness: '' + arena_llm_protocol: '' + arena_reward_threshold: null + arena_reward_transform_fn: '' + arena_llm_route_mode: ${oc.env:ARENA_LLM_ROUTE_MODE,"session_gateway"} + arena_streams: [] + timeout: ${oc.env:ARENA_TASK_WAIT_TIMEOUT,7200} + arena_mixture_epoch_size: ${oc.env:ARENA_MIXTURE_EPOCH_SIZE,0} + arena_task_envs: + ANTHROPIC_CUSTOM_HEADERS: 'x-keepalive-enable: true' + CLAUDE_ENABLE_STREAM_WATCHDOG: '0' + CLAUDE_CODE_DISABLE_EXPERIMENTAL_BETAS: '1' +sglang: + context_length: ${oc.env:CONTEXT_LENGTH,131072} + attention_backend: ${oc.env:SGLANG_ATTENTION_BACKEND,fa3} + max_prefill_tokens: ${oc.env:SGLANG_MAX_PREFILL_TOKENS,8192} + chunked_prefill_size: ${oc.env:SGLANG_CHUNKED_PREFILL_SIZE,8192} + max_running_requests: ${oc.env:SGLANG_MAX_RUNNING_REQUESTS,128} + mem_fraction_static: ${oc.env:SGLANG_MEM_FRACTION_STATIC,0.65} + schedule_policy: fcfs + disable_cuda_graph: ${oc.env:SGLANG_DISABLE_CUDA_GRAPH,false} + sampling_backend: pytorch + enable_memory_saver: true + disable_radix_cache: false + triton_attention_num_kv_splits: 16 + disable_custom_all_reduce: true +train_dataset: + batch_size: ${oc.env:TRAIN_BATCH_SIZE,32} + shuffle: ${oc.env:ARENA_DATASET_SHUFFLE,true} + pin_memory: true + num_workers: 0 + path: arena + type: rl + max_length: ${sglang.context_length} + drop_last: true +valid_dataset: null +seed: 42 +enable_offload: false +cluster: + n_nodes: ${oc.env:N_NODES,8} diff --git a/examples/swe/arena_grpo.yaml b/examples/swe/arena_grpo.yaml new file mode 100644 index 0000000000..57d1a8e741 --- /dev/null +++ b/examples/swe/arena_grpo.yaml @@ -0,0 +1,194 @@ +experiment_name: arena-grpo +trial_name: ${oc.env:TRIAL_NAME,smoke} +seed: 1 +enable_offload: false +total_train_epochs: 1 +tokenizer_path: ${actor.path} +cluster: + n_nodes: 2 + n_gpus_per_node: 8 + fileroot: ${oc.env:AREAL_FILEROOT} + name_resolve: + type: nfs + nfs_record_root: ${cluster.fileroot}/name_resolve +scheduler: + type: slurm +rollout: + backend: sglang:d4p1t2 + experiment_name: ${experiment_name} + trial_name: ${trial_name} + max_concurrent_rollouts: 4 + queue_size: null + consumer_batch_size: ${train_dataset.batch_size} + max_head_offpolicyness: 0 + enable_rollout_tracing: true + scheduling_spec: ${actor.scheduling_spec} + fileroot: ${cluster.fileroot} + tokenizer_path: ${tokenizer_path} + dump_to_file: true + agent: + mode: inline + export_style: concat + turn_discount: 1.0 + admin_api_key: ${oc.env:SWE_RL_ADMIN_API_KEY} + tool_call_parser: qwen25 + reasoning_parser: qwen3 + chat_template_type: concat + engine_max_tokens: ${gconfig.max_tokens} + session_timeout_seconds: 3600 + setup_timeout: 900 + request_timeout: 7200 +gconfig: + n_samples: 2 + min_new_tokens: 0 + max_new_tokens: 8192 + max_tokens: 32767 + greedy: false + temperature: 1.0 +actor: + backend: fsdp:d8p1t1 + experiment_name: ${experiment_name} + trial_name: ${trial_name} + path: ${oc.env:MODEL_PATH} + init_from_scratch: false + disable_dropout: true + gradient_checkpointing: true + dtype: bfloat16 + mb_spec: + max_tokens_per_mb: 32768 + packing_algorithm: ffd + optimizer: + type: adam + lr: 6.0e-06 + weight_decay: 0.017 + 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.4 + temperature: ${gconfig.temperature} + reward_scaling: 1.0 + reward_bias: 0.0 + kl_ctl: 0.0 + ppo_n_minibatches: 1 + recompute_logprob: true + use_decoupled_loss: true + rejection_sampling: + metric: ratio + upper: 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 + gpu: 1 + cpu: 8 + mem: 32 + port_count: 8 + image: ${oc.env:AREAL_IMAGE} + cmd: ${oc.env:AREAL_PYTHON} -m areal.infra.rpc.rpc_server + env_vars: + PYTHONPATH: ${oc.env:AREAL_DIR} + additional_bash_cmds: + - set -a; source "$ARENA_CREDENTIALS_FILE"; set +a + - unset ALL_PROXY all_proxy HTTP_PROXY HTTPS_PROXY http_proxy https_proxy +ref: + backend: ${actor.backend} + 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 + packing_algorithm: ffd + optimizer: null + scheduling_strategy: + type: colocation + target: actor + scheduling_spec: ${actor.scheduling_spec} +sglang: + model_path: ${actor.path} + random_seed: ${seed} + skip_tokenizer_init: true + dtype: ${actor.dtype} + max_running_requests: null + # SGLang reserves one token when checking prompt + completion length. + context_length: 32768 + mem_fraction_static: 0.7 +vllm: + model: ${actor.path} + seed: ${seed} + skip_tokenizer_init: false + dtype: ${actor.dtype} + max_model_len: 32768 + gpu_memory_utilization: 0.8 +train_dataset: + batch_size: 8 + shuffle: false + pin_memory: true + num_workers: 0 + path: arena + type: rl + max_length: 1024 +valid_dataset: + batch_size: 8 + pin_memory: true + num_workers: 0 + path: arena + type: rl + shuffle: false +saver: + experiment_name: ${experiment_name} + trial_name: ${trial_name} + fileroot: ${cluster.fileroot} + freq_epochs: null + freq_steps: null + freq_secs: null +recover: + mode: disabled + experiment_name: ${experiment_name} + trial_name: ${trial_name} + fileroot: ${cluster.fileroot} + freq_epochs: null + freq_steps: null + freq_secs: null +evaluator: + experiment_name: ${experiment_name} + trial_name: ${trial_name} + fileroot: ${cluster.fileroot} + freq_epochs: null + 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 +total_train_steps: 1 +dynamic_bs: false +should_accept_fn: null +econfig: + dataset_source: arena + stream_id: ${oc.env:ARENA_STREAM_ID,""} + arena_base_url: ${oc.env:ARENA_OPENAPI_BASE} + arena_llm_route_mode: session_gateway + arena_result_dump_dir: ${cluster.fileroot}/${experiment_name}/${trial_name}/arena_results + timeout: 1800.0 diff --git a/examples/swe/arena_multi_stream.yaml b/examples/swe/arena_multi_stream.yaml new file mode 100644 index 0000000000..a0eac2920e --- /dev/null +++ b/examples/swe/arena_multi_stream.yaml @@ -0,0 +1,20 @@ +defaults: +- arena_grpo +- _self_ +econfig: + stream_id: '' + arena_streams: + - name: first + stream_id: your-first-stream + sampling_weight: 1 + harness: your-harness@version + expected_reward_ref: + key: your-reward + version: 1.0.0 + - name: second + stream_id: your-second-stream + sampling_weight: 1 + harness: your-harness@version + expected_reward_ref: + key: your-reward + version: 1.0.0 diff --git a/examples/swe/arena_preflight.py b/examples/swe/arena_preflight.py new file mode 100644 index 0000000000..65c0e4b4f9 --- /dev/null +++ b/examples/swe/arena_preflight.py @@ -0,0 +1,447 @@ +"""Validate a multi-Stream Arena configuration with the Arena OpenAPI.""" + +from __future__ import annotations + +import argparse +import json +import os +import sys +import time +from collections.abc import Mapping +from dataclasses import asdict, replace +from typing import Any +from urllib.parse import quote + +import httpx + +from examples.swe.arena_client import ( + ArenaAPIError, + ArenaOpenAPIClient, + infer_llm_protocol_from_harness, +) +from examples.swe.arena_config import load_arena_stream_configs +from examples.swe.arena_types import ArenaRewardRefConfig + + +def _arena_api_base_url(base_url: str) -> str: + base_url = base_url.rstrip("/") + if base_url.endswith("/openapi/v1"): + return base_url + return f"{base_url}/openapi/v1" + + +def _arena_json( + base_url: str, + args: list[str], + *, + retries: int, + client: httpx.Client | None = None, +) -> Mapping[str, Any]: + token = os.getenv("ARENA_OPENAPI_TOKEN", "") + if not token: + raise ArenaAPIError("ARENA_OPENAPI_TOKEN is required for Arena preflight") + + command = tuple(args[:2]) + options = dict(zip(args[3::2], args[4::2])) + api_base = _arena_api_base_url(base_url) + if command == ("stream", "get") and len(args) == 3: + method = "GET" + url = f"{api_base}/streams/{quote(args[2], safe='')}" + params = None + elif command == ("harness", "version-get") and len(args) == 4: + method = "GET" + url = ( + f"{api_base}/harnesses/{quote(args[2], safe='')}/versions/" + f"{quote(args[3], safe='')}" + ) + params = None + elif command == ("stream", "dataset") and len(args) >= 3: + method = "POST" + url = f"{api_base}/streams/{quote(args[2], safe='')}/dataset" + params = { + "limit": options.get("--limit", "1"), + "offset": options.get("--offset", "0"), + } + elif command == ("stream", "tasks") and len(args) >= 3: + method = "GET" + url = f"{api_base}/streams/{quote(args[2], safe='')}/tasks" + params = { + "limit": options.get("--limit", "100"), + "offset": options.get("--offset", "0"), + } + else: + raise ValueError(f"Unsupported Arena preflight operation: {args}") + + headers = {"Authorization": f"Bearer {token}"} + owns_client = client is None + if client is None: + timeout = httpx.Timeout( + float(os.getenv("ARENA_PREFLIGHT_HTTP_TIMEOUT", "180")), + connect=float(os.getenv("ARENA_PREFLIGHT_CONNECT_TIMEOUT", "10")), + ) + client = httpx.Client(timeout=timeout, trust_env=False) + try: + for attempt in range(retries + 1): + try: + response = client.request(method, url, params=params, headers=headers) + except httpx.RequestError as exc: + if attempt == retries: + raise ArenaAPIError( + "Arena preflight request failed after " + f"{attempt + 1} attempts: {type(exc).__name__}" + ) from exc + time.sleep(min(2**attempt, 10)) + continue + retryable = ( + response.status_code == 429 + or response.status_code >= 500 + or ( + response.status_code == 403 + and "spanner-http-ant-group-watch-all" in response.text[:1000] + ) + ) + if retryable and attempt < retries: + response.close() + time.sleep(min(2**attempt, 10)) + continue + try: + response.raise_for_status() + except httpx.HTTPStatusError as exc: + raise ArenaAPIError( + f"Arena OpenAPI returned HTTP {response.status_code}: " + f"{response.text[:500]}" + ) from exc + try: + payload = response.json() + except ValueError as exc: + raise ArenaAPIError( + "Arena OpenAPI returned a non-JSON response" + ) from exc + if not isinstance(payload, Mapping): + raise ArenaAPIError("Arena OpenAPI response must be a JSON object") + return payload + finally: + if owns_client: + client.close() + raise AssertionError("unreachable") + + +def _payload_data(payload: Mapping[str, Any]) -> Any: + return payload.get("data", payload) + + +def _reward_ref(value: Any, *, stream_name: str) -> ArenaRewardRefConfig: + if value is None: + return ArenaRewardRefConfig() + if not isinstance(value, Mapping): + raise ArenaAPIError( + f"Arena Stream {stream_name!r} default_reward_ref must be an object" + ) + key = value.get("key") + version = value.get("version") + if ( + not isinstance(key, str) + or not key + or not isinstance(version, str) + or not version + ): + raise ArenaAPIError( + f"Arena Stream {stream_name!r} default_reward_ref is incomplete" + ) + return ArenaRewardRefConfig(key=key, version=version) + + +def _harness_ref(value: Any) -> str: + if not isinstance(value, Mapping): + return "" + key = value.get("key") + version = value.get("version") + if ( + not isinstance(key, str) + or not key + or not isinstance(version, str) + or not version + ): + return "" + return f"{key}@{version}" + + +def _stream_object(payload: Mapping[str, Any]) -> Mapping[str, Any]: + data = _payload_data(payload) + if isinstance(data, Mapping) and isinstance(data.get("stream"), Mapping): + return data["stream"] + if isinstance(data, Mapping): + return data + raise ArenaAPIError("Arena stream get response is missing the Stream object") + + +def _dataset_total(payload: Mapping[str, Any]) -> int: + data = _payload_data(payload) + total = data.get("total") if isinstance(data, Mapping) else None + if not isinstance(total, int) or isinstance(total, bool) or total < 1: + raise ArenaAPIError("Arena Stream dataset must contain at least one row") + return total + + +def _task_rows(payload: Mapping[str, Any]) -> list[Mapping[str, Any]]: + data = _payload_data(payload) + if isinstance(data, Mapping): + data = data.get("items", []) + if not isinstance(data, list): + return [] + return [item for item in data if isinstance(item, Mapping)] + + +def _recent_health( + rows: list[Mapping[str, Any]], + *, + harness: str, + reward_ref: ArenaRewardRefConfig, +) -> tuple[int, int]: + harness_key, _, harness_version = harness.partition("@") + terminal = 0 + failures = 0 + for item in rows: + task = item.get("task", item) + if not isinstance(task, Mapping): + continue + task_harness = task.get("harness_ref") + if not isinstance(task_harness, Mapping) or ( + task_harness.get("key"), + task_harness.get("version"), + ) != (harness_key, harness_version): + continue + task_reward = task.get("reward_ref") + if reward_ref.key and ( + not isinstance(task_reward, Mapping) + or ( + task_reward.get("key"), + task_reward.get("version"), + ) + != (reward_ref.key, reward_ref.version) + ): + continue + status = str(task.get("status") or "").upper() + if status in {"DONE", "OK"}: + terminal += 1 + elif status in ArenaOpenAPIClient.FAILED_TASK_STATUSES: + terminal += 1 + failures += 1 + return terminal, failures + + +def validate_streams( + streams_file: str = "", + *, + streams_yaml_b64: str = "", + base_url: str, + default_harness: str = "", + retries: int = 3, + recent_tasks: int = 100, + min_terminal_tasks: int = 1, + max_failure_rate: float = 0.8, + allow_unhealthy: bool = False, +) -> list[dict[str, Any]]: + """Resolve and validate every configured Stream without launching a task.""" + + timeout = httpx.Timeout( + float(os.getenv("ARENA_PREFLIGHT_HTTP_TIMEOUT", "180")), + connect=float(os.getenv("ARENA_PREFLIGHT_CONNECT_TIMEOUT", "10")), + ) + with httpx.Client(timeout=timeout, trust_env=False) as client: + return _validate_streams_with_client( + streams_file, + streams_yaml_b64=streams_yaml_b64, + base_url=base_url, + default_harness=default_harness, + retries=retries, + recent_tasks=recent_tasks, + min_terminal_tasks=min_terminal_tasks, + max_failure_rate=max_failure_rate, + allow_unhealthy=allow_unhealthy, + client=client, + ) + + +def _validate_streams_with_client( + streams_file: str = "", + *, + streams_yaml_b64: str = "", + base_url: str, + default_harness: str = "", + retries: int = 3, + recent_tasks: int = 100, + min_terminal_tasks: int = 1, + max_failure_rate: float = 0.8, + allow_unhealthy: bool = False, + client: httpx.Client, +) -> list[dict[str, Any]]: + if retries < 0: + raise ValueError("retries must be non-negative") + if recent_tasks < 1: + raise ValueError("recent_tasks must be positive") + if min_terminal_tasks < 0: + raise ValueError("min_terminal_tasks must be non-negative") + if not 0.0 <= max_failure_rate <= 1.0: + raise ValueError("max_failure_rate must be between 0 and 1") + + streams = load_arena_stream_configs( + { + "arena_streams_file": streams_file, + "arena_streams_yaml_b64": streams_yaml_b64, + "arena_harness": default_harness, + } + ) + summaries: list[dict[str, Any]] = [] + for configured in streams: + stream_payload = _arena_json( + base_url, + ["stream", "get", configured.stream_id], + retries=retries, + client=client, + ) + stream = _stream_object(stream_payload) + actual_stream_id = stream.get("stream_id") + if actual_stream_id != configured.stream_id: + raise ArenaAPIError( + f"Arena Stream lookup returned {actual_stream_id!r}, expected " + f"{configured.stream_id!r}" + ) + if stream.get("status") != "ACTIVE": + raise ArenaAPIError(f"Arena Stream {configured.stream_id!r} is not ACTIVE") + + actual_reward_ref = _reward_ref( + stream.get("default_reward_ref"), stream_name=configured.name + ) + expected = configured.expected_reward_ref + if expected.key and expected != actual_reward_ref: + raise ArenaAPIError( + f"Arena Stream {configured.name!r} reward_ref drifted: expected " + f"{expected.key}@{expected.version}, got " + f"{actual_reward_ref.key}@{actual_reward_ref.version}" + ) + + harness = configured.harness or _harness_ref(stream.get("default_harness_ref")) + harness_key, separator, harness_version = harness.partition("@") + if not separator or not harness_key or not harness_version: + raise ArenaAPIError( + f"Arena Stream {configured.name!r} requires Harness key@version" + ) + harness_payload = _arena_json( + base_url, + ["harness", "version-get", harness_key, harness_version], + retries=retries, + client=client, + ) + harness_data = _payload_data(harness_payload) + harness_status = ( + harness_data.get("status") if isinstance(harness_data, Mapping) else None + ) + if harness_status != "PUBLISHED": + raise ArenaAPIError(f"Arena Harness {harness!r} is not PUBLISHED") + + dataset_payload = _arena_json( + base_url, + ["stream", "dataset", configured.stream_id, "--limit", "1"], + retries=retries, + client=client, + ) + dataset_total = _dataset_total(dataset_payload) + # The task-list endpoint may return either a paginated object or a bare + # list under ``data``. Request the full health window directly; reusing a + # one-row total probe would silently evaluate only one task for the latter. + tasks_payload = _arena_json( + base_url, + [ + "stream", + "tasks", + configured.stream_id, + "--limit", + str(recent_tasks), + "--offset", + "0", + ], + retries=retries, + client=client, + ) + terminal_count, failure_count = _recent_health( + _task_rows(tasks_payload), + harness=harness, + reward_ref=actual_reward_ref, + ) + failure_rate = failure_count / terminal_count if terminal_count else 0.0 + unhealthy = ( + terminal_count < min_terminal_tasks or failure_rate >= max_failure_rate + ) + if unhealthy and not allow_unhealthy: + raise ArenaAPIError( + f"Arena Stream {configured.name!r} failed health check: " + f"terminal={terminal_count}, failures={failure_count}, " + f"failure_rate={failure_rate:.4f}" + ) + + protocol = configured.llm_protocol or infer_llm_protocol_from_harness(harness) + resolved = replace( + configured, + harness=harness, + llm_protocol=protocol, + expected_reward_ref=actual_reward_ref, + ) + summary = asdict(resolved) + # Task env values can contain operational secrets and are irrelevant to + # the preflight report. + summary["task_envs"] = sorted(resolved.task_envs) + summary.update( + { + "dataset_total": dataset_total, + "recent_terminal_tasks": terminal_count, + "recent_failures": failure_count, + "recent_failure_rate": failure_rate, + "unhealthy_override": unhealthy, + } + ) + summaries.append(summary) + return summaries + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + source = parser.add_mutually_exclusive_group(required=True) + source.add_argument("--streams-file") + source.add_argument( + "--streams-env-b64", + metavar="NAME", + help="Read base64-encoded inline Streams YAML from environment variable NAME", + ) + parser.add_argument("--base-url", required=True) + parser.add_argument("--default-harness", default="") + parser.add_argument("--retries", type=int, default=3) + parser.add_argument("--recent-tasks", type=int, default=100) + parser.add_argument("--min-terminal-tasks", type=int, default=1) + parser.add_argument("--max-failure-rate", type=float, default=0.8) + parser.add_argument("--allow-unhealthy", action="store_true") + args = parser.parse_args() + + streams_yaml_b64 = "" + if args.streams_env_b64: + streams_yaml_b64 = os.getenv(args.streams_env_b64, "") + if not streams_yaml_b64.strip(): + parser.error(f"environment variable {args.streams_env_b64!r} is empty") + + summaries = validate_streams( + args.streams_file or "", + streams_yaml_b64=streams_yaml_b64, + base_url=args.base_url, + default_harness=args.default_harness, + retries=args.retries, + recent_tasks=args.recent_tasks, + min_terminal_tasks=args.min_terminal_tasks, + max_failure_rate=args.max_failure_rate, + allow_unhealthy=args.allow_unhealthy, + ) + json.dump({"streams": summaries}, sys.stdout, ensure_ascii=False, indent=2) + sys.stdout.write("\n") + + +if __name__ == "__main__": + main() diff --git a/examples/swe/arena_stream_config_projection.py b/examples/swe/arena_stream_config_projection.py new file mode 100644 index 0000000000..4f157d7890 --- /dev/null +++ b/examples/swe/arena_stream_config_projection.py @@ -0,0 +1,57 @@ +# SPDX-License-Identifier: Apache-2.0 + +"""Project inline Arena Streams from an AReaL training config.""" + +from __future__ import annotations + +import argparse +import base64 +from pathlib import Path +from typing import Any + +import yaml + + +def _reject_interpolations(value: Any, *, path: str) -> None: + """Reject values whose preflight meaning could differ after Hydra resolution.""" + if isinstance(value, str) and "${" in value: + raise ValueError( + f"{path} contains a Hydra interpolation; use a literal value so " + "preflight and training see the same Stream configuration" + ) + if isinstance(value, list): + for index, item in enumerate(value): + _reject_interpolations(item, path=f"{path}[{index}]") + elif isinstance(value, dict): + for key, item in value.items(): + _reject_interpolations(item, path=f"{path}.{key}") + + +def project_arena_streams(config_path: Path) -> bytes: + """Return canonical top-level ``streams`` YAML from a training config.""" + config = yaml.safe_load(config_path.read_text(encoding="utf-8")) + if not isinstance(config, dict): + raise ValueError(f"Multi-Stream config must be a mapping: {config_path}") + econfig = config.get("econfig") + if not isinstance(econfig, dict): + raise ValueError(f"Multi-Stream config is missing econfig: {config_path}") + streams = econfig.get("arena_streams") + if not isinstance(streams, list) or not streams: + raise ValueError( + "Multi-Stream config requires a non-empty " + f"econfig.arena_streams: {config_path}" + ) + _reject_interpolations(streams, path="econfig.arena_streams") + return yaml.safe_dump({"streams": streams}, sort_keys=False).encode("utf-8") + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("config", type=Path) + args = parser.parse_args() + payload = project_arena_streams(args.config) + print(base64.b64encode(payload).decode("ascii"), end="") + + +if __name__ == "__main__": + main() diff --git a/examples/swe/arena_types.py b/examples/swe/arena_types.py new file mode 100644 index 0000000000..672a965f04 --- /dev/null +++ b/examples/swe/arena_types.py @@ -0,0 +1,30 @@ +"""Lightweight Arena configuration types shared by preflight and training.""" + +from __future__ import annotations + +from dataclasses import dataclass, field + + +@dataclass +class ArenaRewardRefConfig: + """Expected Arena reward implementation pinned for one Stream.""" + + key: str = "" + version: str = "" + + +@dataclass +class ArenaStreamConfig: + """Per-Stream routing, mixture weighting, and reward configuration.""" + + name: str = "" + stream_id: str = "" + sampling_weight: float = 1.0 + harness: str = "" + llm_protocol: str = "" + task_envs: dict[str, str] = field(default_factory=dict) + expected_reward_ref: ArenaRewardRefConfig = field( + default_factory=ArenaRewardRefConfig + ) + reward_threshold: float | None = None + reward_transform_fn: str = "" diff --git a/examples/swe/dataset_configs/flash-v3-multi-stream.yaml b/examples/swe/dataset_configs/flash-v3-multi-stream.yaml new file mode 100644 index 0000000000..e4eae62b54 --- /dev/null +++ b/examples/swe/dataset_configs/flash-v3-multi-stream.yaml @@ -0,0 +1,17 @@ +streams: +- name: vn_greenfield + stream_id: vn-greenfield-fixed + sampling_weight: 1 + harness: flash-vn-agent@2.0.3 + llm_protocol: anthropic + expected_reward_ref: + key: vn-creation-x2env-greenfield-reward + version: 2.0.8 +- name: favor_app_online_scratch + stream_id: favor-app-online-scratch-fixed + sampling_weight: 1 + harness: claude-code@1.2.31-native-direct-sfx-260821 + llm_protocol: anthropic + expected_reward_ref: + key: favor-app-online-scratch-reward + version: 2.0.2 diff --git a/examples/swe/filter_function.py b/examples/swe/filter_function.py index f97345cef5..ae15f0473d 100644 --- a/examples/swe/filter_function.py +++ b/examples/swe/filter_function.py @@ -70,7 +70,9 @@ def filter_function(sample): # Use original_rewards if available for filtering and statistics rewards = _get_rewards(sample) accept = bool((rewards != rewards[0]).any().item()) - _record_filter_result(accept=accept, all_correct=bool((rewards > 0).all().item())) + _record_filter_result( + accept=accept, all_correct=bool((rewards >= 1.0 - 1e-3).all().item()) + ) return accept diff --git a/examples/swe/reward_transforms.py b/examples/swe/reward_transforms.py new file mode 100644 index 0000000000..ad34554834 --- /dev/null +++ b/examples/swe/reward_transforms.py @@ -0,0 +1,17 @@ +"""Reusable reward transforms for SWE Arena workflows.""" + +import math +from typing import Any + + +def astra_partial_reward( + reward: float, + _data: dict[str, Any], + reward_threshold: float = 0.98, +) -> float: + """Keep Astra scores below the configured threshold at one tenth strength.""" + if not math.isfinite(reward) or not 0.0 <= reward <= 1.0: + raise ValueError(f"Astra reward must be finite and within [0, 1], got {reward}") + if reward < reward_threshold: + return reward * 0.1 + return 1.0 diff --git a/examples/swe/submit_arena.sh b/examples/swe/submit_arena.sh new file mode 100755 index 0000000000..27cdf9b48b --- /dev/null +++ b/examples/swe/submit_arena.sh @@ -0,0 +1,130 @@ +#!/usr/bin/env bash +# SPDX-License-Identifier: Apache-2.0 +# Submit one CPU controller; AReaL schedules the GPU workers from CONFIG_PATH. +set -euo pipefail +set +x + +case "${1:-}" in + -h|--help) + cat <<'EOF' +Usage: bash examples/swe/submit_arena.sh [--check-arena] + +Required environment: + AREAL_IMAGE, AREAL_PYTHON, MODEL_PATH, AREAL_FILEROOT, + ARENA_CREDENTIALS_FILE, AREAL_CONTROLLER_MOUNTS, AREAL_WORKER_MOUNTS +Optional: + AREAL_DIR (script's checkout), CONFIG_PATH (examples/swe/arena_grpo.yaml), + TRIAL_NAME (unique timestamp), AREAL_CONTAINER_BIN (apptainer), + SBATCH_PARTITION, SBATCH_ACCOUNT, SBATCH_RESERVATION, SBATCH_TIMELIMIT, + ARENA_CONTROLLER_TIME (01:30:00), ARENA_PREFLIGHT_MIN_TASKS (1). + +Use CONFIG_PATH to choose the single- or multi-stream YAML and GPU topology. +--check-arena submits a CPU-only validation job, without starting GPU workers. +Credentials stay in the shared file; never supply tokens as command arguments. +EOF + exit 0 ;; + ""|--check-arena|--controller|--run) ;; + *) echo 'Unknown argument; use --help' >&2; exit 2 ;; +esac +[[ $# -le 1 ]] || { echo 'Expected at most one argument' >&2; exit 2; } + +# sbatch copies this script to its spool: resolve the checkout BEFORE submission. +if [[ ${1:-} != --controller && ${1:-} != --run ]]; then + export AREAL_DIR=${AREAL_DIR:-$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")/../.." && pwd -P)} + export CONFIG_PATH=${CONFIG_PATH:-examples/swe/arena_grpo.yaml} + export TRIAL_NAME=${TRIAL_NAME:-arena-$(date +%Y%m%d-%H%M%S)-$$} + export ARENA_CHECK_ONLY=${1:-} +fi +: "${AREAL_DIR:?Set AREAL_DIR to the shared checkout}" +: "${CONFIG_PATH:?Set CONFIG_PATH to the training YAML}" +: "${AREAL_IMAGE:?Set AREAL_IMAGE to the container image}" +: "${AREAL_PYTHON:?Set AREAL_PYTHON to Python inside the container}" +: "${MODEL_PATH:?Set MODEL_PATH to the shared model directory}" +: "${AREAL_FILEROOT:?Set AREAL_FILEROOT to shared experiment storage}" +: "${ARENA_CREDENTIALS_FILE:?Set ARENA_CREDENTIALS_FILE to a shared private env file}" +: "${AREAL_CONTROLLER_MOUNTS:?Set container binds including shared storage and Slurm access}" +: "${AREAL_WORKER_MOUNTS:?Set worker container binds for code, model, credentials and output}" +export AREAL_DIR CONFIG_PATH AREAL_IMAGE AREAL_PYTHON MODEL_PATH AREAL_FILEROOT +export ARENA_CREDENTIALS_FILE AREAL_CONTROLLER_MOUNTS AREAL_WORKER_MOUNTS +export TRIAL_NAME SLURM_EXPORT_ENV=ALL +export SBATCH_TIMELIMIT=${SBATCH_TIMELIMIT:-01:00:00} +cd -- "$AREAL_DIR" +for name in AREAL_DIR AREAL_IMAGE MODEL_PATH AREAL_FILEROOT ARENA_CREDENTIALS_FILE; do + [[ ${!name} == /* ]] || { echo "$name must be an absolute shared path" >&2; exit 2; } +done +[[ -r $CONFIG_PATH && -r $AREAL_IMAGE && -d $MODEL_PATH && -r $ARENA_CREDENTIALS_FILE ]] || { + echo 'Config, image, model or credentials file is missing/unreadable' >&2; exit 2; +} +export ARENA_SUBMISSION_DIR="$AREAL_FILEROOT/submissions/$TRIAL_NAME" +mkdir -p -- "$ARENA_SUBMISSION_DIR" + +if [[ ${1:-} == --controller ]]; then + # Site-specific Slurm/munge mounts belong in the environment, not this script. + exec "${AREAL_CONTAINER_BIN:-apptainer}" exec --no-eval --pid --writable-tmpfs \ + --bind "$AREAL_CONTROLLER_MOUNTS" "$AREAL_IMAGE" \ + bash "$AREAL_DIR/examples/swe/submit_arena.sh" --run +elif [[ ${1:-} == --run ]]; then + # Older credential files also set deployment options. Explicit launch values win. + declare -A launch_env=() + for name in AREAL_DIR CONFIG_PATH AREAL_IMAGE AREAL_PYTHON MODEL_PATH AREAL_FILEROOT \ + ARENA_CREDENTIALS_FILE AREAL_WORKER_MOUNTS TRIAL_NAME ARENA_SUBMISSION_DIR ARENA_CHECK_ONLY \ + ARENA_STREAM_ID ARENA_OPENAPI_BASE ARENA_PREFLIGHT_MIN_TASKS SBATCH_TIMELIMIT; do + if [[ -v $name ]]; then + launch_env[$name]=${!name} + fi + done + set -a + source "$ARENA_CREDENTIALS_FILE" + set +a + for name in "${!launch_env[@]}"; do + export "$name=${launch_env[$name]}" + done + : "${ARENA_OPENAPI_BASE:?Credentials must export ARENA_OPENAPI_BASE}" + : "${ARENA_OPENAPI_TOKEN:?Credentials must export ARENA_OPENAPI_TOKEN}" + : "${ARENA_LLM_API_KEY:?Credentials must export ARENA_LLM_API_KEY}" + : "${SWE_RL_ADMIN_API_KEY:?Credentials must export SWE_RL_ADMIN_API_KEY}" + export SLURM_EXPORT_ENV=ALL + unset ALL_PROXY all_proxy HTTP_PROXY HTTPS_PROXY http_proxy https_proxy + export PYTHONPATH="$AREAL_DIR${PYTHONPATH:+:$PYTHONPATH}" PYTHONUNBUFFERED=1 + export PATH="$(dirname -- "$AREAL_PYTHON"):$PATH" + export APPTAINER_BIND="$AREAL_WORKER_MOUNTS" SINGULARITY_BIND="$AREAL_WORKER_MOUNTS" + # Resolve exactly the same config used by training, including Hydra defaults. + "$AREAL_PYTHON" - "$CONFIG_PATH" <<'PY' +import base64 +import json +import os +import sys +from dataclasses import asdict +from pathlib import Path + +import yaml + +from examples.swe.arena_config import load_arena_stream_configs +from examples.swe.arena_preflight import validate_streams +from examples.swe.utils import SWEPPOConfig +from areal.api.cli_args import load_expr_config + +args = ["--config", sys.argv[1]] +config, _ = load_expr_config(args, SWEPPOConfig) +if config.scheduler.type != "slurm" or config.econfig.dataset_source != "arena": + raise ValueError("This entrypoint requires the Slurm scheduler and Arena dataset") +streams = load_arena_stream_configs(config.econfig) +payload = yaml.safe_dump({"streams": [asdict(stream) for stream in streams]}) +summary = validate_streams( + streams_yaml_b64=base64.b64encode(payload.encode()).decode(), + base_url=config.econfig.arena_base_url, + min_terminal_tasks=int(os.environ.get("ARENA_PREFLIGHT_MIN_TASKS", "1")), +) +Path(os.environ["ARENA_SUBMISSION_DIR"], "preflight.json").write_text( + json.dumps({"streams": summary}, indent=2) + "\n", encoding="utf-8" +) +if not os.environ.get("ARENA_CHECK_ONLY"): + os.execv(sys.executable, [sys.executable, "-m", "examples.swe.train_swe_rl", *args]) +PY +else + exec sbatch --parsable --nodes=1 --ntasks=1 --cpus-per-task=4 --mem=16G \ + --gres=gpu:0 --time="${ARENA_CONTROLLER_TIME:-01:30:00}" --export=ALL \ + --job-name="$TRIAL_NAME-controller" --chdir="$AREAL_DIR" \ + --output="$ARENA_SUBMISSION_DIR/controller-%j.log" \ + "$AREAL_DIR/examples/swe/submit_arena.sh" --controller +fi diff --git a/examples/swe/submit_arena_flash_v3.sh b/examples/swe/submit_arena_flash_v3.sh new file mode 100755 index 0000000000..ef23c66b75 --- /dev/null +++ b/examples/swe/submit_arena_flash_v3.sh @@ -0,0 +1,67 @@ +#!/usr/bin/env bash +# SPDX-License-Identifier: Apache-2.0 +# Flash V3 Arena profiles from swe-dev; deployment paths come from the environment. +set -euo pipefail +set +x + +case "${1:-}" in + -h|--help) + cat <<'HELP' +Usage: bash examples/swe/submit_arena_flash_v3.sh [--print-profile|--check-arena] + +N_NODES selects 2, 4, 8 (default), or 16 shared training/inference worker nodes. +Defaults: batch 32, 12 samples, 128K context, 500 steps, two Arena streams. +Supply submit_arena.sh's required environment plus ROLLOUT_IMAGE, AWEX_ROOT, +FLASH_LINEAR_ATTENTION_ROOT. MODEL_PATH selects a Flash V3 checkpoint. +ACTOR_RUNTIME_PYTHONPATH and ROLLOUT_RUNTIME_PYTHONPATH select optional overlays. +DATASET_CONFIG overrides the versioned streams YAML. See README_arena.md for +runtime requirements and differences from the complete swe-dev recipe. +HELP + exit 0 ;; + ""|--print-profile|--check-arena) ;; + *) echo 'Unknown argument; use --help' >&2; exit 2 ;; +esac +[[ $# -le 1 ]] || { echo 'Expected at most one argument' >&2; exit 2; } +export N_NODES=${N_NODES:-8} +case "$N_NODES" in + 2) actor='(attn:d1p2t4c2|ffn:d1p2e8)'; rollout=d4t4p1 ;; + 4) actor='(attn:d2p2t4c2|ffn:d2p2e8)'; rollout=d8t4p1 ;; + 8) actor='(attn:d2p2t2c8|ffn:d4p2e8)'; rollout=d16t4p1 ;; + 16) actor='(attn:d4p2t2c8|ffn:d8p2e8)'; rollout=d32t4p1 ;; + *) echo 'N_NODES must be 2, 4, 8 or 16' >&2; exit 2 ;; +esac +export ACTOR_BACKEND="megatron:$actor" ROLLOUT_BACKEND="sglang:$rollout" +export TRAIN_BATCH_SIZE=${TRAIN_BATCH_SIZE:-32} N_SAMPLES=${N_SAMPLES:-12} +export TOTAL_TRAIN_STEPS=${TOTAL_TRAIN_STEPS:-500} +export MAX_CONCURRENT_ROLLOUTS=${MAX_CONCURRENT_ROLLOUTS:-1500} +export CONTEXT_LENGTH=${CONTEXT_LENGTH:-131072} +[[ $CONTEXT_LENGTH =~ ^[1-9][0-9]*$ ]] && (( CONTEXT_LENGTH >= 2 )) || { + echo 'CONTEXT_LENGTH must be an integer of at least 2' >&2; exit 2; +} +export MAX_TOKENS=${MAX_TOKENS:-$((CONTEXT_LENGTH - 1))} +export MAX_NEW_TOKENS=${MAX_NEW_TOKENS:-$MAX_TOKENS} +for name in MAX_TOKENS MAX_NEW_TOKENS; do + [[ ${!name} =~ ^[1-9][0-9]*$ ]] && (( ${!name} < CONTEXT_LENGTH )) || { + echo "$name must be positive and smaller than CONTEXT_LENGTH" >&2; exit 2; + } +done +export ACTOR_MAX_TOKENS_PER_MB=${ACTOR_MAX_TOKENS_PER_MB:-$CONTEXT_LENGTH} +if [[ ${1:-} == --print-profile ]]; then + printf '%s\n' "nodes=$N_NODES" "actor=$ACTOR_BACKEND" "rollout=$ROLLOUT_BACKEND" \ + "context=$CONTEXT_LENGTH" "max_tokens=$MAX_TOKENS" "max_new_tokens=$MAX_NEW_TOKENS" \ + "batch=$TRAIN_BATCH_SIZE" "samples=$N_SAMPLES" \ + "steps=$TOTAL_TRAIN_STEPS" "concurrency=$MAX_CONCURRENT_ROLLOUTS" + exit 0 +fi +script_dir=$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd -P) +export AREAL_DIR=${AREAL_DIR:-$(cd -- "$script_dir/../.." && pwd -P)} +export CONFIG_PATH=${CONFIG_PATH:-$AREAL_DIR/examples/swe/arena_flash_v3.yaml} +export DATASET_CONFIG=${DATASET_CONFIG:-$AREAL_DIR/examples/swe/dataset_configs/flash-v3-multi-stream.yaml} +: "${ROLLOUT_IMAGE:?Set ROLLOUT_IMAGE to the Flash V3 SGLang runtime}" +: "${AWEX_ROOT:?Set AWEX_ROOT to an AWEX checkout with the Flash V3 converter}" +: "${FLASH_LINEAR_ATTENTION_ROOT:?Set FLASH_LINEAR_ATTENTION_ROOT to compatible FLA sources}" +export ROLLOUT_IMAGE AWEX_ROOT FLASH_LINEAR_ATTENTION_ROOT +[[ -r $ROLLOUT_IMAGE && -r $DATASET_CONFIG && -d $AWEX_ROOT && -d $FLASH_LINEAR_ATTENTION_ROOT ]] || { + echo 'Rollout image, streams YAML, AWEX or FLA sources are missing/unreadable' >&2; exit 2; +} +exec bash "$AREAL_DIR/examples/swe/submit_arena.sh" "$@" diff --git a/examples/swe/train_swe_rl.py b/examples/swe/train_swe_rl.py index 312f3d9de2..4a1fc2ca1c 100644 --- a/examples/swe/train_swe_rl.py +++ b/examples/swe/train_swe_rl.py @@ -3,12 +3,24 @@ import json import sys import warnings +from collections import Counter +from dataclasses import asdict, replace from pathlib import Path from typing import Any from datasets import Dataset -from examples.swe.utils import SWEPPOConfig +from examples.swe.arena_client import ( + ArenaAPIError, + ArenaOpenAPIClient, + infer_llm_protocol_from_harness, + resolve_llm_protocol, +) +from examples.swe.arena_config import ( + build_weighted_arena_rows, + load_arena_stream_configs, +) +from examples.swe.utils import ArenaRewardRefConfig, ArenaStreamConfig, SWEPPOConfig from areal import PPOTrainer from areal.api.cli_args import load_expr_config @@ -88,6 +100,123 @@ def group_filter(x: dict[str, Any]): return x["rewards"].mean() <= 0.95 +def _stream_reward_ref(stream: dict[str, Any]) -> ArenaRewardRefConfig: + value = stream.get("default_reward_ref") + if value is None: + return ArenaRewardRefConfig() + if not isinstance(value, dict): + raise ArenaAPIError("Arena Stream default_reward_ref must be an object") + key = value.get("key") + version = value.get("version") + if ( + not isinstance(key, str) + or not key + or not isinstance(version, str) + or not version + ): + raise ArenaAPIError( + "Arena Stream default_reward_ref requires non-empty key and version" + ) + return ArenaRewardRefConfig(key=key, version=version) + + +def _resolve_arena_stream( + client: ArenaOpenAPIClient, + configured: ArenaStreamConfig, +) -> tuple[ArenaStreamConfig, list[dict[str, str]]]: + stream = client.resolve_stream(configured.stream_id) + resolved_stream_id = stream.get("stream_id") + if not isinstance(resolved_stream_id, str) or not resolved_stream_id: + raise ArenaAPIError("The selected Arena Stream is missing a valid stream_id") + actual_reward_ref = _stream_reward_ref(stream) + expected_reward_ref = configured.expected_reward_ref + if expected_reward_ref.key and expected_reward_ref != actual_reward_ref: + raise ArenaAPIError( + f"Arena Stream {configured.name!r} reward_ref drifted: expected " + f"{expected_reward_ref.key}@{expected_reward_ref.version}, got " + f"{actual_reward_ref.key}@{actual_reward_ref.version}" + ) + + if configured.llm_protocol: + llm_protocol = resolve_llm_protocol(stream, configured.llm_protocol) + elif configured.harness: + llm_protocol = infer_llm_protocol_from_harness(configured.harness) + else: + llm_protocol = resolve_llm_protocol(stream) + + resolved = replace( + configured, + stream_id=resolved_stream_id, + llm_protocol=llm_protocol, + expected_reward_ref=actual_reward_ref, + ) + rows = client.get_all_dataset_rows(resolved_stream_id, llm_protocol) + for row in rows: + row.update( + { + "arena_stream_name": resolved.name, + "reward_ref_key": actual_reward_ref.key, + "reward_ref_version": actual_reward_ref.version, + } + ) + return resolved, rows + + +def get_arena_mixture_dataset( + econfig, + *, + size_multiple: int = 1, +) -> tuple[Dataset, list[ArenaStreamConfig]]: + """Load and deterministically mix prompt rows from configured Arena Streams.""" + client = ArenaOpenAPIClient( + base_url=econfig.arena_base_url, + timeout=econfig.arena_request_timeout, + request_retries=econfig.arena_request_retries, + ) + configured_streams = load_arena_stream_configs(econfig) + resolved_streams: list[ArenaStreamConfig] = [] + rows_by_stream: dict[str, list[dict[str, str]]] = {} + for configured in configured_streams: + resolved, rows = _resolve_arena_stream(client, configured) + resolved_streams.append(resolved) + rows_by_stream[resolved.name] = rows + + rows = build_weighted_arena_rows( + rows_by_stream, + resolved_streams, + epoch_size=int(getattr(econfig, "arena_mixture_epoch_size", 0)), + size_multiple=size_multiple, + ) + dataset = Dataset.from_list(rows) + counts = Counter(row["arena_stream_name"] for row in rows) + logger.info( + "Created Arena mixture with %d prompt rows: %s", + len(dataset), + dict(sorted(counts.items())), + ) + if size_multiple > 1 and len(dataset) % size_multiple: + logger.warning( + "Arena raw union has %d rows, not divisible by training batch size %d; " + "drop_last will omit %d tail rows without repeating source data", + len(dataset), + size_multiple, + len(dataset) % size_multiple, + ) + return dataset, resolved_streams + + +def get_arena_dataset(econfig) -> tuple[Dataset, str]: + """Load one Arena Stream while retaining the original public return type.""" + + dataset, resolved_streams = get_arena_mixture_dataset(econfig) + if len(resolved_streams) != 1: + raise ValueError( + "get_arena_dataset supports one Stream; use " + "get_arena_mixture_dataset for a multi-Stream configuration" + ) + return dataset, resolved_streams[0].stream_id + + def _install_aweagent_deps_on_ray_nodes(aweagent_root: str): """Install AReaL-SWEAgent dependencies on all Ray GPU nodes. @@ -159,36 +288,49 @@ def main(args): econfig = config.econfig # When using Ray scheduler, ensure SWEAgent deps are on all nodes - if config.scheduler.type == "ray": + if config.scheduler.type == "ray" and econfig.dataset_source == "jsonl": import ray ray.init(address="auto", ignore_reinit_error=True) _install_aweagent_deps_on_ray_nodes(_resolve_aweagent_root(econfig)) - # Resolve dataset paths from config - train_path = config.train_dataset.path - valid_path = config.valid_dataset.path - - def resolve_path(p: str) -> str: - if Path(p).is_absolute() or Path(p).exists(): + if econfig.dataset_source == "arena": + train_dataset, resolved_streams = get_arena_mixture_dataset( + econfig, size_multiple=config.train_dataset.batch_size + ) + valid_dataset = train_dataset + econfig.arena_streams = resolved_streams + econfig.arena_streams_yaml_b64 = "" + econfig.arena_streams_file = "" + workflow = "examples.swe.arena_agent.ArenaStreamAgentWorkflow" + elif econfig.dataset_source == "jsonl": + # Resolve dataset paths from config + train_path = config.train_dataset.path + valid_path = config.valid_dataset.path + + def resolve_path(p: str) -> str: + if Path(p).is_absolute() or Path(p).exists(): + return p + if econfig.dataset_path: + candidate = Path(econfig.dataset_path) / p + if candidate.exists(): + return str(candidate) return p - if econfig.dataset_path: - candidate = Path(econfig.dataset_path) / p - if candidate.exists(): - return str(candidate) - return p - - train_dataset = get_swe_dataset( - dataset_path=resolve_path(train_path), - split="train", - ) - valid_dataset = get_swe_dataset( - dataset_path=resolve_path(valid_path), - split="test", - ) - # Build workflow kwargs - from dataclasses import asdict + train_dataset = get_swe_dataset( + dataset_path=resolve_path(train_path), + split="train", + ) + valid_dataset = get_swe_dataset( + dataset_path=resolve_path(valid_path), + split="test", + ) + + workflow = "examples.swe.agent.SWEAgentWorkflow" + else: + raise ValueError( + f"Unsupported econfig.dataset_source: {econfig.dataset_source!r}" + ) econfig_dict = asdict(econfig) workflow_kwargs = dict( @@ -213,7 +355,7 @@ def resolve_path(p: str) -> str: valid_dataset=valid_dataset, ) as trainer: trainer.train( - workflow="examples.swe.agent.SWEAgentWorkflow", + workflow=workflow, workflow_kwargs=workflow_kwargs, eval_workflow=None, eval_workflow_kwargs=eval_workflow_kwargs, diff --git a/examples/swe/utils.py b/examples/swe/utils.py index c53881cbff..c8a5ad9033 100644 --- a/examples/swe/utils.py +++ b/examples/swe/utils.py @@ -2,6 +2,9 @@ from dataclasses import dataclass, field +from examples.swe.arena_types import ArenaRewardRefConfig as ArenaRewardRefConfig +from examples.swe.arena_types import ArenaStreamConfig + from areal.api.cli_args import PPOConfig @@ -26,10 +29,178 @@ class SWEEnvConfig: timeout: Maximum time allowed for a single episode in seconds. """ + dataset_source: str = field( + default="jsonl", + metadata={ + "help": "Dataset and agent backend: 'jsonl' or 'arena'.", + "choices": ["jsonl", "arena"], + }, + ) dataset_path: str = field( default="", metadata={"help": "Path to the SWE-bench JSONL dataset file."}, ) + stream_id: str = field( + default="", + metadata={ + "help": ( + "Arena Stream id. When empty, the first active Stream returned by " + "the Arena OpenAPI is used." + ) + }, + ) + arena_streams: list[ArenaStreamConfig] = field( + default_factory=list, + metadata={ + "help": ( + "Optional inline Arena Stream mixture. The default epoch uses " + "every Stream row once; sampling_weight controls deterministic " + "interleaving and explicit subset allocation." + ) + }, + ) + arena_streams_yaml_b64: str = field( + default="", + metadata={ + "help": ( + "Optional base64-encoded inline YAML containing a top-level " + "'streams' list. " + "Mutually exclusive with arena_streams and arena_streams_file." + ) + }, + ) + arena_streams_file: str = field( + default="", + metadata={ + "help": ( + "Optional YAML/JSON file containing a top-level 'streams' list. " + "Mutually exclusive with inline arena_streams." + ) + }, + ) + arena_mixture_epoch_size: int = field( + default=0, + metadata={ + "help": ( + "Prompt rows in a deterministic weighted Arena virtual epoch. The " + "value cannot exceed the unique source rows; zero uses the complete " + "raw union without batch-padding repeats." + ) + }, + ) + arena_result_dump_dir: str = field( + default="", + metadata={ + "help": ( + "Optional directory for mode-0600 per-process Arena result JSONL " + "audit shards. Empty disables full raw-result persistence." + ) + }, + ) + arena_result_dump_max_bytes: int = field( + default=1_000_000, + metadata={ + "help": "Maximum serialized Arena raw payload bytes per audit record." + }, + ) + arena_base_url: str = field( + default="", + metadata={ + "help": ( + "Arena OpenAPI base URL. Defaults to the ARENA_OPENAPI_BASE " + "environment variable." + ) + }, + ) + arena_request_timeout: float = field( + default=60.0, + metadata={"help": "Arena Stream and dataset request timeout in seconds."}, + ) + arena_request_retries: int = field( + default=3, + metadata={"help": "Retries for transient Arena HTTP request failures."}, + ) + arena_poll_interval: float = field( + default=5.0, + metadata={"help": "Arena task-result polling interval in seconds."}, + ) + arena_registration_timeout: float = field( + default=180.0, + metadata={"help": "Arena LLM registration request timeout in seconds."}, + ) + arena_registration_probe_interval: float = field( + default=60.0, + metadata={ + "help": ( + "Seconds between worker-scoped Arena Session Gateway liveness " + "probes while tasks are unfinished." + ) + }, + ) + arena_llm_route_mode: str = field( + default="gateway", + metadata={ + "help": ( + "Arena LLM route: 'gateway' registers one temporary Arena model " + "per rollout; 'session_gateway' reuses one public Arena route per " + "rollout worker and selects each sample with a session capability; " + "'direct' " + "injects the rollout proxy's per-session URL and key into the " + "Harness task." + ), + "choices": ["gateway", "session_gateway", "direct"], + }, + ) + arena_llm_protocol: str = field( + default="", + metadata={ + "help": ( + "Optional Arena Harness client protocol label. Leave empty to infer " + "from the Stream Harness. The AReaL proxy is always registered as " + "an OpenAI Chat Completions upstream." + ), + "choices": ["", "anthropic", "responses", "chat_completions"], + }, + ) + arena_harness: str = field( + default="", + metadata={ + "help": ( + "Optional Arena Harness key and version passed to launch_one_task, " + "for example 'claude-code-with-skills@5.0.1'." + ) + }, + ) + arena_task_envs: dict[str, str] = field( + default_factory=dict, + metadata={ + "help": ( + "Additional environment variables passed to launch_one_task. " + "MODEL_NAME, BASE_URL, and API_KEY are managed by AReaL." + ) + }, + ) + arena_reward_threshold: float | None = field( + default=None, + metadata={ + "help": ( + "If set, map Arena rewards >= this threshold to 1.0 and lower " + "rewards to 0.0 when no transform is configured. Configured " + "transforms receive it as reward_threshold, and it also classifies " + "transformed rewards for pass@k metrics." + ) + }, + ) + arena_reward_transform_fn: str = field( + default="", + metadata={ + "help": ( + "Optional import path for a callable reward transform receiving " + "(reward, data) and returning a float. The callable must accept " + "reward_threshold when arena_reward_threshold is configured." + ) + }, + ) agent_type: str = field( default="swe", metadata={ diff --git a/examples/terminal_bench/agent/camel_terminal_agent.py b/examples/terminal_bench/agent/camel_terminal_agent.py index 5547aba768..e7360dc27d 100644 --- a/examples/terminal_bench/agent/camel_terminal_agent.py +++ b/examples/terminal_bench/agent/camel_terminal_agent.py @@ -11,13 +11,12 @@ from agent_rl_config import TaskTimeouts from camel.messages import BaseMessage from camel.toolkits import FunctionTool, TerminalToolkit -from transformers import PreTrainedTokenizerFast - from terminal_bench.handlers.trial_handler import TrialHandler from terminal_bench.parsers.base_parser import UnitTestStatus from terminal_bench.parsers.parser_factory import ParserFactory from terminal_bench.terminal.docker_compose_manager import DockerComposeManager from terminal_bench.terminal.terminal import Terminal +from transformers import PreTrainedTokenizerFast from areal.experimental.camel.openai_model import AReaLOpenAICompatibleModel from areal.utils.perf_tracer import ( diff --git a/tests/experimental/openai/test_context_overflow_recovery.py b/tests/experimental/openai/test_context_overflow_recovery.py new file mode 100644 index 0000000000..870c395c88 --- /dev/null +++ b/tests/experimental/openai/test_context_overflow_recovery.py @@ -0,0 +1,310 @@ +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +from unittest.mock import AsyncMock + +import pytest + +from areal.experimental.openai.proxy import workflow as workflow_module +from areal.experimental.openai.proxy.server import ( + derive_session_gateway_api_key, + derive_session_gateway_token, +) +from areal.experimental.openai.proxy.workflow import OpenAIProxyWorkflow +from areal.experimental.openai.types import InteractionWithTokenLogpReward +from areal.infra import workflow_context +from areal.infra.workflow_context import WorkflowContext +from areal.utils import stats_tracker + + +class _FailingAgent: + async def run(self, data, **kwargs): + raise RuntimeError("harness failed") + + +class _SuccessfulAgent: + async def run(self, data, **kwargs): + return 1.0 + + +class _RecordingSuccessfulAgent(_SuccessfulAgent): + def __init__(self): + self.persisted_rewards = [] + + async def persist_episode_result(self, data, reward): + self.persisted_rewards.append((data, reward)) + + +class _SystemFailingAgent(_FailingAgent): + @staticmethod + def classify_proxy_failure(error, *, context_overflow, interaction_count): + return "system_failure_reject" + + +class _ModelFailingAgent(_FailingAgent): + @staticmethod + def classify_proxy_failure(error, *, context_overflow, interaction_count): + return "model_failure_zero" + + +class _RecordingFailingAgent(_FailingAgent): + def __init__(self): + self.failure_dispositions = [] + self.persisted_rewards = [] + + async def record_failure_disposition(self, data, error, disposition): + self.failure_dispositions.append((data, error, disposition)) + + async def persist_episode_result(self, data, reward): + self.persisted_rewards.append((data, reward)) + + +class _FakeProxyClient: + context_overflow = True + context_overflow_message = "prompt exceeds context window" + system_error = False + system_error_message = "" + + def __init__(self, *args, interaction_count: int = 1, **kwargs): + self.session_api_key = "session-key" + self.session_id = "session-1" + self.interaction_count = interaction_count + self.interaction = InteractionWithTokenLogpReward(reward=None) + self.last_reward = None + + async def __aenter__(self): + return self + + async def __aexit__(self, exc_type, exc_val, exc_tb): + return None + + async def set_last_reward(self, reward: float): + self.last_reward = reward + self.interaction.reward = reward + + async def set_reward(self, completion_id: str, reward: float): + raise AssertionError("Per-completion rewards are not expected") + + async def export_interactions(self, **kwargs): + if self.interaction_count == 0: + return {} + return {"completion-1": self.interaction} + + +@pytest.mark.asyncio +async def test_inline_agent_receives_proxy_session_id(monkeypatch): + calls = [] + + class RecordingAgent: + async def run(self, data, **kwargs): + calls.append((data, kwargs)) + return 1.0 + + http_client = object() + monkeypatch.setattr( + workflow_context, "get_httpx_client", AsyncMock(return_value=http_client) + ) + workflow = OpenAIProxyWorkflow(mode="inline", agent=RecordingAgent()) + + reward = await workflow._run_agent( + "session-key", + {"data_id": "data-1"}, + session_id="session-1", + worker_runtime="worker-1", + ) + + assert reward == 1.0 + assert calls == [ + ( + {"data_id": "data-1"}, + { + "base_url": "", + "http_client": http_client, + "api_key": "session-key", + "proxy_gateway_api_key": derive_session_gateway_api_key( + workflow._admin_api_key + ), + "session_id": "session-1", + "proxy_session_token": derive_session_gateway_token( + workflow._admin_api_key, "session-1" + ), + "worker_runtime": "worker-1", + }, + ) + ] + + +@pytest.mark.asyncio +async def test_context_overflow_recovers_existing_interactions_with_zero_reward( + monkeypatch, +): + """Harness failure after overflow should return the preceding trajectory.""" + fake_client = _FakeProxyClient() + monkeypatch.setattr( + workflow_module, "OpenAIProxyClient", lambda *args, **kwargs: fake_client + ) + agent = _RecordingFailingAgent() + workflow = OpenAIProxyWorkflow(mode="inline", agent=agent) + monkeypatch.setattr(workflow, "_grant_capacity", AsyncMock()) + monkeypatch.setattr( + workflow_context, "get_aiohttp_session", AsyncMock(return_value=object()) + ) + stats_tracker.export_all(reset=True) + workflow_context.set(WorkflowContext(task_id=1)) + try: + result = await workflow.arun_episode(engine=None, data={}) + finally: + workflow_context.set(WorkflowContext()) + + assert result == {"completion-1": fake_client.interaction} + assert fake_client.last_reward == 0.0 + assert len(agent.failure_dispositions) == 1 + assert agent.failure_dispositions[0][2] == "model_failure_zero" + assert agent.persisted_rewards == [({}, 0.0)] + stats = stats_tracker.export_all(reset=True) + assert stats["rollout/context_overflow"] == 1.0 + assert stats["rollout/reward"] == 0.0 + + +@pytest.mark.asyncio +async def test_context_overflow_overrides_successful_agent_reward_with_zero( + monkeypatch, +): + """Any observed overflow should make the recovered trajectory a failure.""" + fake_client = _FakeProxyClient() + monkeypatch.setattr( + workflow_module, "OpenAIProxyClient", lambda *args, **kwargs: fake_client + ) + workflow = OpenAIProxyWorkflow(mode="inline", agent=_SuccessfulAgent()) + monkeypatch.setattr(workflow, "_grant_capacity", AsyncMock()) + monkeypatch.setattr( + workflow_context, "get_aiohttp_session", AsyncMock(return_value=object()) + ) + workflow_context.set(WorkflowContext(task_id=4)) + try: + result = await workflow.arun_episode(engine=None, data={}) + finally: + workflow_context.set(WorkflowContext()) + + assert result == {"completion-1": fake_client.interaction} + assert fake_client.last_reward == 0.0 + assert fake_client.interaction.reward == 0.0 + + +@pytest.mark.asyncio +async def test_first_request_context_overflow_without_interactions_returns_none( + monkeypatch, +): + """No trajectory can be recovered when overflow precedes all generations.""" + fake_client = _FakeProxyClient(interaction_count=0) + monkeypatch.setattr( + workflow_module, "OpenAIProxyClient", lambda *args, **kwargs: fake_client + ) + workflow = OpenAIProxyWorkflow(mode="inline", agent=_FailingAgent()) + monkeypatch.setattr(workflow, "_grant_capacity", AsyncMock()) + monkeypatch.setattr( + workflow_context, "get_aiohttp_session", AsyncMock(return_value=object()) + ) + workflow_context.set(WorkflowContext(task_id=2)) + try: + result = await workflow.arun_episode(engine=None, data={}) + finally: + workflow_context.set(WorkflowContext()) + + assert result is None + assert fake_client.last_reward is None + + +@pytest.mark.asyncio +async def test_successful_agent_with_empty_export_is_rejected_and_audited(monkeypatch): + """Dropped/orphan-only exports are not usable grouped-rollout slots.""" + fake_client = _FakeProxyClient(interaction_count=0) + fake_client.context_overflow = False + monkeypatch.setattr( + workflow_module, "OpenAIProxyClient", lambda *args, **kwargs: fake_client + ) + agent = _RecordingSuccessfulAgent() + workflow = OpenAIProxyWorkflow(mode="inline", agent=agent) + monkeypatch.setattr(workflow, "_grant_capacity", AsyncMock()) + monkeypatch.setattr( + workflow_context, "get_aiohttp_session", AsyncMock(return_value=object()) + ) + stats_tracker.export_all(reset=True) + workflow_context.set(WorkflowContext(task_id=6)) + try: + result = await workflow.arun_episode(engine=None, data={}) + finally: + workflow_context.set(WorkflowContext()) + + assert result is None + assert agent.persisted_rewards == [({}, None)] + assert "rollout/reward" not in stats_tracker.export_all(reset=True) + + +@pytest.mark.asyncio +async def test_non_context_agent_failure_still_propagates(monkeypatch): + """Infrastructure and unrelated Harness failures must not become reward zero.""" + fake_client = _FakeProxyClient() + fake_client.context_overflow = False + monkeypatch.setattr( + workflow_module, "OpenAIProxyClient", lambda *args, **kwargs: fake_client + ) + workflow = OpenAIProxyWorkflow(mode="inline", agent=_FailingAgent()) + monkeypatch.setattr(workflow, "_grant_capacity", AsyncMock()) + monkeypatch.setattr( + workflow_context, "get_aiohttp_session", AsyncMock(return_value=object()) + ) + workflow_context.set(WorkflowContext(task_id=3)) + try: + with pytest.raises(RuntimeError, match="harness failed"): + await workflow.arun_episode(engine=None, data={}) + finally: + workflow_context.set(WorkflowContext()) + + +@pytest.mark.asyncio +async def test_context_overflow_does_not_mask_classified_system_failure(monkeypatch): + """Authoritative system attribution must override an overflow signal.""" + fake_client = _FakeProxyClient() + monkeypatch.setattr( + workflow_module, "OpenAIProxyClient", lambda *args, **kwargs: fake_client + ) + workflow = OpenAIProxyWorkflow(mode="inline", agent=_SystemFailingAgent()) + monkeypatch.setattr(workflow, "_grant_capacity", AsyncMock()) + monkeypatch.setattr( + workflow_context, "get_aiohttp_session", AsyncMock(return_value=object()) + ) + workflow_context.set(WorkflowContext(task_id=5)) + try: + with pytest.raises(RuntimeError, match="harness failed"): + await workflow.arun_episode(engine=None, data={}) + finally: + workflow_context.set(WorkflowContext()) + + assert fake_client.last_reward is None + + +@pytest.mark.asyncio +async def test_proxy_system_error_overrides_model_failure_classifier(monkeypatch): + """An observed proxy 500 must reject even if Harness reports agent failure.""" + fake_client = _FakeProxyClient() + fake_client.context_overflow = False + fake_client.system_error = True + fake_client.system_error_message = "backend unavailable" + monkeypatch.setattr( + workflow_module, "OpenAIProxyClient", lambda *args, **kwargs: fake_client + ) + workflow = OpenAIProxyWorkflow(mode="inline", agent=_ModelFailingAgent()) + monkeypatch.setattr(workflow, "_grant_capacity", AsyncMock()) + monkeypatch.setattr( + workflow_context, "get_aiohttp_session", AsyncMock(return_value=object()) + ) + workflow_context.set(WorkflowContext(task_id=7)) + try: + with pytest.raises(RuntimeError, match="harness failed"): + await workflow.arun_episode(engine=None, data={}) + finally: + workflow_context.set(WorkflowContext()) + + assert fake_client.last_reward is None diff --git a/tests/experimental/openai/test_proxy_rollout_server.py b/tests/experimental/openai/test_proxy_rollout_server.py index 35c51b497b..7accd0e29d 100644 --- a/tests/experimental/openai/test_proxy_rollout_server.py +++ b/tests/experimental/openai/test_proxy_rollout_server.py @@ -11,8 +11,13 @@ from areal.api import ModelResponse from areal.experimental.openai.cache import InteractionCache +from areal.experimental.openai.client import ContextLengthExceededError from areal.experimental.openai.proxy import proxy_rollout_server as srv -from areal.experimental.openai.proxy.server import SessionData +from areal.experimental.openai.proxy.server import ( + SessionData, + derive_session_gateway_api_key, + derive_session_gateway_token, +) from areal.experimental.openai.proxy.tensor_reference import GroupTensorStoreRegistry from areal.experimental.openai.types import InteractionWithTokenLogpReward from areal.infra.processor_cache import ProcessorCacheRegistry @@ -58,6 +63,97 @@ def _admin_headers(): return {"Authorization": f"Bearer {_ADMIN_KEY}"} +class _Request: + def __init__(self, headers: dict[str, str]): + self.headers = headers + self.url = SimpleNamespace(path="/v1/chat/completions") + + +def test_require_session_key_accepts_active_session_key(): + srv._api_key_to_session["session-key"] = "task-1-0" + + session_id = srv._require_session_key( + _Request({"authorization": "Bearer session-key"}) + ) + + assert session_id == "task-1-0" + + +def test_require_session_key_accepts_gateway_capability_for_active_session(): + srv._session_to_api_key["task-1-0"] = "session-key" + gateway_key = derive_session_gateway_api_key(_ADMIN_KEY) + session_token = derive_session_gateway_token(_ADMIN_KEY, "task-1-0") + + session_id = srv._require_session_key( + _Request( + { + "authorization": "Bearer arena-gateway-key", + "x-api-key": gateway_key, + "x-session-id": "task-1-0", + "x-session-token": session_token, + } + ) + ) + + assert session_id == "task-1-0" + + +@pytest.mark.parametrize("session_id", [None, "missing-session"]) +def test_require_session_key_rejects_gateway_without_active_session(session_id): + headers = { + "authorization": (f"Bearer {derive_session_gateway_api_key(_ADMIN_KEY)}") + } + if session_id is not None: + headers["x-session-id"] = session_id + headers["x-session-token"] = derive_session_gateway_token( + _ADMIN_KEY, session_id + ) + + with pytest.raises(srv.HTTPException) as exc_info: + srv._require_session_key(_Request(headers)) + + assert exc_info.value.status_code == 401 + + +def test_require_session_key_rejects_capability_for_another_active_session(): + srv._session_to_api_key["task-1-0"] = "session-key-1" + srv._session_to_api_key["task-2-0"] = "session-key-2" + + with pytest.raises(srv.HTTPException) as exc_info: + srv._require_session_key( + _Request( + { + "x-api-key": derive_session_gateway_api_key(_ADMIN_KEY), + "x-session-id": "task-2-0", + "x-session-token": derive_session_gateway_token( + _ADMIN_KEY, "task-1-0" + ), + } + ) + ) + + assert exc_info.value.status_code == 401 + + +@pytest.mark.parametrize("path", ["/rl/set_reward", "/rl/end_session"]) +def test_gateway_capability_cannot_mutate_session(path): + """A public gateway must not assign rewards or close a training session.""" + srv._session_to_api_key["task-1-0"] = "session-key" + request = _Request( + { + "x-api-key": derive_session_gateway_api_key(_ADMIN_KEY), + "x-session-id": "task-1-0", + "x-session-token": derive_session_gateway_token(_ADMIN_KEY, "task-1-0"), + } + ) + request.url = SimpleNamespace(path=path) + + with pytest.raises(srv.HTTPException) as exc_info: + srv._require_session_key(request) + + assert exc_info.value.status_code == 403 + + # --------------------------------------------------------------------------- # Tests: health reports forked worker identity # --------------------------------------------------------------------------- @@ -95,6 +191,171 @@ def test_proxy_worker_index_falls_back_to_slurm_env(monkeypatch): assert srv._resolve_worker_index(-1) == 5 +# --------------------------------------------------------------------------- +# Tests: message preprocessing +# --------------------------------------------------------------------------- + + +def test_preprocess_messages_flattens_text_blocks_and_preserves_images(monkeypatch): + """OpenAI-routed Claude text blocks should be flattened before inference.""" + + class RemoveReminder: + def __call__(self, messages): + for message in messages: + if isinstance(message.get("content"), str): + message["content"] = message["content"].replace( + "reminder", "processed" + ) + return messages + + monkeypatch.setattr(srv, "_message_preprocessors", [RemoveReminder()]) + image_content = [ + {"type": "text", "text": "describe"}, + {"type": "image_url", "image_url": {"url": "https://example/image.png"}}, + ] + messages = [ + {"role": "system", "content": image_content}, + { + "role": "user", + "content": [ + {"type": "text", "text": "hello"}, + { + "type": "tool_result", + "content": [{"type": "text", "text": "reminder"}], + }, + ], + }, + {"role": "user", "content": image_content}, + ] + + result = srv._preprocess_messages(messages) + + assert isinstance(result[0]["content"], str) + assert result[1]["content"] == "hello\nprocessed" + assert result[2]["content"] == image_content + + +def test_prepare_request_messages_normalizes_tuple_and_system_content(monkeypatch): + """Tuple-backed request messages should not bypass system normalization.""" + monkeypatch.setattr(srv, "_message_preprocessors", []) + messages = ( + { + "role": "system", + "content": ({"type": "text", "text": "system prompt"},), + }, + {"role": "user", "content": "hello"}, + ) + + result = srv._prepare_request_messages(messages) + + assert isinstance(result, list) + assert result[0]["content"] == "system prompt" + + +def test_prepare_request_messages_preserves_generator_after_unsupported_item( + monkeypatch, +): + """Unsupported generator items must not leave a partially consumed iterator.""" + monkeypatch.setattr(srv, "_message_preprocessors", []) + unsupported = object() + + def message_generator(): + yield {"role": "user", "content": "hello"} + yield unsupported + + result = srv._prepare_request_messages(message_generator()) + + assert isinstance(result, list) + assert result == [{"role": "user", "content": "hello"}, unsupported] + + +@pytest.mark.asyncio +async def test_chat_completions_accepts_image_content_in_tool_message(monkeypatch): + """FastAPI must not lazily validate gateway image blocks as text-only.""" + captured_messages = None + + async def create( + *, messages, areal_cache, model="areal", temperature=1.0, top_p=1.0 + ): + nonlocal captured_messages + del areal_cache, model, temperature, top_p + captured_messages = messages + return {"ok": True} + + monkeypatch.setattr( + srv, + "_openai_client", + SimpleNamespace( + chat=SimpleNamespace(completions=SimpleNamespace(create=create)) + ), + ) + srv._session_cache["image-session"] = SessionData(session_id="image-session") + srv._api_key_to_session["image-key"] = "image-session" + image_part = { + "type": "image_url", + "image_url": {"url": "data:image/png;base64,aW1hZ2U="}, + } + + async with _client() as client: + response = await client.post( + "/v1/chat/completions", + headers={"Authorization": "Bearer image-key"}, + json={ + "model": "areal", + "messages": [ + { + "role": "tool", + "tool_call_id": "tool-1", + "content": [image_part], + } + ], + }, + ) + + assert response.status_code == 200 + assert captured_messages[0]["content"] == [image_part] + + +@pytest.mark.asyncio +async def test_internal_generation_failure_is_reported_when_session_ends(monkeypatch): + """A proxy-side 500 must remain distinguishable from model behavior.""" + + async def create( + *, messages, areal_cache, model="areal", temperature=1.0, top_p=1.0 + ): + del messages, areal_cache, model, temperature, top_p + raise RuntimeError("backend unavailable") + + monkeypatch.setattr( + srv, + "_openai_client", + SimpleNamespace( + chat=SimpleNamespace(completions=SimpleNamespace(create=create)) + ), + ) + srv._session_cache["failed-session"] = SessionData(session_id="failed-session") + srv._api_key_to_session["failed-key"] = "failed-session" + + async with _client() as client: + generation = await client.post( + "/v1/chat/completions", + headers={"Authorization": "Bearer failed-key"}, + json={ + "model": "areal", + "messages": [{"role": "user", "content": "hello"}], + }, + ) + ended = await client.post( + "/rl/end_session", + headers={"Authorization": "Bearer failed-key"}, + json={}, + ) + + assert generation.status_code == 500 + assert ended.json()["system_error"] is True + assert "backend unavailable" in ended.json()["system_error_message"] + + # --------------------------------------------------------------------------- # Tests: start_session with provided api_key # --------------------------------------------------------------------------- @@ -217,6 +478,100 @@ async def test_end_session_returns_interaction_count(self, monkeypatch): data = resp_end.json() assert data["interaction_count"] == 0 + @pytest.mark.asyncio + async def test_context_overflow_is_reported_when_session_ends(self, monkeypatch): + """A context error should persist on the session for workflow recovery.""" + + async def overflow_create(*, areal_cache, **_kwargs): + del areal_cache + raise ContextLengthExceededError("prompt exceeds context window") + + monkeypatch.setattr( + srv, + "_openai_client", + SimpleNamespace( + chat=SimpleNamespace( + completions=SimpleNamespace(create=overflow_create) + ) + ), + ) + monkeypatch.setattr(srv, "_capacity", 1) + async with _client() as client: + start = await client.post( + "/rl/start_session", + headers=_admin_headers(), + json={"task_id": "overflow"}, + ) + api_key = start.json()["api_key"] + headers = {"Authorization": f"Bearer {api_key}"} + + generation = await client.post( + "/chat/completions", + headers=headers, + json={ + "model": "areal", + "messages": [{"role": "user", "content": "long prompt"}], + }, + ) + ended = await client.post("/rl/end_session", headers=headers, json={}) + + assert generation.status_code == 400 + assert generation.json()["detail"]["type"] == "context_length_exceeded" + assert ended.json()["context_overflow"] is True + assert ended.json()["context_overflow_message"] == ( + "prompt exceeds context window" + ) + + @pytest.mark.asyncio + async def test_anthropic_streaming_context_overflow_returns_400(self, monkeypatch): + """Anthropic streaming must preserve the structured context error.""" + + async def overflow_create(*, areal_cache, **_kwargs): + del areal_cache + raise ContextLengthExceededError("prompt exceeds context window") + + monkeypatch.setattr( + srv, + "_openai_client", + SimpleNamespace( + chat=SimpleNamespace( + completions=SimpleNamespace(create=overflow_create) + ) + ), + ) + monkeypatch.setattr( + srv, + "_translate_anthropic_to_openai_request", + lambda _request: { + "model": "areal", + "messages": [{"role": "user", "content": "long prompt"}], + }, + ) + monkeypatch.setattr(srv, "_capacity", 1) + + async with _client() as client: + start = await client.post( + "/rl/start_session", + headers=_admin_headers(), + json={"task_id": "anthropic-stream-overflow"}, + ) + api_key = start.json()["api_key"] + response = await client.post( + "/v1/messages", + headers={"Authorization": f"Bearer {api_key}"}, + json={ + "model": "areal", + "messages": [{"role": "user", "content": "long prompt"}], + "stream": True, + }, + ) + + assert response.status_code == 400 + assert response.json()["detail"] == { + "type": "context_length_exceeded", + "message": "prompt exceeds context window", + } + def test_setup_openai_client_loads_and_passes_vlm_processor(monkeypatch): """The v1 proxy should reuse the model processor for trajectory export.""" diff --git a/tests/experimental/openai/test_streaming_chat_completions.py b/tests/experimental/openai/test_streaming_chat_completions.py index 411554ab51..faed2c0020 100644 --- a/tests/experimental/openai/test_streaming_chat_completions.py +++ b/tests/experimental/openai/test_streaming_chat_completions.py @@ -1,4 +1,4 @@ -"""Tests for streaming and non-streaming behaviour of the ``/chat/completions`` endpoint. +"""Tests for streaming and non-streaming chat-completions behaviour. The proxy rollout server's ``chat_completions`` handler must correctly return a ``StreamingResponse`` (SSE) when ``stream=True`` is requested, and a plain JSON @@ -9,6 +9,7 @@ from __future__ import annotations +import asyncio import json import threading from unittest.mock import MagicMock @@ -18,10 +19,13 @@ from openai.types.chat.chat_completion import Choice from openai.types.chat.chat_completion_chunk import Choice as ChunkChoice from openai.types.chat.chat_completion_chunk import ChoiceDelta +from openai.types.completion_usage import CompletionUsage +from starlette.responses import StreamingResponse from areal.api import ModelResponse from areal.experimental.openai.client import AsyncCompletionsWithReward from areal.experimental.openai.proxy import proxy_rollout_server as srv +from areal.experimental.openai.types import InteractionWithTokenLogpReward # --------------------------------------------------------------------------- # Helpers (same pattern as test_proxy_rollout_server.py) @@ -86,6 +90,10 @@ async def _fake_create( ``inspect.signature``-based filtering of request-body fields. """ if stream: + if areal_cache is not None: + areal_cache["chatcmpl-test"] = InteractionWithTokenLogpReward( + messages=list(messages or []) + ) async def _gen(): yield ChatCompletionChunk( @@ -116,6 +124,7 @@ async def _gen(): created=0, model=model, object="chat.completion", + usage=CompletionUsage(prompt_tokens=1, completion_tokens=1, total_tokens=2), ) @@ -132,12 +141,63 @@ def _mock_openai_client(monkeypatch): # --------------------------------------------------------------------------- +@pytest.mark.asyncio +@pytest.mark.parametrize("path", ["/v1/chat/completions", "/v1/messages"]) +@pytest.mark.parametrize("stream", [False, True]) +async def test_message_preprocessor_runs_once_per_request(monkeypatch, path, stream): + """Non-idempotent preprocessors must not duplicate injected prompt content.""" + injected_message = {"role": "system", "content": "Injected instruction"} + preprocessor = MagicMock(side_effect=lambda messages: [*messages, injected_message]) + captured_messages = [] + + async def create( + *, messages, model, stream=False, temperature=None, top_p=None, areal_cache=None + ): + captured_messages.append(messages) + return await _fake_create( + messages=messages, model=model, stream=stream, areal_cache=areal_cache + ) + + mock_client = MagicMock() + mock_client.chat.completions.create = create + monkeypatch.setattr(srv, "_openai_client", mock_client) + monkeypatch.setattr(srv, "_message_preprocessors", [preprocessor]) + monkeypatch.setattr(srv, "_capacity", 1) + + async with _client() as client: + start = await client.post( + "/rl/start_session", + headers=_admin_headers(), + json={"task_id": "preprocess"}, + ) + response = await client.post( + path, + headers=_session_headers(start.json()["api_key"]), + json={ + "model": "claude-compatible", + "max_tokens": 16, + "messages": [ + {"role": "user", "content": [{"type": "text", "text": "hello"}]} + ], + "stream": stream, + }, + ) + + preprocessor.assert_called_once() + assert captured_messages == [ + [{"role": "user", "content": "hello"}, injected_message] + ] + assert response.status_code == 200 + assert "hello" in response.text + + class TestChatCompletionsEndpoint: - """Verify streaming and non-streaming responses from ``/chat/completions``.""" + """Verify both supported chat-completions paths.""" @pytest.mark.asyncio + @pytest.mark.parametrize("path", ["/chat/completions", "/v1/chat/completions"]) async def test_streaming_returns_sse_response( - self, monkeypatch, _mock_openai_client + self, monkeypatch, _mock_openai_client, path ): """``stream=True`` returns an SSE stream (``text/event-stream``).""" monkeypatch.setattr(srv, "_capacity", 1) @@ -152,7 +212,7 @@ async def test_streaming_returns_sse_response( api_key = resp.json()["api_key"] resp = await client.post( - "/chat/completions", + path, headers=_session_headers(api_key), json={ "messages": [{"role": "user", "content": "hi"}], @@ -170,17 +230,93 @@ async def test_streaming_returns_sse_response( for line in resp.text.strip().split("\n\n") if line.startswith("data: ") ] - assert len(events) >= 2 # at least one chunk + [DONE] + assert len(events) >= 3 # prelude + generated chunk + [DONE] assert events[-1] == "data: [DONE]" - # Verify the data payload is a valid ChatCompletionChunk - chunk = json.loads(events[0].removeprefix("data: ")) - assert chunk["object"] == "chat.completion.chunk" - assert chunk["model"] == "test" - assert chunk["choices"][0]["delta"]["content"] == "hello" + chunks = [json.loads(event.removeprefix("data: ")) for event in events[:-1]] + assert chunks[0]["object"] == "chat.completion.chunk" + assert chunks[0]["model"] == "test" + assert chunks[0]["choices"][0]["delta"] == { + "role": "assistant", + "content": "", + } + assert chunks[1]["choices"][0]["delta"]["content"] == "hello" + assert {chunk["id"] for chunk in chunks} == {chunks[0]["id"]} + assert chunks[0]["id"] != "chatcmpl-test" + + resp = await client.post( + "/rl/set_reward", + headers=_session_headers(api_key), + json={"interaction_id": chunks[0]["id"], "reward": 0.75}, + ) + assert resp.status_code == 200 + + @pytest.mark.asyncio + async def test_streaming_starts_and_heartbeats_before_generation_finishes( + self, monkeypatch, _mock_openai_client + ): + """The proxy must not wait for simulated streaming generation to finish.""" + generation_started = asyncio.Event() + release_generation = asyncio.Event() + + async def slow_create( + *, + messages=None, + model=None, + stream=None, + temperature=None, + top_p=None, + areal_cache=None, + **kwargs, + ): + generation_started.set() + await release_generation.wait() + return await _fake_create( + messages=messages, + model=model, + stream=stream, + temperature=temperature, + top_p=top_p, + areal_cache=areal_cache, + **kwargs, + ) + + srv._openai_client.chat.completions.create = slow_create + monkeypatch.setattr(srv, "_STREAM_HEARTBEAT_INTERVAL_SECONDS", 0.01) + srv._session_cache["slow-session"] = srv.SessionData(session_id="slow-session") + + response = await asyncio.wait_for( + srv.chat_completions( + { + "messages": [{"role": "user", "content": "hi"}], + "model": "test", + "stream": True, + }, + session_id="slow-session", + ), + timeout=0.1, + ) + + assert isinstance(response, StreamingResponse) + iterator = response.body_iterator + first = await asyncio.wait_for(anext(iterator), timeout=0.1) + prelude = json.loads(first.removeprefix("data: ")) + assert prelude["choices"][0]["delta"]["content"] == "" + + next_event = asyncio.create_task(anext(iterator)) + await asyncio.wait_for(generation_started.wait(), timeout=0.1) + heartbeat = await asyncio.wait_for(next_event, timeout=0.1) + assert heartbeat == ": ping\n\n" + + release_generation.set() + remaining = [chunk async for chunk in iterator] + assert remaining[-1] == "data: [DONE]\n\n" @pytest.mark.asyncio - async def test_non_streaming_returns_json(self, monkeypatch, _mock_openai_client): + @pytest.mark.parametrize("path", ["/chat/completions", "/v1/chat/completions"]) + async def test_non_streaming_returns_json( + self, monkeypatch, _mock_openai_client, path + ): """Without ``stream``, returns a JSON ``ChatCompletion``.""" monkeypatch.setattr(srv, "_capacity", 1) @@ -194,7 +330,7 @@ async def test_non_streaming_returns_json(self, monkeypatch, _mock_openai_client api_key = resp.json()["api_key"] resp = await client.post( - "/chat/completions", + path, headers=_session_headers(api_key), json={ "messages": [{"role": "user", "content": "hi"}], diff --git a/tests/test_arena_worker_cleanup.py b/tests/test_arena_worker_cleanup.py new file mode 100644 index 0000000000..fa11518de0 --- /dev/null +++ b/tests/test_arena_worker_cleanup.py @@ -0,0 +1,43 @@ +"""Worker-scoped Arena registration cleanup regressions.""" + +from types import SimpleNamespace + +from areal.api.cli_args import InferenceEngineConfig + + +def test_remote_inf_engine_destroy_runs_each_registered_callback_once() -> None: + from areal.infra.remote_inf_engine import RemoteInfEngine + + engine = RemoteInfEngine( + InferenceEngineConfig(setup_timeout=60), + backend=SimpleNamespace(), + ) + calls = [] + engine.register_destroy_callback("arena", lambda: calls.append("first")) + engine.register_destroy_callback("arena", lambda: calls.append("duplicate")) + + engine.destroy() + engine.destroy() + + assert calls == ["first"] + + +def test_remote_inf_engine_destroy_callback_failure_does_not_mask_cleanup() -> None: + from areal.infra.remote_inf_engine import RemoteInfEngine + + engine = RemoteInfEngine( + InferenceEngineConfig(setup_timeout=60), + backend=SimpleNamespace(), + ) + calls = [] + + def fail() -> None: + calls.append("fail") + raise RuntimeError("cleanup failed") + + engine.register_destroy_callback("later", lambda: calls.append("later")) + engine.register_destroy_callback("fail", fail) + + engine.destroy() + + assert calls == ["fail", "later"] diff --git a/tests/test_config_redaction.py b/tests/test_config_redaction.py new file mode 100644 index 0000000000..fc7c58b097 --- /dev/null +++ b/tests/test_config_redaction.py @@ -0,0 +1,30 @@ +"""Tests for safe experiment configuration serialization.""" + +from areal.utils.config_utils import REDACTED_VALUE, redact_sensitive_config + + +def test_redact_sensitive_config_redacts_nested_credentials_only(): + """Credentials should be removed without hiding token-count configuration.""" + config = { + "admin_api_key": "admin-secret", + "arena": {"access_token": "arena-secret"}, + "workers": [{"password": "worker-secret", "max_tokens": 131072}], + "tokenizer_path": "/models/tokenizer", + } + + redacted = redact_sensitive_config(config) + + assert redacted == { + "admin_api_key": REDACTED_VALUE, + "arena": {"access_token": REDACTED_VALUE}, + "workers": [{"password": REDACTED_VALUE, "max_tokens": 131072}], + "tokenizer_path": "/models/tokenizer", + } + assert config["admin_api_key"] == "admin-secret" + + +def test_redact_sensitive_config_preserves_empty_optional_credentials(): + """Empty optional credential fields should remain empty for readability.""" + config = {"wandb_api_key": "", "refresh_token": None} + + assert redact_sensitive_config(config) == config diff --git a/tests/test_fsdp_microbatch_sync.py b/tests/test_fsdp_microbatch_sync.py new file mode 100644 index 0000000000..26563c9cfd --- /dev/null +++ b/tests/test_fsdp_microbatch_sync.py @@ -0,0 +1,103 @@ +# SPDX-License-Identifier: Apache-2.0 + +from datetime import timedelta +from types import SimpleNamespace + +import pytest +import torch +import torch.distributed as dist +import torch.multiprocessing as mp +from torch.nn.parallel import DistributedDataParallel + +from areal.api.cli_args import MicroBatchSpec +from areal.engine.fsdp_engine import FSDPEngine, FSDPTrainContext +from areal.utils.data import ( + pack_tensor_dict, + pad_mb_list, + split_padded_tensor_dict_into_mb_list, + unsqueeze_mb_list, +) + + +class _TinyModel(torch.nn.Module): + def __init__(self): + super().__init__() + self.weight = torch.nn.Parameter(torch.tensor(2.0)) + self.calls = 0 + + def forward(self, input_ids): + self.calls += 1 + return SimpleNamespace(logits=input_ids.float() * self.weight) + + +class _LoopHarness: + enable_tree_training = False + + def __init__(self, model): + self.model = model + self.cpu_group = dist.group.WORLD + + def _prepare_mb_inputs(self, item): + inputs = {"input_ids": item.padded_mb["input_ids"]} + return inputs, FSDPTrainContext(inputs, item.orig_mb) + + +def _run_uneven_microbatches(rank, rendezvous): + dist.init_process_group( + "gloo", + init_method=rendezvous, + rank=rank, + world_size=2, + timeout=timedelta(seconds=60), + ) + try: + # Rank 0 cannot be repartitioned into rank 1's three nonempty batches. + ids = torch.arange(1, (1 if rank == 0 else 3) + 1).reshape(-1, 1) + data = {"input_ids": ids, "attention_mask": torch.ones_like(ids)} + mb_list = split_padded_tensor_dict_into_mb_list( + data, MicroBatchSpec(n_mbs=1, max_tokens_per_mb=1), sync_mbs=False + ) + mb_list.mbs = [pack_tensor_dict(mb) for mb in mb_list.mbs] + mb_list = unsqueeze_mb_list(pad_mb_list(mb_list)) + model = DistributedDataParallel(_TinyModel()) + engine = _LoopHarness(model) + + for forward_only in (True, False): + outputs = [] + model.module.calls = 0 + model.zero_grad() + + def process_output(logits, ctx): + outputs.append(logits[0].detach()) + return logits.sum() + + with torch.set_grad_enabled(not forward_only): + FSDPEngine.forward_backward_batch( + engine, mb_list, process_output, forward_only=forward_only + ) + + assert model.module.calls == 3 + assert len(outputs) == len(ids) + torch.testing.assert_close( + torch.stack(outputs), ids.flatten().float() * 2, rtol=0, atol=0 + ) + if not forward_only: + # DDP average of real gradients: (1 + (1 + 2 + 3)) / 2. + torch.testing.assert_close( + model.module.weight.grad, torch.tensor(3.5), rtol=0, atol=0 + ) + assert len(mb_list) == len(ids) + finally: + dist.destroy_process_group() + + +@pytest.mark.slow +@pytest.mark.ci +def test_uneven_microbatches_preserve_outputs_and_gradients(tmp_path): + """Exercise the FSDP execution loop with real CPU gradient collectives.""" + mp.spawn( + _run_uneven_microbatches, + args=(f"file://{tmp_path / 'rendezvous'}",), + nprocs=2, + join=True, + ) diff --git a/tests/test_grouped_rollout_workflow.py b/tests/test_grouped_rollout_workflow.py index acbe43eca8..5189182624 100644 --- a/tests/test_grouped_rollout_workflow.py +++ b/tests/test_grouped_rollout_workflow.py @@ -204,3 +204,44 @@ async def _afinalize_processor_cache_group(self, context): finally: task.cancel() await asyncio.gather(task, return_exceptions=True) + + +@pytest.mark.asyncio +async def test_group_metrics_include_failed_slots_before_propagating_error(): + """Completed rewards survive a sibling error in attempted-group metrics.""" + completed = asyncio.Event() + metrics = [] + + class Workflow(RolloutWorkflow): + async def arun_episode(self, engine, data): + if workflow_context.get().sample_idx == 0: + completed.set() + return {"completion": _interaction(0.75)} + await completed.wait() + raise ValueError("failed slot") + + def record_group_metrics(self, data, rewards, group_size): + metrics.append((rewards, group_size)) + + workflow = GroupedRolloutWorkflow(Workflow(), group_size=2, logger=_Logger()) + with pytest.raises(ValueError, match="failed slot"): + await workflow.arun_episode(engine=None, data={}) + + assert metrics == [([0.75, None], 2)] + + +@pytest.mark.asyncio +async def test_group_metrics_observe_rewards_before_normalization(): + """Per-Stream pass rates must use terminal rewards, not normalized values.""" + metrics = [] + agent = _ListWorkflow([{"first": _interaction(0.0)}, {"second": _interaction(1.0)}]) + agent.record_group_metrics = lambda data, rewards, size: metrics.append(rewards) + workflow = GroupedRolloutWorkflow( + agent, group_size=2, logger=_Logger(), reward_normalization=True + ) + + result = await workflow.arun_episode(engine=None, data={}) + + assert metrics == [[0.0, 1.0]] + assert result["first"].reward == pytest.approx(-1.0) + assert result["second"].reward == pytest.approx(1.0) diff --git a/tests/test_network_ipv6_utils.py b/tests/test_network_ipv6_utils.py index 235ceba311..b593fb46b3 100644 --- a/tests/test_network_ipv6_utils.py +++ b/tests/test_network_ipv6_utils.py @@ -1,4 +1,5 @@ import importlib.util +import random from pathlib import Path import pytest @@ -96,3 +97,21 @@ def test_find_free_ports_ignores_out_of_range_excludes(): count=6, port_range=(10000, 10005), exclude_ports={50000} ) assert len(ports) == 6 + + +def test_find_free_ports_seeded_candidates_busy_uses_independent_rng(monkeypatch): + """Port allocation must neither replay nor advance the training RNG.""" + network = _load_network_module() + training_rng = random.Random(1) + state = training_rng.getstate() + occupied = {training_rng.randint(10000, 32767) for _ in range(10)} + training_rng.setstate(state) + free_port = next(port for port in range(10000, 32768) if port not in occupied) + monkeypatch.setattr(network.random, "randint", training_rng.randint) + monkeypatch.setattr( + network.random.SystemRandom, "randint", lambda self, lo, hi: free_port + ) + monkeypatch.setattr(network, "is_port_free", lambda port: port not in occupied) + + assert network.find_free_ports(1) == [free_port] + assert training_rng.getstate() == state diff --git a/tests/test_swe_arena.py b/tests/test_swe_arena.py new file mode 100644 index 0000000000..69f5b5ac27 --- /dev/null +++ b/tests/test_swe_arena.py @@ -0,0 +1,2159 @@ +"""Tests for the Arena Stream dataset and proxy agent integration.""" + +import asyncio +import json + +import httpx +import pytest + +from examples.swe.arena_agent import ArenaStreamAgentWorkflow +from examples.swe.arena_client import ( + ArenaAPIError, + ArenaOpenAPIClient, + ArenaTaskFailedError, + infer_llm_protocol, + resolve_llm_protocol, +) +from examples.swe.filter_function import filter_function + +from areal.experimental.openai.proxy.workflow import OpenAIProxyWorkflow +from areal.infra import workflow_context +from areal.infra.workflow_context import WorkflowContext +from areal.utils import stats_tracker + + +class FakeWorkerRuntime: + def __init__(self) -> None: + self.callbacks = {} + + def register_destroy_callback(self, key, callback) -> None: + self.callbacks.setdefault(key, callback) + + def destroy(self) -> None: + callbacks = list(self.callbacks.values()) + self.callbacks.clear() + for callback in callbacks: + callback() + + +def _reset_stats() -> None: + stats_tracker.export_all(reset=True) + + +def test_resolve_stream_id_when_unspecified_returns_first_active(monkeypatch): + """The first active Stream should be selected when no id is configured.""" + monkeypatch.setenv("ARENA_OPENAPI_TOKEN", "test-token") + + def handler(request: httpx.Request) -> httpx.Response: + assert request.url.params["status"] == "ACTIVE" + assert request.headers["Authorization"] == "Bearer test-token" + return httpx.Response( + 200, + json={ + "items": [ + {"stream_id": "stream-first", "status": "ACTIVE"}, + {"stream_id": "stream-second", "status": "ACTIVE"}, + ] + }, + ) + + client = ArenaOpenAPIClient(base_url="https://arena.example") + with httpx.Client(transport=httpx.MockTransport(handler)) as http_client: + stream_id = client.resolve_stream_id(client=http_client) + + assert stream_id == "stream-first" + + +def test_llm_gateway_api_key_is_required_only_for_gateway_traffic(monkeypatch): + """Dataset discovery may omit the key, but task launches must require it.""" + monkeypatch.setenv("ARENA_OPENAPI_TOKEN", "test-token") + monkeypatch.delenv("ARENA_LLM_API_KEY", raising=False) + client = ArenaOpenAPIClient(base_url="https://arena.example") + + with pytest.raises(ValueError, match="ARENA_LLM_API_KEY"): + _ = client.llm_gateway_api_key + + +def test_resolve_stream_when_id_is_explicit_returns_matching_metadata(monkeypatch): + """An explicit Stream should use direct get, independent of list pagination.""" + monkeypatch.setenv("ARENA_OPENAPI_TOKEN", "test-token") + + def handler(request: httpx.Request) -> httpx.Response: + assert request.url.path.endswith("/streams/stream-selected") + return httpx.Response( + 200, + json={ + "data": { + "stream": { + "stream_id": "stream-selected", + "status": "ACTIVE", + "default_harness_ref": {"key": "claude-code"}, + } + } + }, + ) + + client = ArenaOpenAPIClient(base_url="https://arena.example") + with httpx.Client(transport=httpx.MockTransport(handler)) as http_client: + stream = client.resolve_stream("stream-selected", client=http_client) + + assert stream["default_harness_ref"] == {"key": "claude-code"} + + +def test_resolve_stream_when_direct_get_returns_wrong_id(monkeypatch): + """A stale or misrouted direct response must not select another Stream.""" + monkeypatch.setenv("ARENA_OPENAPI_TOKEN", "test-token") + + def handler(_: httpx.Request) -> httpx.Response: + return httpx.Response(200, json={"stream_id": "stream-other"}) + + client = ArenaOpenAPIClient(base_url="https://arena.example") + with httpx.Client(transport=httpx.MockTransport(handler)) as http_client: + with pytest.raises(ArenaAPIError, match="returned id"): + client.resolve_stream("stream-selected", client=http_client) + + +def test_resolve_stream_async_when_id_is_explicit_uses_direct_get(monkeypatch): + """Rollout-time resolution should not depend on Stream-list pagination.""" + monkeypatch.setenv("ARENA_OPENAPI_TOKEN", "test-token") + + def handler(request: httpx.Request) -> httpx.Response: + assert request.url.path.endswith("/streams/stream-selected") + assert "status" not in request.url.params + return httpx.Response( + 200, + json={ + "stream_id": "stream-selected", + "status": "ACTIVE", + "default_harness_ref": {"key": "claude-code"}, + }, + ) + + async def resolve() -> dict: + transport = httpx.MockTransport(handler) + async with httpx.AsyncClient(transport=transport) as http_client: + client = ArenaOpenAPIClient(base_url="https://arena.example") + return await client.resolve_stream_async( + "stream-selected", + client=http_client, + timeout=10.0, + ) + + stream = asyncio.run(resolve()) + + assert stream["default_harness_ref"] == {"key": "claude-code"} + + +@pytest.mark.parametrize( + ("harness_key", "expected_protocol"), + [ + ("Claude-Code-With-Skills", "anthropic"), + ("openai-codex", "responses"), + ("swe-agent", "chat_completions"), + (None, "chat_completions"), + ], +) +def test_infer_llm_protocol_from_harness_key(harness_key, expected_protocol): + """Harness names should select their native protocol case-insensitively.""" + stream = ( + {"default_harness_ref": {"key": harness_key}} if harness_key is not None else {} + ) + + assert infer_llm_protocol(stream) == expected_protocol + + +def test_resolve_llm_protocol_override_bypasses_harness_inference(): + """Explicit protocol selection should bypass Harness-name inference.""" + stream = {"default_harness_ref": {"key": "claude-code"}} + + assert resolve_llm_protocol(stream, "chat_completions") == "chat_completions" + + +def test_resolve_llm_protocol_rejects_invalid_override(): + with pytest.raises(ValueError, match="Unsupported Arena LLM protocol override"): + resolve_llm_protocol({}, "invalid") + + +@pytest.mark.parametrize("failure", ["timeout", "gateway"]) +@pytest.mark.parametrize("owned_client", [False, True]) +def test_list_streams_retries_with_one_client(monkeypatch, failure, owned_client): + """Transient failures retry using the supplied client or one owned connection pool.""" + monkeypatch.setenv("ARENA_OPENAPI_TOKEN", "test-token") + real_client = httpx.Client + attempts = 0 + clients_created = 0 + + def handler(request: httpx.Request) -> httpx.Response: + nonlocal attempts + attempts += 1 + if attempts == 1: + if failure == "timeout": + raise httpx.ReadTimeout("transient", request=request) + return httpx.Response(504, text="gateway timeout") + return httpx.Response(200, json={"items": [{"stream_id": "stream-1"}]}) + + def create_client(**kwargs): + nonlocal clients_created + clients_created += 1 + return real_client(transport=httpx.MockTransport(handler), **kwargs) + + monkeypatch.setattr("examples.swe.arena_client.httpx.Client", create_client) + client = ArenaOpenAPIClient(base_url="https://arena.example", request_retries=1) + if owned_client: + streams = client.list_streams() + else: + with create_client() as http_client: + streams = client.list_streams(client=http_client) + + assert streams == [{"stream_id": "stream-1"}] + assert attempts == 2 + assert clients_created == 1 + + +def test_get_all_dataset_rows_below_api_limit_uses_one_data_page(monkeypatch): + """A small dataset should need one data page after the total probe.""" + monkeypatch.setenv("ARENA_OPENAPI_TOKEN", "test-token") + requested_limits: list[int] = [] + + def handler(request: httpx.Request) -> httpx.Response: + assert request.method == "POST" + assert request.url.path.endswith("/streams/stream-1/dataset") + limit = int(request.url.params["limit"]) + requested_limits.append(limit) + if limit == 1: + return httpx.Response( + 200, + json={ + "data_ids": ["data-1"], + "count": 1, + "total": 3, + "offset": 0, + "limit": 1, + }, + ) + return httpx.Response( + 200, + json={ + "data_ids": ["data-1", "data-2", "data-3"], + "items": [ + {"data_id": "data-1", "tags": ["domain:swe"]}, + {"data_id": "data-2", "tags": ["domain:math", "level:hard"]}, + {"data_id": "data-3", "tags": ["level:easy"]}, + ], + "count": 3, + "total": 3, + "offset": 0, + "limit": 3, + }, + ) + + client = ArenaOpenAPIClient(base_url="https://arena.example") + with httpx.Client(transport=httpx.MockTransport(handler)) as http_client: + rows = client.get_all_dataset_rows("stream-1", client=http_client) + + assert requested_limits == [1, 3] + assert rows == [ + { + "data_id": "data-1", + "stream_id": "stream-1", + "llm_protocol": "chat_completions", + "arena_task_type": "swe", + }, + { + "data_id": "data-2", + "stream_id": "stream-1", + "llm_protocol": "chat_completions", + "arena_task_type": "math", + }, + { + "data_id": "data-3", + "stream_id": "stream-1", + "llm_protocol": "chat_completions", + "arena_task_type": "unknown", + }, + ] + + +@pytest.mark.parametrize( + "tags", + [ + ["domain:math", "domain:swe"], + ["domain:swe", 1], + {"domain": "swe"}, + ], +) +def test_get_all_dataset_rows_domain_parse_failure_uses_unknown(monkeypatch, tags): + """Malformed or ambiguous domain metadata must not stop the experiment.""" + monkeypatch.setenv("ARENA_OPENAPI_TOKEN", "test-token") + + def handler(_: httpx.Request) -> httpx.Response: + return httpx.Response( + 200, + json={ + "data_ids": ["data-1"], + "items": [{"data_id": "data-1", "tags": tags}], + "count": 1, + "total": 1, + "offset": 0, + "limit": 1, + }, + ) + + client = ArenaOpenAPIClient(base_url="https://arena.example") + with httpx.Client(transport=httpx.MockTransport(handler)) as http_client: + rows = client.get_all_dataset_rows("stream-1", client=http_client) + + assert rows[0]["arena_task_type"] == "unknown" + + +def test_get_all_dataset_rows_over_api_limit_paginates(monkeypatch): + """Dataset loading should paginate Streams larger than the API page limit.""" + monkeypatch.setenv("ARENA_OPENAPI_TOKEN", "test-token") + requests: list[tuple[int, int]] = [] + + def handler(request: httpx.Request) -> httpx.Response: + limit = int(request.url.params["limit"]) + offset = int(request.url.params["offset"]) + requests.append((limit, offset)) + if limit == 1 and offset == 0 and len(requests) == 1: + data_ids = ["data-0"] + else: + data_ids = [ + f"data-{index}" for index in range(offset, min(offset + limit, 1001)) + ] + return httpx.Response( + 200, + json={ + "data_ids": data_ids, + "count": len(data_ids), + "total": 1001, + "offset": offset, + "limit": limit, + }, + ) + + client = ArenaOpenAPIClient(base_url="https://arena.example") + with httpx.Client(transport=httpx.MockTransport(handler)) as http_client: + rows = client.get_all_dataset_rows("stream-1", client=http_client) + + assert requests == [(1, 0), (1000, 0), (1, 1000)] + assert len(rows) == 1001 + assert rows[0]["data_id"] == "data-0" + assert rows[-1]["data_id"] == "data-1000" + + +def test_get_all_dataset_rows_rejects_total_drift_between_pages(monkeypatch): + """A changing Stream must not produce a mixed pagination snapshot.""" + monkeypatch.setenv("ARENA_OPENAPI_TOKEN", "test-token") + requests: list[tuple[int, int]] = [] + + def handler(request: httpx.Request) -> httpx.Response: + limit = int(request.url.params["limit"]) + offset = int(request.url.params["offset"]) + requests.append((limit, offset)) + is_probe = len(requests) == 1 + page_total = 1001 if is_probe or offset == 0 else 1002 + data_ids = ( + ["data-0"] + if is_probe + else [f"data-{index}" for index in range(offset, min(offset + limit, 1001))] + ) + return httpx.Response( + 200, + json={ + "data_ids": data_ids, + "count": len(data_ids), + "total": page_total, + "offset": offset, + "limit": limit, + }, + ) + + client = ArenaOpenAPIClient(base_url="https://arena.example") + with httpx.Client(transport=httpx.MockTransport(handler)) as http_client: + with pytest.raises(ArenaAPIError, match="total changed during pagination"): + client.get_all_dataset_rows("stream-1", client=http_client) + + +def test_get_all_dataset_rows_rejects_duplicate_ids_between_pages(monkeypatch): + """Offset pagination must not silently repeat a row across pages.""" + monkeypatch.setenv("ARENA_OPENAPI_TOKEN", "test-token") + requests = 0 + + def handler(request: httpx.Request) -> httpx.Response: + nonlocal requests + requests += 1 + limit = int(request.url.params["limit"]) + offset = int(request.url.params["offset"]) + if requests == 1: + data_ids = ["data-0"] + elif offset == 0: + data_ids = [f"data-{index}" for index in range(1000)] + else: + data_ids = ["data-999"] + return httpx.Response( + 200, + json={ + "data_ids": data_ids, + "count": len(data_ids), + "total": 1001, + "offset": offset, + "limit": limit, + }, + ) + + client = ArenaOpenAPIClient(base_url="https://arena.example") + with httpx.Client(transport=httpx.MockTransport(handler)) as http_client: + with pytest.raises(ArenaAPIError, match="pagination returned duplicate ids"): + client.get_all_dataset_rows("stream-1", client=http_client) + + +@pytest.mark.parametrize( + ("field", "wrong_value", "message"), + [ + ("offset", 999, "pagination returned offset"), + ("limit", 999, "pagination returned limit"), + ("count", 999, "pagination returned count"), + ], +) +def test_get_all_dataset_rows_rejects_inconsistent_page_metadata( + monkeypatch, field, wrong_value, message +): + """Incorrect page metadata must not silently produce a mixed dataset.""" + monkeypatch.setenv("ARENA_OPENAPI_TOKEN", "test-token") + requests = 0 + + def handler(request: httpx.Request) -> httpx.Response: + nonlocal requests + requests += 1 + limit = int(request.url.params["limit"]) + offset = int(request.url.params["offset"]) + if requests == 1: + data_ids = ["data-0"] + else: + data_ids = [ + f"data-{index}" for index in range(offset, min(offset + limit, 2)) + ] + payload = { + "data_ids": data_ids, + "count": len(data_ids), + "total": 2, + "offset": offset, + "limit": limit, + } + if requests > 1: + payload[field] = wrong_value + return httpx.Response(200, json=payload) + + client = ArenaOpenAPIClient(base_url="https://arena.example") + with httpx.Client(transport=httpx.MockTransport(handler)) as http_client: + with pytest.raises(ArenaAPIError, match=message): + client.get_all_dataset_rows("stream-1", client=http_client) + + +def test_register_and_delete_llm_proxy(monkeypatch): + """OpenAPI registry calls should forward one model and endpoint.""" + monkeypatch.setenv("ARENA_OPENAPI_TOKEN", "arena-token") + monkeypatch.setenv("ARENA_LLM_API_KEY", "test-llm-key") + deployment_id = "deployment-1" + requests: list[httpx.Request] = [] + + def handler(request: httpx.Request) -> httpx.Response: + requests.append(request) + assert request.headers["Authorization"] == "Bearer arena-token" + if request.method == "POST": + assert request.url.path.endswith("/openapi/v1/llm/models") + payload = json.loads(request.content) + assert payload == { + "model_name": "stream-areal-test-1", + "endpoints": [ + { + "endpoint_id": deployment_id, + "upstream_model": "stream-areal-test-1", + "base_url": "http://rollout-proxy", + "api_key": "session-key", + "inbound_protos": ["chat"], + "enabled": True, + } + ], + "enabled": True, + "metadata": {"deployment_id": deployment_id}, + } + return httpx.Response( + 200, + json={"model_name": "stream-areal-test-1"}, + ) + assert request.method == "DELETE" + assert request.url.path.endswith("/openapi/v1/llm/models/stream-areal-test-1") + return httpx.Response(204) + + client = ArenaOpenAPIClient(base_url="https://arena.example") + with httpx.Client(transport=httpx.MockTransport(handler)) as http_client: + registered_url, returned_id = client.register_llm_proxy( + model_name="stream-areal-test-1", + upstream_base_url="http://rollout-proxy", + upstream_api_key="session-key", + deployment_id=deployment_id, + client=http_client, + ) + client.delete_llm_proxy(returned_id, client=http_client) + + assert registered_url == "https://arena.example/api" + assert returned_id == "stream-areal-test-1" + assert len(requests) == 2 + + +def test_probe_llm_proxy_distinguishes_present_and_garbage_collected(monkeypatch): + """A precise model GET should expose Arena GC as a clean false result.""" + monkeypatch.setenv("ARENA_OPENAPI_TOKEN", "arena-token") + + async def handler(request: httpx.Request) -> httpx.Response: + model_name = request.url.path.rsplit("/", 1)[-1] + if model_name == "stream-areal-present": + return httpx.Response(200, json={"model_name": model_name}) + if model_name == "stream-areal-collected": + return httpx.Response(404, json={"detail": "not found"}) + return httpx.Response(200, json={"model_name": "stream-areal-wrong"}) + + async def probe() -> None: + transport = httpx.MockTransport(handler) + async with httpx.AsyncClient(transport=transport) as http_client: + client = ArenaOpenAPIClient(base_url="https://arena.example") + assert await client.llm_proxy_exists_async( + "stream-areal-present", client=http_client + ) + assert not await client.llm_proxy_exists_async( + "stream-areal-collected", client=http_client + ) + with pytest.raises(ArenaAPIError, match="unexpected model_name"): + await client.llm_proxy_exists_async( + "stream-areal-mismatch", client=http_client + ) + + asyncio.run(probe()) + + +@pytest.mark.parametrize( + "protocol", + ["anthropic", "responses", "chat_completions"], +) +def test_register_llm_proxy_always_advertises_openai_chat(monkeypatch, protocol): + """All Harness protocols should use the AReaL proxy's native Chat API.""" + monkeypatch.setenv("ARENA_OPENAPI_TOKEN", "arena-token") + monkeypatch.setenv("ARENA_LLM_API_KEY", "test-llm-key") + + def handler(request: httpx.Request) -> httpx.Response: + assert request.headers["Authorization"] == "Bearer arena-token" + payload = json.loads(request.content) + assert payload["endpoints"] == [ + { + "endpoint_id": "deployment-1", + "upstream_model": "stream-areal-test-1", + "base_url": "http://rollout-proxy", + "api_key": "session-key", + "inbound_protos": ["chat"], + "enabled": True, + } + ] + return httpx.Response(200, json={"model_name": "stream-areal-test-1"}) + + client = ArenaOpenAPIClient(base_url="https://arena.example") + with httpx.Client(transport=httpx.MockTransport(handler)) as http_client: + client.register_llm_proxy( + model_name="stream-areal-test-1", + upstream_base_url="http://rollout-proxy", + upstream_api_key="session-key", + deployment_id="deployment-1", + protocol=protocol, + client=http_client, + ) + + +def test_agent_launches_task_through_proxy_and_returns_reward(monkeypatch): + """The Arena agent should forward proxy credentials and return task reward.""" + monkeypatch.setenv("ARENA_OPENAPI_TOKEN", "arena-token") + monkeypatch.setenv("ARENA_LLM_API_KEY", "test-llm-key") + _reset_stats() + + result_polls = 0 + registered_model_name = "" + + def handler(request: httpx.Request) -> httpx.Response: + nonlocal registered_model_name, result_polls + if request.method == "POST" and request.url.path.endswith( + "/openapi/v1/llm/models" + ): + assert request.headers["Authorization"] == "Bearer arena-token" + payload = json.loads(request.content) + registered_model_name = payload["model_name"] + assert registered_model_name.startswith("stream-areal-") + assert payload["endpoints"] == [ + { + "endpoint_id": payload["metadata"]["deployment_id"], + "upstream_model": registered_model_name, + "base_url": "http://rollout-proxy", + "api_key": "session-key", + "inbound_protos": ["chat"], + "enabled": True, + } + ] + return httpx.Response(200, json={"model_name": registered_model_name}) + if request.url.path.endswith("/streams/stream-1/launch_one_task"): + assert request.headers["Authorization"] == "Bearer arena-token" + assert json.loads(request.content) == { + "data_id": "data-1", + "model_name": registered_model_name, + "base_url": "https://arena.example/api", + "api_key": "test-llm-key", + "harness": "claude-code-with-skills@5.0.1", + "envs": { + "MODEL_NAME": registered_model_name, + "BASE_URL": "https://arena.example/api", + "API_KEY": "test-llm-key", + "CLAUDE_CODE_DISABLE_TERMINAL_TITLE": "1", + }, + } + return httpx.Response( + 202, + json={"accepted": True, "task_id": "task-1", "status": "PENDING"}, + ) + if request.method == "DELETE": + assert request.headers["Authorization"] == "Bearer arena-token" + assert request.url.path.endswith( + f"/openapi/v1/llm/models/{registered_model_name}" + ) + return httpx.Response(204) + assert request.headers["Authorization"] == "Bearer arena-token" + assert request.url.path.endswith("/tasks/task-1/result") + result_polls += 1 + if result_polls == 1: + return httpx.Response(200, json={"status": "AGENT_RUNNING"}) + return httpx.Response(200, json={"status": "OK", "score": 0.75}) + + async def run_agent() -> float: + transport = httpx.MockTransport(handler) + async with httpx.AsyncClient(transport=transport) as http_client: + workflow = ArenaStreamAgentWorkflow( + econfig={ + "arena_base_url": "https://arena.example", + "arena_poll_interval": 0.0, + "arena_harness": "claude-code-with-skills@5.0.1", + "arena_task_envs": {"CLAUDE_CODE_DISABLE_TERMINAL_TITLE": "1"}, + "timeout": 10.0, + } + ) + return await workflow.run( + { + "stream_id": "stream-1", + "data_id": "data-1", + "llm_protocol": "anthropic", + "arena_task_type": "swe", + }, + base_url="http://rollout-proxy", + api_key="session-key", + arena_http_client=http_client, + ) + + reward = asyncio.run(run_agent()) + assert reward == 0.75 + assert result_polls == 2 + stats = stats_tracker.export_all(reset=True) + assert stats["rollout/arena/call_success"] == 1.0 + assert stats["rollout/arena/registration_success"] == 1.0 + assert stats["rollout/arena/launch_success"] == 1.0 + assert stats["rollout/arena/terminal_success"] == 1.0 + assert "rollout/reward" not in stats + assert "rollout/swe/reward" not in stats + assert stats["rollout/arena/cleanup_success"] == 1.0 + assert stats["rollout/arena/call_success__count"] == 1 + + +def test_agent_direct_route_uses_session_credentials_without_registration(monkeypatch): + """Direct mode should preserve per-session isolation without registry churn.""" + monkeypatch.setenv("ARENA_OPENAPI_TOKEN", "arena-token") + monkeypatch.delenv("ARENA_LLM_API_KEY", raising=False) + _reset_stats() + requests: list[httpx.Request] = [] + + def handler(request: httpx.Request) -> httpx.Response: + requests.append(request) + assert request.url.path.endswith("/streams/stream-1/launch_one_task") + payload = json.loads(request.content) + assert payload["session_id"] == "proxy-session-1" + assert payload["base_url"] == "http://rollout-proxy" + assert payload["api_key"] == "session-key" + assert payload["model_name"].startswith("stream-areal-") + assert payload["envs"] == { + "MODEL_NAME": payload["model_name"], + "BASE_URL": "http://rollout-proxy", + "API_KEY": "session-key", + "ANTHROPIC_BASE_URL": "http://rollout-proxy", + "ANTHROPIC_API_KEY": "session-key", + "OPENAI_BASE_URL": "http://rollout-proxy", + "OPENAI_API_KEY": "session-key", + "CLAUDE_CODE_DISABLE_TERMINAL_TITLE": "1", + "NO_PROXY": "rollout-proxy", + "no_proxy": "rollout-proxy", + } + return httpx.Response( + 200, + json={"task_id": "task-direct", "status": "DONE", "score": 0.25}, + ) + + async def run_agent() -> float: + transport = httpx.MockTransport(handler) + async with httpx.AsyncClient(transport=transport) as http_client: + workflow = ArenaStreamAgentWorkflow( + econfig={ + "arena_base_url": "https://arena.example", + "arena_llm_route_mode": "direct", + "arena_harness": "claude-code@1", + "arena_task_envs": { + "ANTHROPIC_BASE_URL": "https://arena.example/api", + "CLAUDE_CODE_DISABLE_TERMINAL_TITLE": "1", + }, + } + ) + return await workflow.run( + { + "stream_id": "stream-1", + "data_id": "data-1", + "llm_protocol": "anthropic", + }, + base_url="http://rollout-proxy", + api_key="session-key", + session_id="proxy-session-1", + arena_http_client=http_client, + ) + + assert asyncio.run(run_agent()) == pytest.approx(0.25) + assert len(requests) == 1 + stats = stats_tracker.export_all(reset=True) + assert stats["rollout/arena/direct_route"] == 1.0 + assert stats["rollout/arena/call_success"] == 1.0 + assert "rollout/arena/registration_success" not in stats + assert "rollout/arena/cleanup_success" not in stats + + +def test_agent_direct_route_preserves_configured_no_proxy_hosts(monkeypatch): + monkeypatch.setenv("ARENA_OPENAPI_TOKEN", "arena-token") + + def handler(request: httpx.Request) -> httpx.Response: + payload = json.loads(request.content) + assert payload["envs"]["NO_PROXY"] == "localhost,rollout-proxy" + assert payload["envs"]["no_proxy"] == "localhost,rollout-proxy" + return httpx.Response( + 200, + json={"task_id": "task-direct", "status": "DONE", "score": 0.0}, + ) + + async def run_agent() -> float: + transport = httpx.MockTransport(handler) + async with httpx.AsyncClient(transport=transport) as http_client: + workflow = ArenaStreamAgentWorkflow( + econfig={ + "arena_base_url": "https://arena.example", + "arena_llm_route_mode": "direct", + "arena_task_envs": {"NO_PROXY": "localhost"}, + } + ) + return await workflow.run( + {"stream_id": "stream-1", "data_id": "data-1"}, + base_url="http://rollout-proxy:1234", + api_key="session-key", + session_id="proxy-session-1", + arena_http_client=http_client, + ) + + assert asyncio.run(run_agent()) == 0.0 + + +def test_agent_direct_route_requires_proxy_session_id(monkeypatch): + monkeypatch.setenv("ARENA_OPENAPI_TOKEN", "arena-token") + workflow = ArenaStreamAgentWorkflow( + econfig={ + "arena_base_url": "https://arena.example", + "arena_llm_route_mode": "direct", + } + ) + + with pytest.raises(ValueError, match="session_id is required"): + asyncio.run( + workflow.run( + {"stream_id": "stream-1", "data_id": "data-1"}, + base_url="http://rollout-proxy", + api_key="session-key", + ) + ) + + +def test_agent_session_gateway_reuses_worker_registration_and_binds_sessions( + monkeypatch, +): + """Concurrent sessions share one route while retaining explicit isolation.""" + monkeypatch.setenv("ARENA_OPENAPI_TOKEN", "arena-token") + monkeypatch.setenv("ARENA_LLM_API_KEY", "arena-gateway-key") + _reset_stats() + registration_count = 0 + deletion_count = 0 + launched_sessions: set[str] = set() + both_launched = asyncio.Event() + + async def handler(request: httpx.Request) -> httpx.Response: + nonlocal registration_count, deletion_count + if request.method == "POST" and request.url.path.endswith( + "/openapi/v1/llm/models" + ): + registration_count += 1 + payload = json.loads(request.content) + assert payload["model_name"].startswith("stream-areal-session-") + assert payload["endpoints"] == [ + { + "endpoint_id": payload["metadata"]["deployment_id"], + "upstream_model": payload["model_name"], + "base_url": "http://rollout-proxy", + "api_key": "proxy-gateway-key", + "inbound_protos": ["chat"], + "enabled": True, + } + ] + return httpx.Response(201, json={"model_name": payload["model_name"]}) + if request.method == "POST" and request.url.path.endswith( + "/streams/stream-1/launch_one_task" + ): + payload = json.loads(request.content) + session_id = payload["session_id"] + launched_sessions.add(session_id) + assert payload["base_url"] == "https://arena.example/api" + assert payload["api_key"] == "arena-gateway-key" + assert payload["envs"]["ANTHROPIC_CUSTOM_HEADERS"] == ( + f"x-keepalive-enable: true\nX-Session-Id: {session_id}\n" + f"X-Session-Token: token-{session_id}" + ) + if len(launched_sessions) == 2: + both_launched.set() + await asyncio.wait_for(both_launched.wait(), timeout=1.0) + return httpx.Response( + 200, + json={ + "task_id": f"task-{session_id}", + "status": "DONE", + "score": 0.5, + }, + ) + if request.method == "DELETE": + deletion_count += 1 + return httpx.Response(204) + raise AssertionError(f"Unexpected request: {request.method} {request.url}") + + async def run_agents() -> list[float]: + transport = httpx.MockTransport(handler) + async with httpx.AsyncClient(transport=transport) as http_client: + workflow = ArenaStreamAgentWorkflow( + econfig={ + "arena_base_url": "https://arena.example", + "arena_llm_route_mode": "session_gateway", + "arena_task_envs": { + "ANTHROPIC_CUSTOM_HEADERS": "x-keepalive-enable: true" + }, + } + ) + + async def run_one(session_id: str, data_id: str) -> float: + return await workflow.run( + {"stream_id": "stream-1", "data_id": data_id}, + base_url="http://rollout-proxy", + api_key=f"key-{session_id}", + proxy_gateway_api_key="proxy-gateway-key", + proxy_session_token=f"token-{session_id}", + session_id=session_id, + arena_http_client=http_client, + ) + + return await asyncio.gather( + run_one("proxy-session-1", "data-1"), + run_one("proxy-session-2", "data-2"), + ) + + assert asyncio.run(run_agents()) == [0.5, 0.5] + assert launched_sessions == {"proxy-session-1", "proxy-session-2"} + assert registration_count == 1 + assert deletion_count == 1 + stats = stats_tracker.export_all(reset=True) + assert stats["rollout/arena/session_gateway_route"] == 1.0 + assert stats["rollout/arena/session_gateway_route__count"] == 2 + + +def test_agent_session_gateway_registration_lives_across_prompt_workflows(monkeypatch): + """Separate prompt workflows on one worker must reuse one Arena model.""" + monkeypatch.setenv("ARENA_OPENAPI_TOKEN", "arena-token") + monkeypatch.setenv("ARENA_LLM_API_KEY", "arena-gateway-key") + _reset_stats() + registrations: list[str] = [] + async_deletions: list[str] = [] + cleanup_deletions: list[str] = [] + launches: list[tuple[str, str, str]] = [] + + async def handler(request: httpx.Request) -> httpx.Response: + if request.method == "POST" and request.url.path.endswith( + "/openapi/v1/llm/models" + ): + payload = json.loads(request.content) + registrations.append(payload["model_name"]) + return httpx.Response(201, json={"model_name": payload["model_name"]}) + if request.method == "POST" and request.url.path.endswith( + "/streams/stream-1/launch_one_task" + ): + payload = json.loads(request.content) + launches.append( + ( + payload["model_name"], + payload["session_id"], + payload["envs"]["ANTHROPIC_CUSTOM_HEADERS"], + ) + ) + return httpx.Response( + 200, + json={ + "task_id": f"task-{payload['session_id']}", + "status": "DONE", + "score": 0.5, + }, + ) + if request.method == "DELETE": + async_deletions.append(request.url.path) + return httpx.Response(204) + raise AssertionError(f"Unexpected request: {request.method} {request.url}") + + worker_runtime = FakeWorkerRuntime() + monkeypatch.setattr( + ArenaOpenAPIClient, + "delete_llm_proxy", + lambda self, model_id, **kwargs: cleanup_deletions.append(model_id), + ) + + async def run_prompts() -> list[float]: + transport = httpx.MockTransport(handler) + async with httpx.AsyncClient(transport=transport) as http_client: + workflows = [ + ArenaStreamAgentWorkflow( + econfig={ + "arena_base_url": "https://arena.example", + "arena_llm_route_mode": "session_gateway", + } + ) + for _ in range(2) + ] + rewards = [] + for index, workflow in enumerate(workflows, start=1): + session_id = f"proxy-session-{index}" + rewards.append( + await workflow.run( + { + "stream_id": "stream-1", + "data_id": f"data-{index}", + }, + base_url="http://rollout-proxy", + api_key=f"session-key-{index}", + proxy_gateway_api_key="proxy-gateway-key", + proxy_session_token=f"token-{index}", + session_id=session_id, + worker_runtime=worker_runtime, + arena_http_client=http_client, + ) + ) + return rewards + + assert asyncio.run(run_prompts()) == [0.5, 0.5] + assert len(registrations) == 1 + assert async_deletions == [] + assert [launch[0] for launch in launches] == registrations * 2 + assert [launch[1] for launch in launches] == [ + "proxy-session-1", + "proxy-session-2", + ] + assert "X-Session-Token: token-1" in launches[0][2] + assert "X-Session-Token: token-2" in launches[1][2] + assert len(worker_runtime.callbacks) == 1 + + worker_runtime.destroy() + assert cleanup_deletions == registrations + + +def test_agent_session_gateway_restores_gc_route_before_next_task(monkeypatch): + """A stale cached route should be probed and restored under its original name.""" + monkeypatch.setenv("ARENA_OPENAPI_TOKEN", "arena-token") + monkeypatch.setenv("ARENA_LLM_API_KEY", "arena-gateway-key") + _reset_stats() + requests: list[tuple[str, str]] = [] + registrations: list[str] = [] + route_exists = True + + async def handler(request: httpx.Request) -> httpx.Response: + nonlocal route_exists + if request.method == "POST" and request.url.path.endswith( + "/openapi/v1/llm/models" + ): + model_name = json.loads(request.content)["model_name"] + registrations.append(model_name) + requests.append(("register", model_name)) + route_exists = True + return httpx.Response(201, json={"model_name": model_name}) + if request.method == "GET" and "/openapi/v1/llm/models/" in request.url.path: + model_name = request.url.path.rsplit("/", 1)[-1] + requests.append(("probe", model_name)) + if route_exists: + return httpx.Response(200, json={"model_name": model_name}) + return httpx.Response(404, json={"detail": "not found"}) + if request.method == "POST" and request.url.path.endswith( + "/streams/stream-1/launch_one_task" + ): + model_name = json.loads(request.content)["model_name"] + requests.append(("launch", model_name)) + return httpx.Response( + 200, + json={"task_id": "task-1", "status": "DONE", "score": 0.5}, + ) + raise AssertionError(f"Unexpected request: {request.method} {request.url}") + + async def run_tasks() -> None: + nonlocal route_exists + transport = httpx.MockTransport(handler) + async with httpx.AsyncClient(transport=transport) as http_client: + workflow = ArenaStreamAgentWorkflow( + econfig={ + "arena_base_url": "https://arena.example", + "arena_llm_route_mode": "session_gateway", + "arena_registration_probe_interval": 0.01, + } + ) + worker_runtime = FakeWorkerRuntime() + + async def run_one(index: int) -> float: + return await workflow.run( + {"stream_id": "stream-1", "data_id": f"data-{index}"}, + base_url="http://rollout-proxy", + api_key=f"session-key-{index}", + proxy_gateway_api_key="proxy-gateway-key", + proxy_session_token=f"token-{index}", + session_id=f"proxy-session-{index}", + worker_runtime=worker_runtime, + arena_http_client=http_client, + ) + + assert await run_one(1) == 0.5 + route_exists = False + await asyncio.sleep(0.02) + assert await run_one(2) == 0.5 + + asyncio.run(run_tasks()) + + assert len(registrations) == 2 + assert registrations[0] == registrations[1] + assert requests[-3:] == [ + ("probe", registrations[0]), + ("register", registrations[0]), + ("launch", registrations[0]), + ] + + +def test_agent_session_gateway_throttles_concurrent_prelaunch_probes(monkeypatch): + """Concurrent task acquisition should singleflight one probe per worker route.""" + monkeypatch.setenv("ARENA_OPENAPI_TOKEN", "arena-token") + monkeypatch.setenv("ARENA_LLM_API_KEY", "arena-gateway-key") + _reset_stats() + registration_count = 0 + probe_count = 0 + + async def handler(request: httpx.Request) -> httpx.Response: + nonlocal registration_count, probe_count + if request.method == "POST" and request.url.path.endswith( + "/openapi/v1/llm/models" + ): + registration_count += 1 + model_name = json.loads(request.content)["model_name"] + return httpx.Response(201, json={"model_name": model_name}) + if request.method == "GET" and "/openapi/v1/llm/models/" in request.url.path: + probe_count += 1 + await asyncio.sleep(0.02) + model_name = request.url.path.rsplit("/", 1)[-1] + return httpx.Response(200, json={"model_name": model_name}) + if request.method == "POST" and request.url.path.endswith( + "/streams/stream-1/launch_one_task" + ): + payload = json.loads(request.content) + return httpx.Response( + 200, + json={ + "task_id": f"task-{payload['session_id']}", + "status": "DONE", + "score": 0.5, + }, + ) + raise AssertionError(f"Unexpected request: {request.method} {request.url}") + + async def run_tasks() -> None: + transport = httpx.MockTransport(handler) + async with httpx.AsyncClient(transport=transport) as http_client: + workflow = ArenaStreamAgentWorkflow( + econfig={ + "arena_base_url": "https://arena.example", + "arena_llm_route_mode": "session_gateway", + "arena_registration_probe_interval": 1.0, + } + ) + worker_runtime = FakeWorkerRuntime() + + async def run_one(index: int) -> float: + return await workflow.run( + {"stream_id": "stream-1", "data_id": f"data-{index}"}, + base_url="http://rollout-proxy", + api_key=f"session-key-{index}", + proxy_gateway_api_key="proxy-gateway-key", + proxy_session_token=f"token-{index}", + session_id=f"proxy-session-{index}", + worker_runtime=worker_runtime, + arena_http_client=http_client, + ) + + assert await run_one(0) == 0.5 + registry = getattr(worker_runtime, "_arena_session_gateway_registry_v1") + entry = next(iter(registry._targets.values())) + entry.last_probe_at -= 2.0 + rewards = await asyncio.gather(*(run_one(i) for i in range(1, 33))) + assert rewards == [0.5] * 32 + + asyncio.run(run_tasks()) + + assert registration_count == 1 + assert probe_count == 1 + + +def test_agent_session_gateway_throttles_slow_probe_failures(monkeypatch): + """A slow failed probe should throttle from completion, not request start.""" + monkeypatch.setenv("ARENA_OPENAPI_TOKEN", "arena-token") + monkeypatch.setenv("ARENA_LLM_API_KEY", "arena-gateway-key") + _reset_stats() + probe_count = 0 + + async def handler(request: httpx.Request) -> httpx.Response: + if request.method == "POST" and request.url.path.endswith( + "/openapi/v1/llm/models" + ): + model_name = json.loads(request.content)["model_name"] + return httpx.Response(201, json={"model_name": model_name}) + if request.method == "POST" and request.url.path.endswith( + "/streams/stream-1/launch_one_task" + ): + payload = json.loads(request.content) + return httpx.Response( + 200, + json={ + "task_id": f"task-{payload['session_id']}", + "status": "DONE", + "score": 0.5, + }, + ) + raise AssertionError(f"Unexpected request: {request.method} {request.url}") + + async def run_tasks() -> None: + nonlocal probe_count + transport = httpx.MockTransport(handler) + async with httpx.AsyncClient(transport=transport) as http_client: + workflow = ArenaStreamAgentWorkflow( + econfig={ + "arena_base_url": "https://arena.example", + "arena_llm_route_mode": "session_gateway", + "arena_registration_probe_interval": 0.05, + } + ) + worker_runtime = FakeWorkerRuntime() + + async def run_one(index: int) -> float: + return await workflow.run( + {"stream_id": "stream-1", "data_id": f"data-{index}"}, + base_url="http://rollout-proxy", + api_key=f"session-key-{index}", + proxy_gateway_api_key="proxy-gateway-key", + proxy_session_token=f"token-{index}", + session_id=f"proxy-session-{index}", + worker_runtime=worker_runtime, + arena_http_client=http_client, + ) + + assert await run_one(0) == 0.5 + registry = getattr(worker_runtime, "_arena_session_gateway_registry_v1") + entry = next(iter(registry._targets.values())) + entry.last_probe_at -= 1.0 + + async def fail_slow_probe(*args, **kwargs) -> bool: + nonlocal probe_count + probe_count += 1 + await asyncio.sleep(0.06) + raise ArenaAPIError("temporary probe failure") + + monkeypatch.setattr( + workflow.client, "llm_proxy_exists_async", fail_slow_probe + ) + assert await run_one(1) == 0.5 + assert await run_one(2) == 0.5 + + asyncio.run(run_tasks()) + + assert probe_count == 1 + + +def test_agent_session_gateway_does_not_launch_known_missing_route(monkeypatch): + """A failed restore must reject cached acquires instead of launching a 404 route.""" + monkeypatch.setenv("ARENA_OPENAPI_TOKEN", "arena-token") + monkeypatch.setenv("ARENA_LLM_API_KEY", "arena-gateway-key") + _reset_stats() + launch_count = 0 + probe_count = 0 + restore_count = 0 + + async def handler(request: httpx.Request) -> httpx.Response: + nonlocal launch_count + if request.method == "POST" and request.url.path.endswith( + "/openapi/v1/llm/models" + ): + model_name = json.loads(request.content)["model_name"] + return httpx.Response(201, json={"model_name": model_name}) + if request.method == "POST" and request.url.path.endswith( + "/streams/stream-1/launch_one_task" + ): + launch_count += 1 + payload = json.loads(request.content) + return httpx.Response( + 200, + json={ + "task_id": f"task-{payload['session_id']}", + "status": "DONE", + "score": 0.5, + }, + ) + raise AssertionError(f"Unexpected request: {request.method} {request.url}") + + async def run_tasks() -> None: + nonlocal probe_count, restore_count + transport = httpx.MockTransport(handler) + async with httpx.AsyncClient(transport=transport) as http_client: + workflow = ArenaStreamAgentWorkflow( + econfig={ + "arena_base_url": "https://arena.example", + "arena_llm_route_mode": "session_gateway", + "arena_registration_probe_interval": 60.0, + } + ) + worker_runtime = FakeWorkerRuntime() + + async def run_one(index: int) -> float: + return await workflow.run( + {"stream_id": "stream-1", "data_id": f"data-{index}"}, + base_url="http://rollout-proxy", + api_key=f"session-key-{index}", + proxy_gateway_api_key="proxy-gateway-key", + proxy_session_token=f"token-{index}", + session_id=f"proxy-session-{index}", + worker_runtime=worker_runtime, + arena_http_client=http_client, + ) + + assert await run_one(0) == 0.5 + registry = getattr(worker_runtime, "_arena_session_gateway_registry_v1") + entry = next(iter(registry._targets.values())) + entry.last_probe_at -= 61.0 + + async def missing_route(*args, **kwargs) -> bool: + nonlocal probe_count + probe_count += 1 + return False + + async def fail_restore(*args, **kwargs): + nonlocal restore_count + restore_count += 1 + raise ArenaAPIError("temporary restore failure") + + monkeypatch.setattr( + workflow.client, "llm_proxy_exists_async", missing_route + ) + monkeypatch.setattr( + workflow.client, "register_llm_proxy_async", fail_restore + ) + with pytest.raises(ArenaAPIError, match="temporary restore failure"): + await run_one(1) + with pytest.raises(ArenaAPIError, match="restore retries are throttled"): + await run_one(2) + + asyncio.run(run_tasks()) + + assert launch_count == 1 + assert probe_count == 2 + assert restore_count == 1 + + +@pytest.mark.parametrize("probe_interval", [0.0, -1.0, float("nan"), float("inf")]) +def test_agent_session_gateway_rejects_invalid_probe_interval(probe_interval): + """Invalid intervals must not disable probes or create a busy loop.""" + with pytest.raises(ValueError, match="finite and positive"): + ArenaStreamAgentWorkflow( + econfig={"arena_registration_probe_interval": probe_interval} + ) + + +def test_agent_session_gateway_restores_gc_route_while_task_is_queued(monkeypatch): + """The background probe should protect tasks waiting in Arena's queue.""" + monkeypatch.setenv("ARENA_OPENAPI_TOKEN", "arena-token") + monkeypatch.setenv("ARENA_LLM_API_KEY", "arena-gateway-key") + _reset_stats() + registrations: list[str] = [] + restored = asyncio.Event() + route_exists = True + + async def handler(request: httpx.Request) -> httpx.Response: + nonlocal route_exists + if request.method == "POST" and request.url.path.endswith( + "/openapi/v1/llm/models" + ): + model_name = json.loads(request.content)["model_name"] + registrations.append(model_name) + route_exists = True + if len(registrations) == 2: + restored.set() + return httpx.Response(201, json={"model_name": model_name}) + if request.method == "GET" and "/openapi/v1/llm/models/" in request.url.path: + model_name = request.url.path.rsplit("/", 1)[-1] + if route_exists: + return httpx.Response(200, json={"model_name": model_name}) + return httpx.Response(404, json={"detail": "not found"}) + if request.method == "POST" and request.url.path.endswith( + "/streams/stream-1/launch_one_task" + ): + route_exists = False + await asyncio.wait_for(restored.wait(), timeout=1.0) + return httpx.Response( + 200, + json={"task_id": "task-queued", "status": "DONE", "score": 0.5}, + ) + raise AssertionError(f"Unexpected request: {request.method} {request.url}") + + async def run_task() -> float: + transport = httpx.MockTransport(handler) + async with httpx.AsyncClient(transport=transport) as http_client: + workflow = ArenaStreamAgentWorkflow( + econfig={ + "arena_base_url": "https://arena.example", + "arena_llm_route_mode": "session_gateway", + "arena_registration_probe_interval": 0.01, + } + ) + return await workflow.run( + {"stream_id": "stream-1", "data_id": "data-1"}, + base_url="http://rollout-proxy", + api_key="session-key", + proxy_gateway_api_key="proxy-gateway-key", + proxy_session_token="session-token", + session_id="proxy-session-1", + worker_runtime=FakeWorkerRuntime(), + arena_http_client=http_client, + ) + + assert asyncio.run(run_task()) == 0.5 + assert len(registrations) == 2 + assert registrations[0] == registrations[1] + + +def test_agent_session_gateway_cleans_uncertain_registration(monkeypatch): + """A lost POST response must not leak the pre-generated model name.""" + monkeypatch.setenv("ARENA_OPENAPI_TOKEN", "arena-token") + monkeypatch.setenv("ARENA_LLM_API_KEY", "arena-gateway-key") + _reset_stats() + posted_names: list[str] = [] + deleted_names: list[str] = [] + + async def handler(request: httpx.Request) -> httpx.Response: + if request.method == "POST" and request.url.path.endswith( + "/openapi/v1/llm/models" + ): + posted_names.append(json.loads(request.content)["model_name"]) + raise httpx.ReadTimeout("lost registration response", request=request) + if request.method == "DELETE": + deleted_names.append(request.url.path.rsplit("/", 1)[-1]) + return httpx.Response(404) + raise AssertionError(f"Unexpected request: {request.method} {request.url}") + + async def run_agent() -> None: + transport = httpx.MockTransport(handler) + async with httpx.AsyncClient(transport=transport) as http_client: + workflow = ArenaStreamAgentWorkflow( + econfig={ + "arena_base_url": "https://arena.example", + "arena_llm_route_mode": "session_gateway", + "arena_request_retries": 0, + } + ) + await workflow.run( + {"stream_id": "stream-1", "data_id": "data-1"}, + base_url="http://rollout-proxy", + api_key="session-key", + proxy_gateway_api_key="proxy-gateway-key", + proxy_session_token="session-token", + session_id="proxy-session-1", + arena_http_client=http_client, + ) + + with pytest.raises(ArenaAPIError, match="request failed"): + asyncio.run(run_agent()) + + assert len(posted_names) == 1 + assert deleted_names == posted_names + stats = stats_tracker.export_all(reset=True) + assert stats["rollout/arena/registration_success"] == 0.0 + assert stats["rollout/arena/cleanup_success"] == 1.0 + + +def test_agent_session_gateway_requires_routing_capability(monkeypatch): + monkeypatch.setenv("ARENA_OPENAPI_TOKEN", "arena-token") + workflow = ArenaStreamAgentWorkflow( + econfig={ + "arena_base_url": "https://arena.example", + "arena_llm_route_mode": "session_gateway", + } + ) + + with pytest.raises(ValueError, match="proxy_gateway_api_key.*proxy_session_token"): + asyncio.run( + workflow.run( + {"stream_id": "stream-1", "data_id": "data-1"}, + base_url="http://rollout-proxy", + api_key="session-key", + session_id="proxy-session-1", + ) + ) + + +@pytest.mark.parametrize( + ("raw_reward", "expected_reward"), + [(0.98, 1.0), (0.9799, 0.0), (1.0, 1.0)], +) +def test_agent_applies_configured_reward_threshold( + monkeypatch, raw_reward, expected_reward +): + monkeypatch.setenv("ARENA_OPENAPI_TOKEN", "arena-token") + monkeypatch.setenv("ARENA_LLM_API_KEY", "test-llm-key") + workflow = ArenaStreamAgentWorkflow( + econfig={ + "arena_base_url": "https://arena.example", + "arena_reward_threshold": 0.98, + } + ) + + async def fake_register(**_kwargs): + return "http://arena.example/proxy", "model-1" + + async def fake_launch(**_kwargs): + return raw_reward + + async def fake_delete(*_args, **_kwargs): + return None + + monkeypatch.setattr(workflow.client, "register_llm_proxy_async", fake_register) + monkeypatch.setattr(workflow.client, "launch_one_task_result", fake_launch) + monkeypatch.setattr(workflow.client, "delete_llm_proxy_async", fake_delete) + + async def run_agent(): + return await workflow.run( + {"stream_id": "stream-1", "data_id": "data-1"}, + base_url="http://rollout-proxy", + api_key="session-key", + ) + + assert asyncio.run(run_agent()) == expected_reward + + +@pytest.mark.parametrize( + ("raw_reward", "expected_reward"), + [ + (0.0, 0.0), + (0.2, 0.02), + (0.75, 0.075), + (0.9799, 0.09799), + (0.98, 1.0), + (1.0, 1.0), + ], +) +def test_astra_reward_transform_scales_partial_rewards(raw_reward, expected_reward): + """Partial scores should stay continuous at one tenth strength.""" + from examples.swe.reward_transforms import astra_partial_reward + + assert astra_partial_reward(raw_reward, {}) == pytest.approx(expected_reward) + + +@pytest.mark.parametrize( + ("raw_reward", "expected_reward"), + [(0.8999, 0.08999), (0.9, 1.0), (0.95, 1.0)], +) +def test_agent_passes_configured_threshold_to_astra_transform( + monkeypatch, raw_reward, expected_reward +): + """The configured Arena threshold should control Astra reward shaping.""" + monkeypatch.setenv("ARENA_OPENAPI_TOKEN", "arena-token") + workflow = ArenaStreamAgentWorkflow( + econfig={ + "arena_base_url": "https://arena.example", + "arena_reward_threshold": 0.9, + "arena_reward_transform_fn": ( + "examples.swe.reward_transforms.astra_partial_reward" + ), + } + ) + + assert workflow._transform_reward(raw_reward, {}) == pytest.approx(expected_reward) + + +@pytest.mark.parametrize("raw_reward", [-0.01, 1.01, float("nan"), float("inf")]) +def test_astra_reward_transform_rejects_invalid_raw_reward(raw_reward): + """The transform should fail fast on non-probability rewards.""" + from examples.swe.reward_transforms import astra_partial_reward + + with pytest.raises(ValueError, match=r"within \[0, 1\]"): + astra_partial_reward(raw_reward, {}) + + +@pytest.mark.parametrize( + ("rewards", "expected"), + [ + ([0.0, 0.0], False), + ([1.0, 1.0], False), + ([0.1, 0.1], False), + ([0.0, 1.0], True), + ([0.0, 0.1], True), + ([0.1, 1.0], True), + ], +) +def test_arena_filter_rejects_all_equal_shaped_groups(rewards, expected): + """Mean centering gives all-equal groups zero learning signal.""" + _reset_stats() + + assert filter_function({"original_rewards": rewards}) is expected + + +def test_arena_filter_uses_shaped_original_rewards_before_centering(): + """Centered rewards must not hide variance in shaped original rewards.""" + _reset_stats() + sample = { + "rewards": [-0.01, 0.01], + "original_rewards": [0.02, 0.04], + } + + assert filter_function(sample) is True + + +@pytest.mark.parametrize( + ("rewards", "expected_all_correct"), + [ + ([0.05, 0.05], 0), + ([0.999, 0.999], 1), + ([0.9989, 0.9989], 0), + ], +) +def test_arena_filter_uses_tolerance_for_all_correct_metrics( + monkeypatch, rewards, expected_all_correct +): + """Near-one floating rewards should be classified with a stable tolerance.""" + metrics = {} + tracker = type( + "Tracker", + (), + {"scalar": lambda _, **values: metrics.update(values)}, + )() + monkeypatch.setattr(stats_tracker, "get", lambda _: tracker) + + assert filter_function({"original_rewards": rewards}) is False + assert metrics["rejected_by_all_correct"] == expected_all_correct + assert metrics["rejected_by_all_wrong"] == 1 - expected_all_correct + + +def test_agent_passes_configured_threshold_to_custom_reward_transform(monkeypatch): + monkeypatch.setenv("ARENA_OPENAPI_TOKEN", "arena-token") + monkeypatch.setenv("ARENA_LLM_API_KEY", "test-llm-key") + + def transform(reward, data, *, reward_threshold): + assert reward == 0.25 + assert data["data_id"] == "data-1" + assert reward_threshold == 0.98 + return 0.875 + + monkeypatch.setattr( + "examples.swe.arena_agent.import_from_string", lambda _: transform + ) + workflow = ArenaStreamAgentWorkflow( + econfig={ + "arena_base_url": "https://arena.example", + "arena_reward_threshold": 0.98, + "arena_reward_transform_fn": "tests.test_swe_arena.transform", + } + ) + + async def fake_register(**_kwargs): + return "http://arena.example/proxy", "model-1" + + async def fake_launch(**_kwargs): + return 0.25 + + async def fake_delete(*_args, **_kwargs): + return None + + monkeypatch.setattr(workflow.client, "register_llm_proxy_async", fake_register) + monkeypatch.setattr(workflow.client, "launch_one_task_result", fake_launch) + monkeypatch.setattr(workflow.client, "delete_llm_proxy_async", fake_delete) + + async def run_agent(): + return await workflow.run( + {"stream_id": "stream-1", "data_id": "data-1"}, + base_url="http://rollout-proxy", + api_key="session-key", + ) + + assert asyncio.run(run_agent()) == 0.875 + + +def test_agent_records_global_and_domain_raw_reward_after_transform(monkeypatch): + """Reward metrics should expose values on both sides of the transform.""" + monkeypatch.setenv("ARENA_OPENAPI_TOKEN", "arena-token") + monkeypatch.setenv("ARENA_LLM_API_KEY", "test-llm-key") + _reset_stats() + + def transform(reward, _data): + return reward * 0.1 + + monkeypatch.setattr( + "examples.swe.arena_agent.import_from_string", lambda _: transform + ) + workflow = ArenaStreamAgentWorkflow( + econfig={ + "arena_base_url": "https://arena.example", + "arena_reward_transform_fn": "tests.test_swe_arena.transform", + } + ) + + async def fake_register(**_kwargs): + return "http://arena.example/proxy", "model-1" + + async def fake_launch(**_kwargs): + return 0.75 + + async def fake_delete(*_args, **_kwargs): + return None + + monkeypatch.setattr(workflow.client, "register_llm_proxy_async", fake_register) + monkeypatch.setattr(workflow.client, "launch_one_task_result", fake_launch) + monkeypatch.setattr(workflow.client, "delete_llm_proxy_async", fake_delete) + + data = { + "stream_id": "stream-1", + "data_id": "data-1", + "arena_task_type": "swe", + } + + async def run_and_record_metrics(): + async with httpx.AsyncClient(trust_env=False) as http_client: + reward = await workflow.run( + data, + base_url="http://rollout-proxy", + api_key="session-key", + arena_http_client=http_client, + ) + stats_tracker.get(workflow_context.stat_scope()).scalar(reward=reward) + workflow.record_episode_metrics(data, reward) + return reward + + assert asyncio.run(run_and_record_metrics()) == pytest.approx(0.075) + stats = stats_tracker.export_all(reset=True) + assert stats["rollout/reward"] == pytest.approx(0.075) + assert stats["rollout/raw_reward"] == pytest.approx(0.75) + assert stats["rollout/swe/reward"] == pytest.approx(0.075) + assert stats["rollout/swe/raw_reward"] == pytest.approx(0.75) + assert stats["debug/task/data-1/raw_reward"] == pytest.approx(0.75) + assert "rollout/task/data-1/raw_reward" not in stats + + +def test_agent_records_global_and_domain_pass_at_k(monkeypatch): + """A group passes when at least one of its k Arena rollouts succeeds.""" + monkeypatch.setenv("ARENA_OPENAPI_TOKEN", "arena-token") + _reset_stats() + workflow = ArenaStreamAgentWorkflow( + econfig={"arena_base_url": "https://arena.example"} + ) + + workflow.record_group_metrics( + {"arena_task_type": "swe"}, + [0.0, None, 1.0, 0.0], + group_size=4, + ) + + stats = stats_tracker.export_all(reset=True) + assert stats["rollout/pass@k"] == 1.0 + assert stats["rollout/swe/pass@k"] == 1.0 + + +def test_agent_records_failed_pass_at_k_group(monkeypatch): + """A group fails pass@k when none of its rollouts has positive reward.""" + monkeypatch.setenv("ARENA_OPENAPI_TOKEN", "arena-token") + _reset_stats() + workflow = ArenaStreamAgentWorkflow( + econfig={"arena_base_url": "https://arena.example"} + ) + + workflow.record_group_metrics( + {"arena_task_type": "math"}, + [0.0, None, -1.0], + group_size=3, + ) + + stats = stats_tracker.export_all(reset=True) + assert stats["rollout/pass@k"] == 0.0 + assert stats["rollout/math/pass@k"] == 0.0 + + +def test_agent_does_not_count_middle_bucket_as_pass(monkeypatch): + """Partial bucket rewards should not inflate Arena pass@k metrics.""" + monkeypatch.setenv("ARENA_OPENAPI_TOKEN", "arena-token") + _reset_stats() + workflow = ArenaStreamAgentWorkflow( + econfig={ + "arena_base_url": "https://arena.example", + "arena_reward_threshold": 0.98, + "arena_reward_transform_fn": "examples.swe.reward_transforms.astra_partial_reward", + } + ) + + workflow.record_group_metrics( + {"arena_task_type": "swe"}, + [0.0, 0.05, 0.09], + group_size=3, + ) + + stats = stats_tracker.export_all(reset=True) + assert stats["rollout/pass@k"] == 0.0 + assert stats["rollout/swe/pass@k"] == 0.0 + + +def test_agent_counts_transformed_success_without_counting_partial_reward(monkeypatch): + """The metric threshold distinguishes transformed success from partial credit.""" + monkeypatch.setenv("ARENA_OPENAPI_TOKEN", "arena-token") + _reset_stats() + workflow = ArenaStreamAgentWorkflow( + econfig={ + "arena_base_url": "https://arena.example", + "arena_reward_threshold": 0.98, + "arena_reward_transform_fn": "examples.swe.reward_transforms.astra_partial_reward", + } + ) + + workflow.record_group_metrics( + {"arena_task_type": "swe"}, + [0.09799, 1.0], + group_size=2, + ) + + stats = stats_tracker.export_all(reset=True) + assert stats["rollout/pass@k"] == 1.0 + assert stats["rollout/group_pass_1"] == 1.0 + + +@pytest.mark.parametrize( + ("threshold", "rewards", "expected_pass"), + [(2.0, [1.0, 0.0], 1.0), (0.0, [0.0, 0.0], 0.0)], +) +def test_agent_threshold_only_pass_metrics_use_binary_reward( + monkeypatch, threshold, rewards, expected_pass +): + """Threshold-only transforms should classify their already-binary output.""" + monkeypatch.setenv("ARENA_OPENAPI_TOKEN", "arena-token") + _reset_stats() + workflow = ArenaStreamAgentWorkflow( + econfig={ + "arena_base_url": "https://arena.example", + "arena_reward_threshold": threshold, + } + ) + + workflow.record_group_metrics({}, rewards, group_size=2) + + stats = stats_tracker.export_all(reset=True) + assert stats["rollout/pass@k"] == expected_pass + + +def test_agent_group_metrics_follow_eval_scope(monkeypatch): + """Arena metrics should use eval-rollout during evaluation.""" + monkeypatch.setenv("ARENA_OPENAPI_TOKEN", "arena-token") + _reset_stats() + workflow = ArenaStreamAgentWorkflow( + econfig={"arena_base_url": "https://arena.example"} + ) + workflow_context.set(WorkflowContext(is_eval=True)) + try: + workflow.record_group_metrics( + {"arena_task_type": "swe"}, + [0.0, 1.0], + group_size=2, + ) + finally: + workflow_context.set(WorkflowContext()) + + stats = stats_tracker.export_all(reset=True) + assert stats["eval-rollout/pass@k"] == 1.0 + assert stats["eval-rollout/swe/pass@k"] == 1.0 + assert "rollout/pass@k" not in stats + + +def test_proxy_records_domain_reward_after_export_boundary(monkeypatch): + """The proxy callback should record domain reward exactly once in eval scope.""" + monkeypatch.setenv("ARENA_OPENAPI_TOKEN", "arena-token") + _reset_stats() + agent = ArenaStreamAgentWorkflow( + econfig={"arena_base_url": "https://arena.example"} + ) + proxy = OpenAIProxyWorkflow(mode="inline", agent=agent) + workflow_context.set(WorkflowContext(is_eval=True)) + try: + proxy.record_episode_metrics({"arena_task_type": "swe"}, 0.75) + finally: + workflow_context.set(WorkflowContext()) + + stats = stats_tracker.export_all(reset=True) + assert stats["eval-rollout/swe/reward"] == 0.75 + assert stats["eval-rollout/swe/reward__count"] == 1 + assert "rollout/swe/reward" not in stats + + +@pytest.mark.parametrize("arena_task_type", [None, "", " ", 1]) +def test_agent_invalid_task_type_records_unknown_metrics(monkeypatch, arena_task_type): + """Invalid task types should be aggregated under unknown instead of failing.""" + monkeypatch.setenv("ARENA_OPENAPI_TOKEN", "arena-token") + _reset_stats() + workflow = ArenaStreamAgentWorkflow( + econfig={"arena_base_url": "https://arena.example"} + ) + + workflow.record_group_metrics( + {"arena_task_type": arena_task_type}, + [0.0, 1.0], + group_size=2, + ) + + stats = stats_tracker.export_all(reset=True) + assert stats["rollout/unknown/pass@k"] == 1.0 + + +def test_agent_records_group_pass_count_distribution(monkeypatch): + """Each group contributes one-hot values across every pass-count bucket.""" + monkeypatch.setenv("ARENA_OPENAPI_TOKEN", "arena-token") + _reset_stats() + workflow = ArenaStreamAgentWorkflow( + econfig={"arena_base_url": "https://arena.example"} + ) + data = {"arena_task_type": "swe"} + + for rewards in ( + [0.0, 0.0, 0.0, 0.0], + [1.0, 0.0, 0.0, 0.0], + [1.0, 1.0, 0.0, 0.0], + [1.0, 1.0, 1.0, 1.0], + ): + workflow.record_group_metrics(data, rewards, group_size=4) + + stats = stats_tracker.export_all(reset=True) + assert stats["rollout/group_pass_0"] == 0.25 + assert stats["rollout/group_pass_1"] == 0.25 + assert stats["rollout/group_pass_2"] == 0.25 + assert stats["rollout/group_pass_3"] == 0.0 + assert stats["rollout/group_pass_4"] == 0.25 + assert stats["rollout/group_pass_1__count"] == 4 + + +def test_agent_records_arena_terminal_failure_metrics(monkeypatch): + """Arena terminal failures should be counted separately from reward=0.""" + monkeypatch.setenv("ARENA_OPENAPI_TOKEN", "arena-token") + monkeypatch.setenv("ARENA_LLM_API_KEY", "test-llm-key") + _reset_stats() + + registered_model_name = "" + + def handler(request: httpx.Request) -> httpx.Response: + nonlocal registered_model_name + if request.method == "POST" and request.url.path.endswith( + "/openapi/v1/llm/models" + ): + registered_model_name = json.loads(request.content)["model_name"] + return httpx.Response(200, json={"model_name": registered_model_name}) + if request.url.path.endswith("/streams/stream-1/launch_one_task"): + return httpx.Response(202, json={"task_id": "task-1"}) + if request.method == "DELETE": + return httpx.Response(204) + assert request.url.path.endswith("/tasks/task-1/result") + return httpx.Response(200, json={"status": "COLLECT_FAILED", "score": 0}) + + async def run_agent() -> None: + transport = httpx.MockTransport(handler) + async with httpx.AsyncClient(transport=transport) as http_client: + workflow = ArenaStreamAgentWorkflow( + econfig={ + "arena_base_url": "https://arena.example", + "arena_poll_interval": 0.0, + "timeout": 10.0, + } + ) + await workflow.run( + {"stream_id": "stream-1", "data_id": "data-1"}, + base_url="http://rollout-proxy", + api_key="session-key", + arena_http_client=http_client, + ) + + with pytest.raises(ArenaAPIError, match="COLLECT_FAILED"): + asyncio.run(run_agent()) + + stats = stats_tracker.export_all(reset=True) + assert stats["rollout/arena/call_success"] == 0.0 + assert stats["rollout/arena/registration_success"] == 1.0 + assert stats["rollout/arena/launch_success"] == 1.0 + assert stats["rollout/arena/terminal_success"] == 0.0 + assert stats["rollout/arena/cleanup_success"] == 1.0 + assert stats["rollout/arena/call_success__count"] == 1 + + +def test_agent_timeout_logs_task_and_data_ids(monkeypatch): + monkeypatch.setenv("ARENA_OPENAPI_TOKEN", "arena-token") + monkeypatch.setenv("ARENA_LLM_API_KEY", "test-llm-key") + _reset_stats() + workflow = ArenaStreamAgentWorkflow( + econfig={"arena_base_url": "https://arena.example", "timeout": 0.01} + ) + warnings: list[tuple[str, tuple[object, ...]]] = [] + + async def register_proxy(**_kwargs): + return ( + "https://arena.example/api", + "stream-areal-test", + ) + + async def launch_task(*, on_launch_success, **_kwargs): + on_launch_success("task-123") + await asyncio.Event().wait() + + async def delete_proxy(*_args, **_kwargs): + return None + + monkeypatch.setattr(workflow.client, "register_llm_proxy_async", register_proxy) + monkeypatch.setattr(workflow.client, "launch_one_task_result", launch_task) + monkeypatch.setattr(workflow.client, "delete_llm_proxy_async", delete_proxy) + monkeypatch.setattr( + "examples.swe.arena_agent.logger.warning", + lambda message, *args: warnings.append((message, args)), + ) + + async def run_agent() -> None: + async with httpx.AsyncClient() as http_client: + await workflow.run( + {"stream_id": "stream-1", "data_id": "data-456"}, + base_url="http://rollout-proxy", + api_key="session-key", + arena_http_client=http_client, + ) + + with pytest.raises(TimeoutError): + asyncio.run(run_agent()) + + assert warnings == [ + ( + "Arena task wait timed out: task_id=%s data_id=%s timeout_seconds=%s", + ("task-123", "data-456", 0.01), + ) + ] + + +def test_agent_cleanup_failure_does_not_mask_registration_failure(monkeypatch): + """Registration errors should survive a second failure during cleanup.""" + monkeypatch.setenv("ARENA_OPENAPI_TOKEN", "arena-token") + monkeypatch.setenv("ARENA_LLM_API_KEY", "test-llm-key") + _reset_stats() + workflow = ArenaStreamAgentWorkflow( + econfig={"arena_base_url": "https://arena.example"} + ) + deleted_models: list[str] = [] + + async def fail_registration(**_kwargs): + raise ArenaAPIError("registration failed") + + async def fail_cleanup(model_name, **_kwargs): + deleted_models.append(model_name) + raise ArenaAPIError("cleanup failed") + + monkeypatch.setattr(workflow.client, "register_llm_proxy_async", fail_registration) + monkeypatch.setattr(workflow.client, "delete_llm_proxy_async", fail_cleanup) + + async def run_agent() -> None: + async with httpx.AsyncClient() as http_client: + await workflow.run( + {"stream_id": "stream-1", "data_id": "data-1"}, + base_url="http://rollout-proxy", + api_key="session-key", + arena_http_client=http_client, + ) + + with pytest.raises(ArenaAPIError, match="registration failed"): + asyncio.run(run_agent()) + assert len(deleted_models) == 1 + assert deleted_models[0].startswith("stream-areal-") + stats = stats_tracker.export_all(reset=True) + assert stats["rollout/arena/call_success"] == 0.0 + assert stats["rollout/arena/registration_success"] == 0.0 + assert stats["rollout/arena/cleanup_success"] == 0.0 + + +def test_agent_cleanup_failure_preserves_valid_task_reward(monkeypatch): + """Cleanup failure must not invalidate an already-computed Arena reward.""" + monkeypatch.setenv("ARENA_OPENAPI_TOKEN", "arena-token") + monkeypatch.setenv("ARENA_LLM_API_KEY", "test-llm-key") + _reset_stats() + workflow = ArenaStreamAgentWorkflow( + econfig={"arena_base_url": "https://arena.example"} + ) + + async def register(**_kwargs): + return "http://arena.example/proxy", "model-1" + + async def launch(**_kwargs): + return 1.0 + + async def fail_cleanup(*_args, **_kwargs): + raise ArenaAPIError("cleanup failed") + + monkeypatch.setattr(workflow.client, "register_llm_proxy_async", register) + monkeypatch.setattr(workflow.client, "launch_one_task_result", launch) + monkeypatch.setattr(workflow.client, "delete_llm_proxy_async", fail_cleanup) + + async def run_agent() -> float: + return await workflow.run( + { + "stream_id": "stream-1", + "data_id": "data-1", + "arena_task_type": "swe", + }, + base_url="http://rollout-proxy", + api_key="session-key", + ) + + assert asyncio.run(run_agent()) == 1.0 + + stats = stats_tracker.export_all(reset=True) + assert stats["rollout/arena/cleanup_success"] == 0.0 + assert "rollout/reward" not in stats + assert "rollout/raw_reward" not in stats + assert "rollout/swe/reward" not in stats + assert "rollout/swe/raw_reward" not in stats + assert "debug/task/data-1/raw_reward" not in stats + + +def test_launch_one_task_failed_result_raises(monkeypatch): + """Terminal infrastructure failures must not silently become zero reward.""" + monkeypatch.setenv("ARENA_OPENAPI_TOKEN", "arena-token") + + def handler(request: httpx.Request) -> httpx.Response: + if request.method == "POST": + return httpx.Response( + 202, + json={"task_id": "task-1", "status": "PENDING"}, + ) + return httpx.Response(200, json={"status": "HARNESS_FAILED"}) + + async def launch_task() -> None: + client = ArenaOpenAPIClient( + base_url="https://arena.example", + poll_interval=0.0, + ) + async with httpx.AsyncClient( + transport=httpx.MockTransport(handler) + ) as http_client: + await client.launch_one_task( + stream_id="stream-1", + data_id="data-1", + model_name="deployment-1", + proxy_base_url="http://rollout-proxy/v1", + proxy_api_key="session-key", + client=http_client, + ) + + with pytest.raises(ArenaAPIError, match="HARNESS_FAILED"): + asyncio.run(launch_task()) + + +def test_launch_one_task_pending_launch_score_polls_result(monkeypatch): + """Accepted launch envelopes with placeholder scores must still be polled.""" + monkeypatch.setenv("ARENA_OPENAPI_TOKEN", "arena-token") + get_count = 0 + + def handler(request: httpx.Request) -> httpx.Response: + nonlocal get_count + if request.method == "POST": + return httpx.Response( + 202, + json={"task_id": "task-1", "status": "PENDING", "score": 0}, + ) + get_count += 1 + return httpx.Response(200, json={"status": "DONE", "score": 1}) + + async def launch_task() -> float: + client = ArenaOpenAPIClient( + base_url="https://arena.example", + poll_interval=0.0, + ) + async with httpx.AsyncClient( + transport=httpx.MockTransport(handler) + ) as http_client: + return await client.launch_one_task( + stream_id="stream-1", + data_id="data-1", + model_name="deployment-1", + proxy_base_url="http://rollout-proxy/v1", + proxy_api_key="session-key", + client=http_client, + ) + + result = asyncio.run(launch_task()) + assert result == 1.0 + assert get_count == 1 + + +def test_launch_one_task_failed_launch_score_raises(monkeypatch): + """Failed launch envelopes with scores must not become zero rewards.""" + monkeypatch.setenv("ARENA_OPENAPI_TOKEN", "arena-token") + + def handler(request: httpx.Request) -> httpx.Response: + assert request.method == "POST" + return httpx.Response( + 202, + json={"task_id": "task-1", "status": "HARNESS_FAILED", "score": 0}, + ) + + async def launch_task() -> None: + client = ArenaOpenAPIClient( + base_url="https://arena.example", + poll_interval=0.0, + ) + async with httpx.AsyncClient( + transport=httpx.MockTransport(handler) + ) as http_client: + await client.launch_one_task( + stream_id="stream-1", + data_id="data-1", + model_name="deployment-1", + proxy_base_url="http://rollout-proxy/v1", + proxy_api_key="session-key", + client=http_client, + ) + + with pytest.raises(ArenaTaskFailedError, match="HARNESS_FAILED"): + asyncio.run(launch_task()) diff --git a/tests/test_swe_arena_multi_stream.py b/tests/test_swe_arena_multi_stream.py new file mode 100644 index 0000000000..615addc679 --- /dev/null +++ b/tests/test_swe_arena_multi_stream.py @@ -0,0 +1,1171 @@ +"""Unit tests for Arena multi-Stream terminal-reward routing.""" + +import asyncio +import base64 +import json +from collections import Counter +from pathlib import Path +from types import SimpleNamespace + +import httpx +import pytest +import yaml + +from examples.swe.arena_agent import ArenaStreamAgentWorkflow +from examples.swe.arena_client import ( + ArenaAPIError, + ArenaOpenAPIClient, + ArenaTaskFailedError, + ArenaTaskResult, +) +from examples.swe.arena_config import ( + build_weighted_arena_rows, + load_arena_stream_configs, + parse_arena_stream_config, +) +from examples.swe.arena_preflight import _arena_json, _recent_health, validate_streams +from examples.swe.train_swe_rl import get_arena_dataset, get_arena_mixture_dataset +from examples.swe.utils import ArenaRewardRefConfig, ArenaStreamConfig + +from areal.utils import stats_tracker + +REPO_ROOT = Path(__file__).resolve().parents[1] +MULTI_STREAM_CONFIG = REPO_ROOT / "examples/swe/arena_multi_stream.yaml" + + +@pytest.mark.parametrize("source", ["file", "base64", "inline"]) +def test_load_arena_stream_configs_preserves_reward_specs(tmp_path, source): + """All input transports retain per-stream settings and common task envs.""" + entries = [ + { + "name": "astra", + "stream_id": "stream-astra", + "sampling_weight": 3, + "harness": "claude-code@1.0.0", + "expected_reward_ref": { + "key": "astrocode-bench-reward", + "version": "1.0.14", + }, + "reward_threshold": 0.98, + "reward_transform_fn": "examples.swe.reward_transforms.astra_partial_reward", + }, + { + "name": "tbench", + "stream_id": "stream-tbench", + "sampling_weight": 1, + "harness": "codex@2.0.0", + "expected_reward_ref": {"key": "tb2-reward", "version": "4.0.0"}, + }, + ] + payload = yaml.safe_dump({"streams": entries}) + config = {"arena_task_envs": {"COMMON": "1"}} + if source == "file": + path = tmp_path / "streams.yaml" + path.write_text(payload, encoding="utf-8") + config["arena_streams_file"] = str(path) + elif source == "base64": + config["arena_streams_yaml_b64"] = base64.b64encode(payload.encode()).decode() + else: + config["arena_streams"] = entries + + streams = load_arena_stream_configs(config) + + assert [stream.name for stream in streams] == ["astra", "tbench"] + assert streams[0].sampling_weight == 3.0 + assert streams[0].expected_reward_ref == ArenaRewardRefConfig( + key="astrocode-bench-reward", version="1.0.14" + ) + assert streams[0].task_envs == {"COMMON": "1"} + assert streams[1].reward_transform_fn == "" + + +@pytest.mark.parametrize("source", ["file", "base64"]) +@pytest.mark.parametrize( + ("payload", "error"), + [ + ("streams: []", "non-empty"), + ("streams: wrong", "non-empty"), + ("unexpected: []", "unknown top-level fields"), + ], +) +def test_load_arena_stream_configs_rejects_invalid_yaml( + tmp_path, source, payload, error +): + """Both YAML transports reject malformed stream containers.""" + if source == "file": + path = tmp_path / "streams.yaml" + path.write_text(payload, encoding="utf-8") + config = {"arena_streams_file": str(path)} + else: + config = {"arena_streams_yaml_b64": base64.b64encode(payload.encode()).decode()} + + with pytest.raises(ValueError, match=error): + load_arena_stream_configs(config) + + +def test_arena_stream_profile_preserves_raw_rewards(monkeypatch): + """Streams without a transform should retain their raw Arena rewards.""" + monkeypatch.setenv("ARENA_OPENAPI_TOKEN", "arena-token") + streams = yaml.safe_load(MULTI_STREAM_CONFIG.read_text())["econfig"][ + "arena_streams" + ] + workflow = ArenaStreamAgentWorkflow( + econfig={ + "arena_base_url": "https://arena.example", + "arena_streams": streams, + } + ) + data = { + "arena_stream_name": "first", + "stream_id": "your-first-stream", + } + + transformed = [ + workflow._transform_reward(raw_reward, data) for raw_reward in (0.2, 0.75, 0.98) + ] + + assert transformed == pytest.approx([0.2, 0.75, 0.98]) + + +def test_load_arena_stream_configs_rejects_invalid_base64(): + """Corrupted launcher transport must fail instead of becoming legacy config.""" + with pytest.raises(ValueError, match="not valid base64 UTF-8"): + load_arena_stream_configs({"arena_streams_yaml_b64": "not-base64%%%"}) + + +def test_preflight_openapi_maps_read_only_operations(monkeypatch): + """Multi-Stream preflight should not require an external Arena CLI.""" + monkeypatch.setenv("ARENA_OPENAPI_TOKEN", "test-token") + real_client = httpx.Client + requests: list[httpx.Request] = [] + + def handler(request: httpx.Request) -> httpx.Response: + requests.append(request) + assert request.headers["Authorization"] == "Bearer test-token" + return httpx.Response(200, json={"path": request.url.path}) + + def client_factory(*args, **kwargs): + kwargs["transport"] = httpx.MockTransport(handler) + return real_client(*args, **kwargs) + + monkeypatch.setattr("examples.swe.arena_preflight.httpx.Client", client_factory) + + _arena_json( + "https://arena.example", + ["stream", "get", "stream/a"], + retries=0, + ) + _arena_json( + "https://arena.example", + ["harness", "version-get", "claude/code", "1.2.4"], + retries=0, + ) + _arena_json( + "https://arena.example", + ["stream", "dataset", "stream/a", "--limit", "1"], + retries=0, + ) + _arena_json( + "https://arena.example", + [ + "stream", + "tasks", + "stream/a", + "--limit", + "100", + "--offset", + "23", + ], + retries=0, + ) + + assert [ + (request.method, request.url.raw_path.decode().split("?", 1)[0]) + for request in requests + ] == [ + ("GET", "/openapi/v1/streams/stream%2Fa"), + ("GET", "/openapi/v1/harnesses/claude%2Fcode/versions/1.2.4"), + ("POST", "/openapi/v1/streams/stream%2Fa/dataset"), + ("GET", "/openapi/v1/streams/stream%2Fa/tasks"), + ] + assert dict(requests[-1].url.params) == {"limit": "100", "offset": "23"} + + +def test_preflight_retries_known_transient_403(monkeypatch): + monkeypatch.setenv("ARENA_OPENAPI_TOKEN", "test-token") + monkeypatch.setattr("examples.swe.arena_preflight.time.sleep", lambda _delay: None) + attempts = 0 + + def handler(_request: httpx.Request) -> httpx.Response: + nonlocal attempts + attempts += 1 + if attempts == 1: + return httpx.Response( + 403, + text="spanner-http-ant-group-watch-all transient rejection", + ) + return httpx.Response(200, json={"data": {"status": "PUBLISHED"}}) + + with httpx.Client(transport=httpx.MockTransport(handler)) as client: + payload = _arena_json( + "https://arena.example", + ["harness", "version-get", "claude", "1"], + retries=1, + client=client, + ) + + assert payload == {"data": {"status": "PUBLISHED"}} + assert attempts == 2 + + +def test_parse_arena_stream_config_rejects_metric_unsafe_name(): + """Stream names must be safe to embed in metric keys and dump metadata.""" + with pytest.raises(ValueError, match="may only contain"): + parse_arena_stream_config({"name": "bad/name", "stream_id": "stream-a"}) + + +def test_load_arena_stream_configs_rejects_duplicate_stream_ids(): + """A copied Stream id must not silently receive multiple mixture weights.""" + with pytest.raises(ValueError, match="ids must be unique"): + load_arena_stream_configs( + { + "arena_streams": [ + {"name": "a", "stream_id": "stream-a"}, + {"name": "b", "stream_id": "stream-a"}, + ] + } + ) + + +def test_build_weighted_arena_rows_uses_attempted_group_weights(): + """An explicit subset should realize weights without repeating rows.""" + streams = [ + ArenaStreamConfig(name="astra", stream_id="a", sampling_weight=3.0), + ArenaStreamConfig(name="tbench", stream_id="b", sampling_weight=1.0), + ] + rows = build_weighted_arena_rows( + { + "astra": [ + {"arena_stream_name": "astra", "data_id": f"a-{index}"} + for index in range(9) + ], + "tbench": [ + {"arena_stream_name": "tbench", "data_id": f"b-{index}"} + for index in range(9) + ], + }, + streams, + epoch_size=12, + ) + + assert Counter(row["arena_stream_name"] for row in rows) == { + "astra": 9, + "tbench": 3, + } + assert len({row["data_id"] for row in rows}) == 12 + assert [row["arena_stream_name"] for row in rows[:4]] == [ + "astra", + "astra", + "astra", + "tbench", + ] + + +@pytest.mark.parametrize( + ("counts", "weights", "size_multiple", "prefix"), + [ + ([4, 1], [1.0, 1.0], 1, ["a", "b"]), + ([56, 441, 20], [1.0, 1.0, 1.0], 4, ["a", "b", "c", "a"]), + ([2, 3], [2.0, 1.0], 1, ["a", "a", "b"]), + ], + ids=["unequal-datasets", "unrounded-batch", "weighted"], +) +def test_build_weighted_arena_rows_full_union_keeps_each_row_once( + counts, weights, size_multiple, prefix +): + """Weights control order; the default epoch never repeats or pads source rows.""" + names = list("abc"[: len(counts)]) + streams = [ + ArenaStreamConfig(name=name, stream_id=name, sampling_weight=weight) + for name, weight in zip(names, weights) + ] + rows_by_stream = { + name: [ + {"arena_stream_name": name, "data_id": f"{name}-{index}"} + for index in range(count) + ] + for name, count in zip(names, counts) + } + + rows = build_weighted_arena_rows( + rows_by_stream, streams, size_multiple=size_multiple + ) + + assert Counter( + (row["arena_stream_name"], row["data_id"]) for row in rows + ) == Counter( + (row["arena_stream_name"], row["data_id"]) + for source in rows_by_stream.values() + for row in source + ) + assert [row["arena_stream_name"] for row in rows[: len(prefix)]] == prefix + + +def test_build_weighted_arena_rows_rejects_oversized_static_epoch(): + """An explicit static epoch may not repeat source rows.""" + streams = [ + ArenaStreamConfig(name="large", stream_id="a", sampling_weight=1.0), + ArenaStreamConfig(name="small", stream_id="b", sampling_weight=1.0), + ] + with pytest.raises(ValueError, match="cannot exceed the 5 unique source rows"): + build_weighted_arena_rows( + { + "large": [ + {"arena_stream_name": "large", "data_id": f"a-{index}"} + for index in range(4) + ], + "small": [{"arena_stream_name": "small", "data_id": "b-0"}], + }, + streams, + epoch_size=6, + ) + + +def test_build_weighted_arena_rows_common_weight_scale_is_invariant(): + """Only relative Stream weights should affect the virtual mixture.""" + rows_by_stream = { + "a": [{"arena_stream_name": "a", "data_id": "a-0"}], + "b": [ + {"arena_stream_name": "b", "data_id": "b-0"}, + {"arena_stream_name": "b", "data_id": "b-1"}, + ], + } + + baseline = build_weighted_arena_rows( + rows_by_stream, + [ + ArenaStreamConfig(name="a", stream_id="a", sampling_weight=0.5), + ArenaStreamConfig(name="b", stream_id="b", sampling_weight=1.0), + ], + ) + scaled = build_weighted_arena_rows( + rows_by_stream, + [ + ArenaStreamConfig(name="a", stream_id="a", sampling_weight=1.0), + ArenaStreamConfig(name="b", stream_id="b", sampling_weight=2.0), + ], + ) + + assert scaled == baseline + + +def test_build_weighted_arena_rows_rejects_mass_overflow(): + """Finite per-Stream weights must not produce an infinite aggregate mass.""" + with pytest.raises(ValueError, match="overflowed the mixture mass"): + build_weighted_arena_rows( + { + "a": [{"arena_stream_name": "a", "data_id": "a-0"}], + "b": [{"arena_stream_name": "b", "data_id": "b-0"}], + }, + [ + ArenaStreamConfig(name="a", stream_id="a", sampling_weight=1e308), + ArenaStreamConfig(name="b", stream_id="b", sampling_weight=1e308), + ], + ) + + +def test_get_arena_mixture_dataset_resolves_and_mixes_streams(monkeypatch): + """Dataset rows should carry Stream and reward-ref provenance.""" + monkeypatch.setenv("ARENA_OPENAPI_TOKEN", "arena-token") + configured = [ + ArenaStreamConfig( + name="astra", + stream_id="stream-a", + sampling_weight=2.0, + harness="claude-code@1.0.0", + expected_reward_ref=ArenaRewardRefConfig("reward-a", "1.0.0"), + ), + ArenaStreamConfig( + name="tbench", + stream_id="stream-b", + sampling_weight=1.0, + harness="codex@2.0.0", + expected_reward_ref=ArenaRewardRefConfig("reward-b", "2.0.0"), + ), + ] + + def resolve_stream(_self, stream_id, **_kwargs): + suffix = stream_id[-1] + return { + "stream_id": stream_id, + "default_reward_ref": { + "key": f"reward-{suffix}", + "version": "1.0.0" if suffix == "a" else "2.0.0", + }, + } + + def dataset_rows(_self, stream_id, llm_protocol, **_kwargs): + return [ + { + "stream_id": stream_id, + "data_id": f"{stream_id}-{index}", + "llm_protocol": llm_protocol, + "arena_task_type": "swe", + } + for index in range(6) + ] + + monkeypatch.setattr(ArenaOpenAPIClient, "resolve_stream", resolve_stream) + monkeypatch.setattr(ArenaOpenAPIClient, "get_all_dataset_rows", dataset_rows) + econfig = SimpleNamespace( + arena_base_url="https://arena.example", + arena_request_timeout=1.0, + arena_request_retries=0, + arena_streams=configured, + arena_streams_file="", + arena_mixture_epoch_size=9, + ) + + dataset, resolved = get_arena_mixture_dataset(econfig) + + assert Counter(dataset["arena_stream_name"]) == {"astra": 6, "tbench": 3} + assert set(dataset["reward_ref_key"]) == {"reward-a", "reward-b"} + assert [stream.llm_protocol for stream in resolved] == ["anthropic", "responses"] + + +def test_get_arena_dataset_retains_legacy_stream_id_return(monkeypatch): + """The original single-Stream helper remains scalar-return compatible.""" + monkeypatch.setenv("ARENA_OPENAPI_TOKEN", "arena-token") + monkeypatch.setattr( + ArenaOpenAPIClient, + "resolve_stream", + lambda _self, stream_id, **_kwargs: { + "stream_id": stream_id, + "default_reward_ref": None, + }, + ) + monkeypatch.setattr( + ArenaOpenAPIClient, + "get_all_dataset_rows", + lambda _self, stream_id, llm_protocol, **_kwargs: [ + { + "stream_id": stream_id, + "data_id": "data-1", + "llm_protocol": llm_protocol, + "arena_task_type": "swe", + } + ], + ) + econfig = SimpleNamespace( + arena_base_url="https://arena.example", + arena_request_timeout=1.0, + arena_request_retries=0, + arena_streams=[], + arena_streams_file="", + arena_mixture_epoch_size=0, + stream_id="stream-a", + arena_harness="", + arena_llm_protocol="chat_completions", + arena_task_envs={}, + arena_reward_threshold="0.98", + arena_reward_transform_fn="", + ) + + dataset, stream_id = get_arena_dataset(econfig) + + assert len(dataset) == 1 + assert stream_id == "stream-a" + + +def test_multi_stream_preflight_pins_reward_ref_and_checks_health( + monkeypatch, tmp_path +): + """Preflight should validate every routing tuple without exposing env values.""" + streams_file = tmp_path / "streams.yaml" + streams_file.write_text( + """ +streams: + - name: astra + stream_id: stream-a + harness: claude@1 + task_envs: + PRIVATE_VALUE: hidden + expected_reward_ref: + key: reward-a + version: "1" +""".strip(), + encoding="utf-8", + ) + + def arena_json(_base_url, args, *, retries, client): + assert retries == 0 + assert isinstance(client, httpx.Client) + if args[:2] == ["stream", "get"]: + return { + "data": { + "stream": { + "stream_id": "stream-a", + "status": "ACTIVE", + "default_reward_ref": {"key": "reward-a", "version": "1"}, + } + } + } + if args[:2] == ["harness", "version-get"]: + return {"data": {"status": "PUBLISHED"}} + if args[:2] == ["stream", "dataset"]: + return {"data": {"total": 10}} + if args[:2] == ["stream", "tasks"]: + assert args[-4:] == ["--limit", "100", "--offset", "0"] + return { + "data": [ + { + "task": { + "status": "DONE", + "harness_ref": {"key": "claude", "version": "1"}, + "reward_ref": {"key": "reward-a", "version": "1"}, + } + } + ] + } + raise AssertionError(args) + + monkeypatch.setattr("examples.swe.arena_preflight._arena_json", arena_json) + + summaries = validate_streams( + str(streams_file), + base_url="https://arena.example", + retries=0, + ) + + assert summaries[0]["expected_reward_ref"] == { + "key": "reward-a", + "version": "1", + } + assert summaries[0]["llm_protocol"] == "anthropic" + assert summaries[0]["task_envs"] == ["PRIVATE_VALUE"] + assert summaries[0]["recent_failure_rate"] == 0.0 + + +def test_preflight_health_ignores_tasks_without_exact_provenance(): + """Health must not mix tasks from another or unidentifiable grader tuple.""" + rows = [ + {"status": "DONE"}, + { + "status": "DONE", + "harness_ref": {"key": "claude", "version": "1"}, + }, + { + "status": "FAILED", + "harness_ref": {"key": "claude", "version": "1"}, + "reward_ref": {"key": "reward-a", "version": "1"}, + }, + ] + + terminal, failures = _recent_health( + rows, + harness="claude@1", + reward_ref=ArenaRewardRefConfig("reward-a", "1"), + ) + + assert (terminal, failures) == (1, 1) + + +def test_workflow_routes_harness_transform_and_raw_metrics(monkeypatch, tmp_path): + """A row should select all rollout and reward settings from its Stream.""" + stats_tracker.export_all(reset=True) + monkeypatch.setenv("ARENA_OPENAPI_TOKEN", "arena-token") + monkeypatch.setenv("ARENA_LLM_API_KEY", "llm-key") + transforms = { + "transform.a": lambda reward, _data: reward * 0.1, + "transform.b": lambda reward, _data: reward + 0.25, + } + monkeypatch.setattr( + "examples.swe.arena_agent.import_from_string", transforms.__getitem__ + ) + workflow = ArenaStreamAgentWorkflow( + econfig={ + "arena_base_url": "https://arena.example", + "arena_result_dump_dir": str(tmp_path / "arena-results"), + "arena_streams": [ + { + "name": "a", + "stream_id": "stream-a", + "harness": "claude@1", + "reward_transform_fn": "transform.a", + }, + { + "name": "b", + "stream_id": "stream-b", + "harness": "codex@2", + "task_envs": {"MODE": "b"}, + "expected_reward_ref": {"key": "reward-b", "version": "2"}, + "reward_transform_fn": "transform.b", + }, + ], + } + ) + launched = {} + + async def resolve_stream(stream_id, **_kwargs): + return { + "stream_id": stream_id, + "default_reward_ref": {"key": "reward-b", "version": "2"}, + } + + monkeypatch.setattr( + workflow.client, + "resolve_stream_async", + resolve_stream, + ) + + async def register(**_kwargs): + return "https://arena.example/api", "model-b" + + async def launch(**kwargs): + launched.update(kwargs) + kwargs["on_launch_success"]("task-b") + return ArenaTaskResult( + task_id="task-b", + status="DONE", + score=0.5, + raw={"exit_code": 0, "reward_txt": "ok"}, + trace_id="trace-b", + ) + + async def delete(*_args, **_kwargs): + return None + + monkeypatch.setattr(workflow.client, "register_llm_proxy_async", register) + monkeypatch.setattr(workflow.client, "launch_one_task_result", launch) + monkeypatch.setattr(workflow.client, "delete_llm_proxy_async", delete) + data = { + "arena_stream_name": "b", + "stream_id": "stream-b", + "data_id": "data-b", + "reward_ref_key": "reward-b", + "reward_ref_version": "2", + } + + async def run_workflow(): + transport = httpx.MockTransport(lambda _: httpx.Response(500)) + async with httpx.AsyncClient(transport=transport) as http_client: + reward = await workflow.run( + data, + base_url="http://rollout-proxy", + api_key="session-key", + arena_http_client=http_client, + ) + await workflow.persist_episode_result(data, reward) + workflow.record_episode_metrics(data, reward) + return reward + + reward = asyncio.run(run_workflow()) + + assert reward == pytest.approx(0.75) + assert launched["stream_id"] == "stream-b" + assert launched["harness"] == "codex@2" + assert launched["task_envs"] == {"MODE": "b"} + stats = stats_tracker.export_all(reset=True) + assert stats["rollout/arena_score"] == pytest.approx(0.5) + assert stats["rollout/training_score"] == pytest.approx(0.75) + assert stats["rollout/arena_raw_present"] == 1.0 + assert stats["rollout/stream/b/arena_score"] == pytest.approx(0.5) + dump_files = list((tmp_path / "arena-results").glob("arena_results_*.jsonl")) + assert len(dump_files) == 1 + dump_record = json.loads(dump_files[0].read_text(encoding="utf-8")) + assert dump_record["raw"] == {"exit_code": 0, "reward_txt": "ok"} + assert dump_record["expected_reward_ref"] == { + "key": "reward-b", + "version": "2", + } + assert dump_files[0].stat().st_mode & 0o777 == 0o600 + + +def test_workflow_rejects_row_reward_ref_mismatch(monkeypatch): + """A row produced under another grader version must fail before launch.""" + monkeypatch.setenv("ARENA_OPENAPI_TOKEN", "arena-token") + monkeypatch.setenv("ARENA_LLM_API_KEY", "llm-key") + workflow = ArenaStreamAgentWorkflow( + econfig={ + "arena_base_url": "https://arena.example", + "arena_streams": [ + { + "name": "a", + "stream_id": "stream-a", + "expected_reward_ref": {"key": "reward-a", "version": "1"}, + } + ], + } + ) + + with pytest.raises(ValueError, match="reward_ref mismatch"): + asyncio.run( + workflow.run( + { + "arena_stream_name": "a", + "stream_id": "stream-a", + "data_id": "data-a", + "reward_ref_key": "reward-a", + "reward_ref_version": "2", + }, + base_url="http://rollout-proxy", + api_key="session-key", + ) + ) + + +def test_workflow_rejects_live_reward_ref_drift(monkeypatch): + """A grader default changed after dataset discovery must block launch.""" + monkeypatch.setenv("ARENA_OPENAPI_TOKEN", "arena-token") + monkeypatch.setenv("ARENA_LLM_API_KEY", "llm-key") + workflow = ArenaStreamAgentWorkflow( + econfig={ + "arena_base_url": "https://arena.example", + "arena_streams": [ + { + "name": "a", + "stream_id": "stream-a", + "llm_protocol": "anthropic", + "expected_reward_ref": {"key": "reward-a", "version": "1"}, + } + ], + } + ) + + async def resolve_stream(_stream_id, **_kwargs): + return { + "stream_id": "stream-a", + "default_reward_ref": {"key": "reward-a", "version": "2"}, + } + + monkeypatch.setattr( + workflow.client, + "resolve_stream_async", + resolve_stream, + ) + + with pytest.raises(ArenaAPIError, match="drifted before launch"): + asyncio.run( + workflow.run( + { + "arena_stream_name": "a", + "stream_id": "stream-a", + "data_id": "data-a", + "llm_protocol": "anthropic", + "reward_ref_key": "reward-a", + "reward_ref_version": "1", + }, + base_url="http://rollout-proxy", + api_key="session-key", + ) + ) + + +@pytest.mark.parametrize( + ("status", "context_overflow", "interaction_count", "expected"), + [ + ("HARNESS_FAILED", True, 1, "model_failure_zero"), + ("HARNESS_FAILED", True, 0, "unknown_failure_reject"), + ("HARNESS_FAILED", False, 1, "unknown_failure_reject"), + ("SETUP_FAILED", True, 1, "system_failure_reject"), + ("TIMEOUT", False, 1, "model_failure_zero"), + ("TIMEOUT", True, 0, "unknown_failure_reject"), + ("NO_OUTPUT", False, 1, "model_failure_zero"), + ("NO_OUTPUT", False, 0, "unknown_failure_reject"), + ], +) +def test_arena_failure_classifier_is_conservative( + status, context_overflow, interaction_count, expected +): + """Only a typed, recoverable Arena model overflow may become reward zero.""" + error = ArenaTaskFailedError(task_id="task-1", status=status) + + disposition = ArenaStreamAgentWorkflow.classify_proxy_failure( + error, + context_overflow=context_overflow, + interaction_count=interaction_count, + ) + + assert disposition == expected + + +def test_arena_failure_classifier_rejects_non_terminal_api_error(): + """Transport and result-processing errors are system failures.""" + disposition = ArenaStreamAgentWorkflow.classify_proxy_failure( + ArenaAPIError("request failed"), + context_overflow=True, + interaction_count=1, + ) + + assert disposition == "system_failure_reject" + + +def test_arena_failure_classifier_keeps_explicit_claude_agent_phase_failure(): + """A Harness-attributed Claude agent failure is a trainable zero reward.""" + error = ArenaTaskFailedError( + task_id="task-1", + status="HARNESS_FAILED", + result=ArenaTaskResult( + task_id="task-1", + status="HARNESS_FAILED", + score=0.0, + raw={ + "error": "harness: harness agent phase exited with code 1: " + "harness: agent phase error: claude reported error:" + }, + ), + ) + + disposition = ArenaStreamAgentWorkflow.classify_proxy_failure( + error, + context_overflow=False, + interaction_count=3, + ) + + assert disposition == "model_failure_zero" + + +def test_arena_failure_classifier_rejects_ambiguous_harness_failure_raw(): + """Unattributed Harness text must not silently become training data.""" + error = ArenaTaskFailedError( + task_id="task-1", + status="HARNESS_FAILED", + result=ArenaTaskResult( + task_id="task-1", + status="HARNESS_FAILED", + score=0.0, + raw={"error": "harness failed"}, + ), + ) + + disposition = ArenaStreamAgentWorkflow.classify_proxy_failure( + error, + context_overflow=False, + interaction_count=3, + ) + + assert disposition == "unknown_failure_reject" + + +def test_workflow_audits_recovered_failure_with_final_zero_reward( + monkeypatch, tmp_path +): + """The audit row must agree with proxy-level model-failure recovery.""" + monkeypatch.setenv("ARENA_OPENAPI_TOKEN", "arena-token") + workflow = ArenaStreamAgentWorkflow( + econfig={ + "arena_base_url": "https://arena.example", + "arena_result_dump_dir": str(tmp_path / "arena-results"), + "stream_id": "stream-a", + } + ) + data = {"stream_id": "stream-a", "data_id": "data-a"} + error = ArenaTaskFailedError( + task_id="task-a", + status="HARNESS_FAILED", + result=ArenaTaskResult( + task_id="task-a", + status="HARNESS_FAILED", + score=0.0, + raw={"failure": "context_overflow"}, + ), + ) + + async def persist_result() -> None: + await workflow.record_failure_disposition(data, error, "model_failure_zero") + await workflow.persist_episode_result(data, 0.0) + + asyncio.run(persist_result()) + + dump_file = next((tmp_path / "arena-results").glob("arena_results_*.jsonl")) + records = [json.loads(line) for line in dump_file.read_text().splitlines()] + assert len(records) == 1 + assert records[0]["status"] == "HARNESS_FAILED" + assert records[0]["arena_score"] == 0.0 + assert records[0]["training_score"] == 0.0 + assert records[0]["raw"] == {"failure": "context_overflow"} + + +def test_workflow_audits_rejected_failure_without_training_reward( + monkeypatch, tmp_path +): + """System failures remain distinguishable from trainable zero rewards.""" + monkeypatch.setenv("ARENA_OPENAPI_TOKEN", "arena-token") + workflow = ArenaStreamAgentWorkflow( + econfig={ + "arena_base_url": "https://arena.example", + "arena_result_dump_dir": str(tmp_path / "arena-results"), + "stream_id": "stream-a", + } + ) + data = {"stream_id": "stream-a", "data_id": "data-a"} + error = ArenaTaskFailedError( + task_id="task-a", + status="SETUP_FAILED", + result=ArenaTaskResult( + task_id="task-a", + status="SETUP_FAILED", + score=0.0, + raw={"failure": "sandbox"}, + ), + ) + + asyncio.run( + workflow.record_failure_disposition(data, error, "system_failure_reject") + ) + + dump_file = next((tmp_path / "arena-results").glob("arena_results_*.jsonl")) + record = json.loads(dump_file.read_text()) + assert record["status"] == "SETUP_FAILED" + assert record["training_score"] is None + + +def test_workflow_audits_terminal_result_when_transform_fails(monkeypatch, tmp_path): + """A bad local transform must not discard an already-computed Arena result.""" + stats_tracker.export_all(reset=True) + monkeypatch.setenv("ARENA_OPENAPI_TOKEN", "arena-token") + monkeypatch.setenv("ARENA_LLM_API_KEY", "llm-key") + monkeypatch.setattr( + "examples.swe.arena_agent.import_from_string", + lambda _path: lambda _reward, _data: float("nan"), + ) + workflow = ArenaStreamAgentWorkflow( + econfig={ + "arena_base_url": "https://arena.example", + "arena_result_dump_dir": str(tmp_path / "arena-results"), + "arena_streams": [ + { + "name": "a", + "stream_id": "stream-a", + "reward_transform_fn": "transform.bad", + } + ], + } + ) + + async def register(**_kwargs): + return "https://arena.example/api", "model-a" + + async def launch(**kwargs): + kwargs["on_launch_success"]("task-a") + return ArenaTaskResult( + task_id="task-a", + status="DONE", + score=0.5, + raw={"grader": "complete"}, + ) + + async def delete(*_args, **_kwargs): + return None + + monkeypatch.setattr(workflow.client, "register_llm_proxy_async", register) + monkeypatch.setattr(workflow.client, "launch_one_task_result", launch) + monkeypatch.setattr(workflow.client, "delete_llm_proxy_async", delete) + + with pytest.raises(ValueError, match="non-finite"): + asyncio.run( + workflow.run( + { + "arena_stream_name": "a", + "stream_id": "stream-a", + "data_id": "data-a", + }, + base_url="http://rollout-proxy", + api_key="session-key", + ) + ) + + dump_file = next((tmp_path / "arena-results").glob("arena_results_*.jsonl")) + dump_record = json.loads(dump_file.read_text(encoding="utf-8")) + assert dump_record["raw"] == {"grader": "complete"} + assert dump_record["training_score"] is None + stats = stats_tracker.export_all(reset=True) + assert stats["rollout/arena/terminal_success"] == 1.0 + assert stats["rollout/arena/result_processing_success"] == 0.0 + + +@pytest.mark.parametrize( + "raw", + [ + None, + {"astrocode_report": {"case": {"resolved": True}}}, + {"exit_code": 0, "reward_txt": "1"}, + ["grader", "specific"], + ], +) +def test_task_result_preserves_heterogeneous_raw(monkeypatch, raw): + """Different grader payloads should not affect top-level score parsing.""" + monkeypatch.setenv("ARENA_OPENAPI_TOKEN", "arena-token") + + def handler(request: httpx.Request) -> httpx.Response: + if request.method == "POST": + return httpx.Response(202, json={"task_id": "task-1", "status": "PENDING"}) + return httpx.Response( + 200, + json={ + "task_id": "task-1", + "status": "DONE", + "score": 0.25, + "raw": raw, + "artifacts_uri": "oss://artifact", + "trace_id": "trace-1", + "computed_at": "2026-08-17T00:00:00Z", + }, + ) + + async def launch(): + client = ArenaOpenAPIClient(base_url="https://arena.example", poll_interval=0.0) + async with httpx.AsyncClient( + transport=httpx.MockTransport(handler) + ) as http_client: + return await client.launch_one_task_result( + stream_id="stream-1", + data_id="data-1", + model_name="model-1", + proxy_base_url="http://proxy", + proxy_api_key="key", + client=http_client, + ) + + result = asyncio.run(launch()) + + assert result.score == pytest.approx(0.25) + assert result.raw == raw + assert result.trace_id == "trace-1" + + +@pytest.mark.parametrize( + "payload", + [ + 0.25, + {"data": {"status": "DONE", "score": 0.25}}, + {"result": {"output": {"status": "DONE", "score": 0.25}}}, + ], +) +def test_task_result_accepts_legacy_control_envelopes(monkeypatch, payload): + """Known control envelopes remain compatible without inspecting raw.""" + monkeypatch.setenv("ARENA_OPENAPI_TOKEN", "arena-token") + + def handler(_request: httpx.Request) -> httpx.Response: + return httpx.Response(200, json=payload) + + async def launch(): + client = ArenaOpenAPIClient(base_url="https://arena.example") + async with httpx.AsyncClient( + transport=httpx.MockTransport(handler) + ) as http_client: + return await client.launch_one_task_result( + stream_id="stream-1", + data_id="data-1", + model_name="model-1", + proxy_base_url="http://proxy", + proxy_api_key="key", + client=http_client, + ) + + result = asyncio.run(launch()) + + assert result.status == "DONE" + assert result.score == pytest.approx(0.25) + + +def test_result_dump_rejects_symlinked_directory(monkeypatch, tmp_path): + """A sudo-launched audit writer must not follow attacker-controlled links.""" + monkeypatch.setenv("ARENA_OPENAPI_TOKEN", "arena-token") + stats_tracker.export_all(reset=True) + target = tmp_path / "target" + target.mkdir() + dump_link = tmp_path / "arena-results" + dump_link.symlink_to(target, target_is_directory=True) + workflow = ArenaStreamAgentWorkflow( + econfig={ + "arena_base_url": "https://arena.example", + "arena_result_dump_dir": str(dump_link), + "stream_id": "stream-a", + } + ) + + asyncio.run( + workflow._dump_task_result( + {"stream_id": "stream-a", "data_id": "data-a"}, + workflow.stream_configs["default"], + ArenaTaskResult(task_id="task-a", status="DONE", score=1.0), + training_score=1.0, + ) + ) + + assert list(target.iterdir()) == [] + stats = stats_tracker.export_all(reset=True) + assert stats["rollout/arena/result_dump_success"] == 0.0 + + +def test_task_result_does_not_extract_score_from_raw(monkeypatch): + """A grader-local score field must not replace a missing envelope score.""" + monkeypatch.setenv("ARENA_OPENAPI_TOKEN", "arena-token") + + def handler(request: httpx.Request) -> httpx.Response: + if request.method == "POST": + return httpx.Response(202, json={"task_id": "task-1", "status": "PENDING"}) + return httpx.Response( + 200, + json={ + "task_id": "task-1", + "status": "DONE", + "raw": {"score": 1.0, "reward": 1.0}, + }, + ) + + async def launch(): + client = ArenaOpenAPIClient(base_url="https://arena.example", poll_interval=0.0) + async with httpx.AsyncClient( + transport=httpx.MockTransport(handler) + ) as http_client: + await client.launch_one_task( + stream_id="stream-1", + data_id="data-1", + model_name="model-1", + proxy_base_url="http://proxy", + proxy_api_key="key", + client=http_client, + ) + + with pytest.raises(ArenaAPIError, match="top-level 'score'"): + asyncio.run(launch()) + + +def test_failed_task_retains_raw_in_typed_error(monkeypatch): + """Failure diagnostics should survive without becoming reward zero.""" + monkeypatch.setenv("ARENA_OPENAPI_TOKEN", "arena-token") + + def handler(request: httpx.Request) -> httpx.Response: + if request.method == "POST": + return httpx.Response(202, json={"task_id": "task-1", "status": "PENDING"}) + return httpx.Response( + 200, + json={ + "task_id": "task-1", + "status": "HARNESS_FAILED", + "score": 0, + "raw": {"error": "grader failed", "arca": {}}, + }, + ) + + async def launch(): + client = ArenaOpenAPIClient(base_url="https://arena.example", poll_interval=0.0) + async with httpx.AsyncClient( + transport=httpx.MockTransport(handler) + ) as http_client: + await client.launch_one_task( + stream_id="stream-1", + data_id="data-1", + model_name="model-1", + proxy_base_url="http://proxy", + proxy_api_key="key", + client=http_client, + ) + + with pytest.raises(ArenaTaskFailedError) as exc_info: + asyncio.run(launch()) + + assert exc_info.value.result is not None + assert exc_info.value.result.raw == {"error": "grader failed", "arca": {}}