Enhance interchange server functionality with experiment directory support

Author: binary-husky

ajet/backbone/main_trinity.py

@@ -1,15 +1,26 @@
 import ray
+import os
 from trinity.cli.launcher import main
 from trinity.common.config import Config
 from trinity.explorer.explorer import Explorer
 from trinity.trainer.trainer import Trainer
 
+from ajet.utils.config_utils import read_ajet_config_with_cache
 from ajet.utils.core_env_vars import get_runtime_env
 from ajet.utils.launch_utils import set_loguru_default_color
 
+
 set_loguru_default_color()
 
 
+def get_ajet_config_from_trinity_side():
+    yaml_path = os.environ.get("AJET_CONFIG_REDIRECT", None)
+    if yaml_path is None:
+        raise ValueError("AJET_CONFIG_REDIRECT is not set in environment variables")
+    ajet_config = read_ajet_config_with_cache(yaml_path)
+    return ajet_config
+
+
 def patch_runtime_env_to_get_actor():
     """Patch the classmethod of Explorer and Trainer to pass in the runtime env."""
     runtime_env = get_runtime_env(is_trinity=True)
@@ -39,7 +50,10 @@ def patched_trainer_get_actor(cls, config: Config):
     Explorer.get_actor = classmethod(patched_explorer_get_actor)
     Trainer.get_actor = classmethod(patched_trainer_get_actor)
 
-
+    ajet_config = get_ajet_config_from_trinity_side()
+    if ajet_config.ajet.enable_experimental_reverse_proxy:
+        from ajet.tuner_lib.weight_tuner.experimental.as_oai_model_server import start_interchange_server
+        start_interchange_server(ajet_config.ajet.experiment_dir)
 
 
 if __name__ == "__main__":

ajet/backbone/main_verl.py

@@ -249,7 +249,7 @@ def run(self, config):
 
         if config.ajet.enable_experimental_reverse_proxy:
             from ajet.tuner_lib.weight_tuner.experimental.as_oai_model_server import start_interchange_server
