feat: add Arena Stream rollout integration - #1547
Conversation
Load training rows from the Arena Stream API and delegate SWE rollouts through dynamically registered AReaL proxy sessions. Key changes: - add Stream discovery, dataset loading, task launch, and model lifecycle clients - add rollout-only and two-node Slurm entrypoints with W&B support - normalize proxy message containers and redact serialized credentials - cover Arena, proxy preprocessing, and task-specific rollout collection
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces Arena Stream mode integration for SWE-bench RL training, allowing data loading and task delegation directly to the Arena online Stream OpenAPI. Key additions include an Arena OpenAPI client, an agent workflow, a rollout-only execution script, and a utility to redact sensitive credentials from serialized configurations. Feedback on these changes highlights opportunities to prevent silent data loss when processing message iterators, optimize HTTP client reuse within retry loops, avoid masking original exceptions with cleanup failures in finally blocks, and fix nested double-quoting issues in the Slurm launch script.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
| normalized: list[dict] = [] | ||
| for message in messages: | ||
| if isinstance(message, BaseModel): | ||
| message = message.model_dump() | ||
| elif isinstance(message, Mapping): | ||
| message = dict(message) | ||
| else: | ||
| return messages | ||
| normalized.append(_materialize_iterables(message)) | ||
| return _preprocess_messages(normalized) |
There was a problem hiding this comment.
If messages is a one-shot iterator or generator, iterating over it directly will consume its elements. If an unsupported message type is encountered and the function returns messages early, the caller will receive a partially consumed iterator, leading to silent data loss. Converting messages to a list first avoids this issue.
| normalized: list[dict] = [] | |
| for message in messages: | |
| if isinstance(message, BaseModel): | |
| message = message.model_dump() | |
| elif isinstance(message, Mapping): | |
| message = dict(message) | |
| else: | |
| return messages | |
| normalized.append(_materialize_iterables(message)) | |
| return _preprocess_messages(normalized) | |
| messages_list = list(messages) | |
| normalized: list[dict] = [] | |
| for message in messages_list: | |
| if isinstance(message, BaseModel): | |
| message = message.model_dump() | |
| elif isinstance(message, Mapping): | |
| message = dict(message) | |
| else: | |
| return messages_list | |
| normalized.append(_materialize_iterables(message)) | |
| return _preprocess_messages(normalized) |
| for attempt in range(self.request_retries + 1): | ||
| try: | ||
| if client is not None: | ||
| response = client.request(method, url, **kwargs) | ||
| else: | ||
| with httpx.Client(timeout=self.timeout) as owned_client: | ||
| response = owned_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 | ||
| time.sleep(min(2**attempt, 10)) | ||
| continue | ||
| if not _is_retryable_status(response.status_code): | ||
| return response | ||
| if attempt == self.request_retries: | ||
| return response | ||
| response.close() | ||
| time.sleep(min(2**attempt, 10)) |
There was a problem hiding this comment.
Instantiating a new httpx.Client inside the retry loop on every single attempt is highly inefficient. It defeats connection pooling and keep-alive, incurring significant TCP/TLS handshake overhead on each retry. Reusing a single client instance across all attempts is much more efficient.
actual_client = client or httpx.Client(timeout=self.timeout)
try:
for attempt in range(self.request_retries + 1):
try:
response = actual_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
time.sleep(min(2**attempt, 10))
continue
if not _is_retryable_status(response.status_code):
return response
if attempt == self.request_retries:
return response
response.close()
time.sleep(min(2**attempt, 10))
finally:
if client is None:
actual_client.close()| finally: | ||
| try: | ||
| await self.client.delete_llm_proxy_async( | ||
| registered_model_id, | ||
| client=client, | ||
| timeout=self.registration_timeout, | ||
| ) | ||
| logger.info( | ||
| "Deleted Arena LLM proxy registration: model_id=%s", | ||
| registered_model_id, | ||
| ) | ||
| finally: | ||
| if owns_client: | ||
| await client.aclose() |
There was a problem hiding this comment.
If delete_llm_proxy_async raises an exception (e.g., due to network issues or if the registration never succeeded), it will propagate and mask any original exception raised in the main try block (such as a timeout or registration failure). Wrapping the cleanup call in a try-except block ensures the original exception is preserved.
| finally: | |
| try: | |
| await self.client.delete_llm_proxy_async( | |
| registered_model_id, | |
| client=client, | |
| timeout=self.registration_timeout, | |
| ) | |
| logger.info( | |
| "Deleted Arena LLM proxy registration: model_id=%s", | |
| registered_model_id, | |
| ) | |
| finally: | |
| if owns_client: | |
| await client.aclose() | |
| finally: | |
| try: | |
| await self.client.delete_llm_proxy_async( | |
| registered_model_id, | |
| client=client, | |
| timeout=self.registration_timeout, | |
| ) | |
| logger.info( | |
| "Deleted Arena LLM proxy registration: model_id=%s", | |
| registered_model_id, | |
| ) | |
| except Exception as exc: | |
| logger.error( | |
| "Failed to delete Arena LLM proxy registration %s: %s", | |
| registered_model_id, | |
| exc, | |
| ) | |
| finally: | |
| if owns_client: | |
| await client.aclose() |
| active_exception = sys.exc_info()[0] is not None | ||
| delete_error: Exception | None = None | ||
| try: | ||
| arena_client.delete_llm_proxy(registered_model_id) | ||
| logger.info( | ||
| "Arena LLM registration deleted: model_id=%s", | ||
| registered_model_id, | ||
| ) | ||
| except Exception as exc: | ||
| delete_error = exc | ||
| logger.error( | ||
| "Failed to delete Arena LLM registration %s: %s", | ||
| deployment_id, | ||
| exc, | ||
| ) | ||
| finally: | ||
| end_response = proxy_client.post( | ||
| f"{proxy_base_url.rstrip('/')}/rl/end_session", | ||
| headers={"Authorization": f"Bearer {session_api_key}"}, | ||
| ) | ||
| end_response.raise_for_status() | ||
| if delete_error is not None and not active_exception: | ||
| raise delete_error |
There was a problem hiding this comment.
If end_response.raise_for_status() raises an exception, it will propagate out of the finally block and mask any original exception from the main try block or the delete_llm_proxy error. Wrapping the session termination call in a try-except block prevents masking of more critical errors.
active_exception = sys.exc_info()[0] is not None
delete_error: Exception | None = None
try:
arena_client.delete_llm_proxy(registered_model_id)
logger.info(
"Arena LLM registration deleted: model_id=%s",
registered_model_id,
)
except Exception as exc:
delete_error = exc
logger.error(
"Failed to delete Arena LLM registration %s: %s",
deployment_id,
exc,
)
finally:
try:
end_response = proxy_client.post(
f"{proxy_base_url.rstrip('/')}/rl/end_session",
headers={"Authorization": f"Bearer {session_api_key}"},
)
end_response.raise_for_status()
except Exception as exc:
logger.error("Failed to end proxy session: %s", exc)
if delete_error is not None and not active_exception:
raise delete_error| --num-rollouts "${ARENA_NUM_ROLLOUTS}" \ | ||
| --config examples/swe/qwen3_30b_a3b_grpo.yaml \ | ||
| scheduler.type=local \ | ||
| cluster.n_nodes=1 \ | ||
| cluster.n_gpus_per_node=8 \ | ||
| rollout.backend=sglang:d1t8p1 \ | ||
| gconfig.n_samples=1 \ | ||
| rollout.consumer_batch_size="${ARENA_NUM_ROLLOUTS}" \ | ||
| rollout.max_concurrent_rollouts="${ARENA_MAX_CONCURRENT_ROLLOUTS}" \ | ||
| rollout.setup_timeout=3600.0 \ | ||
| sglang.context_length="${SGLANG_CONTEXT_LENGTH:-133120}" \ | ||
| sglang.max_prefill_tokens="${SGLANG_MAX_PREFILL_TOKENS:-133119}" \ | ||
| sglang.max_running_requests="${ARENA_MAX_CONCURRENT_ROLLOUTS}" \ | ||
| sglang.cuda_graph_max_bs="${ARENA_MAX_CONCURRENT_ROLLOUTS}" \ | ||
| econfig.arena_request_timeout=30.0 \ | ||
| train_dataset.batch_size="${ARENA_NUM_ROLLOUTS}" \ | ||
| train_dataset.num_workers=0 \ | ||
| valid_dataset.batch_size="${ARENA_NUM_ROLLOUTS}" \ | ||
| valid_dataset.num_workers=0 |
There was a problem hiding this comment.
Using unescaped double quotes inside a double-quoted string passed to bash -lc prematurely terminates and restarts the outer double-quoted string. Since these parameters are numeric or simple environment variables, removing the inner double quotes entirely is cleaner and avoids potential shell parsing bugs.
| --num-rollouts "${ARENA_NUM_ROLLOUTS}" \ | |
| --config examples/swe/qwen3_30b_a3b_grpo.yaml \ | |
| scheduler.type=local \ | |
| cluster.n_nodes=1 \ | |
| cluster.n_gpus_per_node=8 \ | |
| rollout.backend=sglang:d1t8p1 \ | |
| gconfig.n_samples=1 \ | |
| rollout.consumer_batch_size="${ARENA_NUM_ROLLOUTS}" \ | |
| rollout.max_concurrent_rollouts="${ARENA_MAX_CONCURRENT_ROLLOUTS}" \ | |
| rollout.setup_timeout=3600.0 \ | |
| sglang.context_length="${SGLANG_CONTEXT_LENGTH:-133120}" \ | |
| sglang.max_prefill_tokens="${SGLANG_MAX_PREFILL_TOKENS:-133119}" \ | |
| sglang.max_running_requests="${ARENA_MAX_CONCURRENT_ROLLOUTS}" \ | |
| sglang.cuda_graph_max_bs="${ARENA_MAX_CONCURRENT_ROLLOUTS}" \ | |
| econfig.arena_request_timeout=30.0 \ | |
| train_dataset.batch_size="${ARENA_NUM_ROLLOUTS}" \ | |
| train_dataset.num_workers=0 \ | |
| valid_dataset.batch_size="${ARENA_NUM_ROLLOUTS}" \ | |
| valid_dataset.num_workers=0 | |
| --num-rollouts ${ARENA_NUM_ROLLOUTS} \ | |
| --config examples/swe/qwen3_30b_a3b_grpo.yaml \ | |
| scheduler.type=local \ | |
| cluster.n_nodes=1 \ | |
| cluster.n_gpus_per_node=8 \ | |
| rollout.backend=sglang:d1t8p1 \ | |
| gconfig.n_samples=1 \ | |
| rollout.consumer_batch_size=${ARENA_NUM_ROLLOUTS} \ | |
| rollout.max_concurrent_rollouts=${ARENA_MAX_CONCURRENT_ROLLOUTS} \ | |
| rollout.setup_timeout=3600.0 \ | |
| sglang.context_length=${SGLANG_CONTEXT_LENGTH:-133120} \ | |
| sglang.max_prefill_tokens=${SGLANG_MAX_PREFILL_TOKENS:-133119} \ | |
| sglang.max_running_requests=${ARENA_MAX_CONCURRENT_ROLLOUTS} \ | |
| sglang.cuda_graph_max_bs=${ARENA_MAX_CONCURRENT_ROLLOUTS} \ | |
| econfig.arena_request_timeout=30.0 \ | |
| train_dataset.batch_size=${ARENA_NUM_ROLLOUTS} \ | |
| train_dataset.num_workers=0 \ | |
| valid_dataset.batch_size=${ARENA_NUM_ROLLOUTS} \ | |
| valid_dataset.num_workers=0 |
Provide a checked-in Slurm entrypoint while keeping cluster paths and credentials in an ignored local environment file.
Prevent LiteLLM background probes from entering task-scoped rollout sessions after dynamic model registration.
Preserve each Stream's default Harness metadata so rollout tasks can use the native LLM protocol without lossy gateway conversions.\n\nKey changes:\n- map Claude Harnesses to Anthropic and Codex Harnesses to Responses\n- propagate the selected protocol through Arena dataset rows\n- register protocol-specific LiteLLM deployments with health checks disabled\n- cover Stream resolution and sync/async registration behavior
Allow Arena rollouts to inject explicit non-secret environment variables into the Harness sandbox while keeping proxy credentials managed by AReaL.\n\nKey changes:\n- add arena_task_envs to the SWE environment config\n- configure Claude Code terminal-title suppression in YAML\n- merge custom variables into launch_one_task requests\n- document and test the task environment flow
GitHub Actions
Failure
GitHub Actions
Failure
AReaL rollout proxies expose an OpenAI Chat Completions API. Advertise that native upstream capability so Arena converts Claude and Codex client requests before forwarding them.
|
This pull request has been automatically marked as stale because it has not had recent activity within the last 14 days. Please add a comment or push new commits to keep it active. Thank you for your contribution! |
Description
Add an Arena Stream-backed SWE rollout mode that discovers online datasets, registers the current AReaL rollout proxy with the Arena LLM gateway, launches tasks, polls their results, and cleans up model registrations.
Key changes:
Related Issue
N/A
Type of Change
Checklist
pre-commit run --all-files)./docs/build_all.sh)main/review-prcommand/create-prBreaking Change Details (if applicable):
N/A
Additional Context
Validation:
pytest -q tests/test_swe_arena.py tests/test_config_redaction.py tests/experimental/openai/test_proxy_rollout_server.py— 25 passedpre-commit run --all-files— all code, format, documentation, private-key, and generated CLI checks passed;uv-lockcould not reach the external PyTorch AArch64 wheel metadata endpointSKIP=uv-lock pre-commit run --all-files— passedKnown limitation: Stream datasets above the API's single-request limit are rejected; pagination and dataset sharding are intentionally deferred.