.

Author: garrett4wade

areal/experimental/openai/proxy/server.py

@@ -318,7 +318,12 @@ def __init__(
         self.host_ip = gethostip()
         self._localhost = "0.0.0.0"
         self.server_config = uvicorn.Config(
-            self.app, host=self._localhost, port=self.port, log_level=uvicorn_log_level
+            self.app,
+            host=self._localhost,
+            port=self.port,
+            log_level=uvicorn_log_level,
+            timeout_keep_alive=300,
+            workers=4,
         )
         self.server = uvicorn.Server(self.server_config)
         self.thread = threading.Thread(target=self.server.run, daemon=True)
@@ -378,15 +383,12 @@ async def fetch_next_session(self) -> str:
             except Empty:
                 await asyncio.sleep(0.1)
 
-    async def wait_for_session(
-        self, session_id: str, discount: float = 1.0, style: str = "individual"
-    ) -> SessionData:
+    async def wait_for_session(self, session_id: str) -> SessionData:
         if session_id not in self.session_cache:
             raise KeyError(f"Session {session_id} not found")
         # Wait for session to be completed using event
         await self.session_cache[session_id].wait_for_finish()
-        session = self.session_cache.pop(session_id)
-        return session.export_interactions(discount=discount, style=style)
+        return self.session_cache.pop(session_id)
 
     def set_reward(self, session_id: str, completion_id: str, reward: float):
         """Set reward for a specific completion/response by its ID."""

areal/experimental/openai/proxy/workflow.py

@@ -1,11 +1,17 @@
 from __future__ import annotations
 
+import asyncio
+import atexit
+import os
+import threading
 import traceback
+from concurrent.futures import ProcessPoolExecutor
 from typing import TYPE_CHECKING
 
 from areal.api.workflow_api import AgentWorkflow, RolloutWorkflow
 from areal.core import workflow_context
-from areal.utils import logging
+from areal.utils import logging, stats_tracker
+from areal.utils.perf_tracer import session_context, trace_session
 
 from .client_session import OpenAIProxyClientSession
 from .server import OpenAIProxyServer
@@ -16,6 +22,46 @@
 logger = logging.getLogger("OpenAIProxyWorkflow")
 
 
+# Lazy-initialized thread pool for async HTTP requests
+_executor: ProcessPoolExecutor | None = None
+_executor_lock = threading.Lock()
+
+
+def _get_executor() -> ProcessPoolExecutor:
+    """Get or create the shared process pool executor."""
+    global _executor
+    if _executor is None:
+        with _executor_lock:
+            if _executor is None:
+                _executor = ProcessPoolExecutor(max_workers=4)
+                # Register cleanup on process exit
+                atexit.register(_shutdown_executor)
+    return _executor
+
+
+def _shutdown_executor() -> None:
+    """Shutdown the shared thread pool executor if it exists.
+
+    Called via atexit at process exit, when no other threads should be
+    accessing the executor.
+    """
+    global _executor
+    if _executor is not None:
+        _executor.shutdown(wait=False)
+        _executor = None
+
+
+def _wrap_run(agent, data, extra_envs):
+    for key, value in extra_envs.items():
+        os.environ[key] = value
+
+    try:
+        return asyncio.run(agent.run(None, data))
+    except Exception:
+        logger.error(f"Agent task failed: {traceback.format_exc()}")
+        raise
+
+
 class OpenAIProxyWorkflow(RolloutWorkflow):
     def __init__(
         self,
@@ -32,6 +78,20 @@ def __init__(
         self.discount = discount
         self.export_style = export_style
 
+    @trace_session("run_agent")
+    async def _run_agent(self, base_url: str, data: dict):
+        extra_envs = {
+            "OPENAI_BASE_URL": base_url,
+        }
+        executor = _get_executor()
+        fut = executor.submit(_wrap_run, self.agent, data, extra_envs)
+        try:
+            return await asyncio.wrap_future(fut)
+        except Exception:
+            logger.error(f"Agent task failed: {traceback.format_exc()}")
+            raise
+
+    @session_context()
     async def arun_episode(self, engine: TRolloutEngine, data):
         # Ensure that we own the same engine instance
         task_id = workflow_context.get().task_id
@@ -48,11 +108,8 @@ async def arun_episode(self, engine: TRolloutEngine, data):
             async with OpenAIProxyClientSession(
                 base_url=self.proxy_server.public_addr, task_id=str(task_id)
             ) as session:
-                try:
-                    rewards = await self.agent.run(session.session_url, data)
-                except Exception:
-                    logger.error(f"Agent task failed: {traceback.format_exc()}")
-                    raise
+                rewards = await self._run_agent(session.session_url, data)
+
                 session_id = session.session_id
                 if isinstance(rewards, dict):
                     for completion_id, reward in rewards.items():
@@ -65,6 +122,14 @@ async def arun_episode(self, engine: TRolloutEngine, data):
             # Pop a session id from the server queue and ignore it.
             _ = await self.proxy_server.fetch_next_session()
 
-        return await self.proxy_server.wait_for_session(
-            session_id, discount=self.discount, style=self.export_style
+        session_data = await self.proxy_server.wait_for_session(session_id)
+        last_id = session_data.completions.last_interaction_id
+        interactions = session_data.completions.export_interactions(
+            reward_discount=self.discount, style=self.export_style
         )
+
+        # Record the last reward in wandb/tensorboard
+        last_reward = interactions[last_id].reward
+        stats_tracker.get(workflow_context.stat_scope()).scalar(reward=last_reward)
+
+        return interactions

areal/tests/experimental/openai/test_proxy.py

@@ -122,7 +122,8 @@ async def test_session_lifecycle(proxy_server):
             assert resp.status == 200
 
     # 3. after end session, can fetch results with `wait_for_session`
-    interactions = await proxy_server.wait_for_session(session_id)
+    session_data = await proxy_server.wait_for_session(session_id)
+    interactions = session_data.completions.export_interactions
     assert len(interactions) >= 1
 
 

examples/experimental/proxy/gsm8k_agent.py

@@ -1,28 +1,24 @@
+from math_verify import parse, verify
 from openai import AsyncOpenAI
 from openai.types.chat import ChatCompletion
 
 from areal.api.workflow_api import AgentWorkflow
-from areal.core import workflow_context
-from areal.utils import stats_tracker
 
 
 class GSM8kAgent(AgentWorkflow):
     def __init__(self, **kwargs):
         self.kwargs = kwargs
 
     async def run(self, base_url: str, data: dict):
-        async with AsyncOpenAI(base_url=base_url) as client:
+        # custom_timeout = httpx.Timeout(30.0, read=600.0)
+        # async with AsyncOpenAI(base_url=base_url, max_retries=0,
+        # timeout=custom_timeout) as client:
+        async with AsyncOpenAI(max_retries=0) as client:
             comp: ChatCompletion = await client.chat.completions.create(
                 messages=data["messages"], model="default", **self.kwargs
             )
 
-        # compute reward with areal's existing implementation
-        # Use the following wrapper to suppress the annoying warning of math-verify
-        from areal.api.reward_api import AsyncRewardWrapper
-        from areal.reward.gsm8k import gsm8k_reward_fn
-
-        reward = await AsyncRewardWrapper(gsm8k_reward_fn)(
-            None, comp.choices[0].message.content, None, None, answer=data["answer"]
-        )
-        stats_tracker.get(workflow_context.stat_scope()).scalar(reward=reward)
-        return reward
+        ans = parse(comp.choices[0].message.content)
+        gold = parse(data["answer"])
+        reward = verify(ans, gold)
+        return float(reward)