-            start_interchange_server()
+            start_interchange_server(config.ajet.experiment_dir)
 
         # Initialize the PPO trainer.
         trainer = AjetRayPPOTrainer(

ajet/backbone/main_vllm.py

@@ -137,9 +137,10 @@ def main(config):
 
     runtime_env = get_runtime_env()
     os.environ.update(runtime_env["env_vars"])
+
     if config.ajet.enable_experimental_reverse_proxy:
         from ajet.tuner_lib.weight_tuner.experimental.as_oai_model_server import start_interchange_server
-        start_interchange_server()
+        start_interchange_server(config.ajet.experiment_dir)
 
     def companion_launch():
         import torch

ajet/tuner_lib/weight_tuner/experimental/as_oai_model_server.py

@@ -56,7 +56,7 @@ async def coro_task_1_lookup_dict_received__send_loop(key, websocket: WebSocket,
         while not stop_event.is_set():
             # Check for new requests in ajet_remote_handler_received
             if (key in ajet_remote_handler_received) and len(ajet_remote_handler_received[key]) > 0:
-                logger.warning(f"Sending new request to client for key: {key}")
+                # logger.warning(f"Sending new request to client for key: {key}")
 
                 timeline_uuid = list(ajet_remote_handler_received[key].keys())[0]
 
@@ -95,9 +95,9 @@ async def coro_task_2_lookup_dict_received__receive_loop(key, websocket: WebSock
             # Wait for client response:
             #   ajet/tuner_lib/weight_tuner/experimental/as_oai_model_client.py
             #       await websocket.send(pickle.dumps(response))
-            logger.warning(f"Waiting for response from client for key: {key}")
+            # logger.warning(f"Waiting for response from client for key: {key}")
             response_data = pickle.loads(await websocket.receive_bytes())
-            logger.warning(f"Received response from client for key: {key}")
+            # logger.warning(f"Received response from client for key: {key}")
 
             if not isinstance(response_data, ChatCompletion):
                 stop_event.set()
@@ -152,7 +152,7 @@ async def context_tracker_client_listen(websocket: WebSocket):
         assert episode_uuid_str.startswith("episode_uuid:")
         episode_uuid = episode_uuid_str.split("episode_uuid:")[-1]
 
-        logger.warning(f"WebSocket client connected for episode_uuid: {episode_uuid}")
+        # logger.warning(f"WebSocket client connected for episode_uuid: {episode_uuid}")
 
         key = f"episode_uuid:{episode_uuid}"
         active_websockets[key] = websocket
@@ -166,7 +166,7 @@ async def context_tracker_client_listen(websocket: WebSocket):
         logger.exception(f"Error in websocket connection setup: {e}")
 
     finally:
-        logger.warning(f"WebSocket client disconnected for key: {key}")
+        # logger.warning(f"WebSocket client disconnected for key: {key}")
         if key:
             # Clean up any in-progress requests for this key
             for container in [
@@ -220,7 +220,7 @@ async def chat_completions(request: Request, authorization: str = Header(None)):
     # Create timeline UUID
     timeline_uuid = uuid.uuid4().hex
     # Add to received queue
-    logger.warning(f"Received new chat completion request for agent: {agent_name}, target_tag: {target_tag}, episode_uuid: {episode_uuid}, timeline_uuid: {timeline_uuid}")
+    # logger.warning(f"Received new chat completion request for agent: {agent_name}, target_tag: {target_tag}, episode_uuid: {episode_uuid}, timeline_uuid: {timeline_uuid}")
     ajet_remote_handler_received[key][timeline_uuid] = InterchangeCompletionRequest(
         completion_request = new_req,
         agent_name = agent_name,
@@ -259,7 +259,7 @@ async def reset():
     """
     Reset endpoint to clear all state and disconnect all websockets.
     """
-    logger.warning("Resetting interchange endpoint server state.")
+    # logger.warning("Resetting interchange endpoint server state.")
     # Disconnect all websockets
     for key, ws in list(active_websockets.items()):
         try:
@@ -278,7 +278,7 @@ async def reset():
     return {"status": "reset_complete"}
 
 
-async def monitor_debug_state():
+async def monitor_debug_state(experiment_dir):
     """
     Background task to write debug state to ./interchange_debug.txt every 1 second.
     """
@@ -292,14 +292,14 @@ async def monitor_debug_state():
                 'active_websockets': list(active_websockets.keys())
             }
 
-            with open('./interchange_debug.txt', 'w') as f:
+            with open(f'{experiment_dir}/interchange_debug.txt', 'w') as f:
                 f.write(pformat(debug_info, width=120, indent=2))
                 f.write('\n')
 
-            await asyncio.sleep(1)
+            await asyncio.sleep(2)
         except Exception as e:
             logger.error(f"Error in monitor_debug_state: {e}")
-            await asyncio.sleep(1)
+            await asyncio.sleep(2)
 
 
 def ensure_dat_interchange_server_cache_clear():
@@ -332,7 +332,7 @@ def __init__(self):
         self.server_thread = None
         self.server = None
 
-    def start(self) -> int:
+    def start(self, experiment_dir) -> int:
         """
         Start the FastAPI server on a free port.
 
@@ -346,14 +346,14 @@ def start(self) -> int:
         def run_server():
             async def serve_with_monitor():
                 # Start the monitor task
-                monitor_task = asyncio.create_task(monitor_debug_state())
+                monitor_task = asyncio.create_task(monitor_debug_state(experiment_dir))
 
                 # Start the server
                 config = uvicorn.Config(
                     app=app,
                     host="0.0.0.0",
                     port=self.port,
-                    log_level="info"
+                    log_level="error"
                 )
                 server = uvicorn.Server(config)
                 await server.serve()
@@ -374,14 +374,14 @@ def stop(self):
 
 
 # Convenience function for quick server startup
-def start_interchange_server() -> int:
+def start_interchange_server(experiment_dir) -> int:
     """
     Start the interchange endpoint server and return the port number.
 
     Returns:
         int: The port number the server is running on.
     """
     server = InterchangeEndpointServer()
-    port = server.start()
+    port = server.start(experiment_dir)
     return port