Fix trainer bugs in v0.1

agentlightning/trainer.py

@@ -162,6 +162,9 @@ def _worker_main_loop(self, agent: LitAgent, worker_id: int, is_async: bool):
 
         # Now we are in child processes, so we can safely set up the environment.
         agent.set_trainer(self)
+        # TODO: this should be set elsewhere
+        if agent.trained_agents:
+            self.triplet_exporter.agent_match = agent.trained_agents
         self._initialize_worker_env(worker_id)
 
         mode = "Async" if is_async else "Sync"

agentlightning/verl/daemon.py

@@ -102,12 +102,13 @@ def __init__(
         mini_batch_size,
         pad_token_id,
         reward_fillna_value=0.0,
+        llm_timeout_seconds=600.0,
     ):
         # Server and Task Configuration
         self.server_port = port
-        self.task_timeout_seconds = 180
+        self.llm_timeout_seconds = llm_timeout_seconds
         self.server = AgentLightningServer(
-            host="0.0.0.0", port=self.server_port, task_timeout_seconds=self.task_timeout_seconds
+            host="0.0.0.0", port=self.server_port, task_timeout_seconds=self.llm_timeout_seconds
         )
         self.proxy_port = _find_available_port()  # Run proxy on a different port
 
@@ -168,7 +169,7 @@ def proxy(path):
                     data=request.get_data(),
                     cookies=request.cookies,
                     allow_redirects=False,
-                    timeout=self.task_timeout_seconds,
+                    timeout=self.llm_timeout_seconds,
                 )
                 # Filter out hop-by-hop headers before returning the response
                 excluded_headers = [