Make the number of tasks on the server and client consistent. (#187)

examples/apo/legacy_apo_client.py

@@ -43,5 +43,7 @@ if __name__ == "__main__":
     configure_logger()
     dotenv.load_dotenv()
     agent = SimpleAgent()
-    trainer = Trainer(n_workers=2)
+    # Use 2 workers to simulate multiple clients
+    # max_tasks is optional, limit to 2 tasks here for a quick demo.
+    trainer = Trainer(n_workers=2, max_tasks=2)
     trainer.fit_v0(agent, "http://127.0.0.1:9997")

examples/apo/legacy_apo_server.py

@@ -23,6 +23,7 @@ async def example_apo():
         "You are a helpful assistant.",
         "You are a knowledgeable AI.",
         "You are a friendly chatbot.",
+        "You are an experienced expert.",
     ]
 
     prompt_and_rewards: list[tuple[str, float]] = []
@@ -40,7 +41,7 @@ async def example_apo():
         print(f"[Algo] Task '{task_id}' is now available for clients.")
 
         # 3. The algorithm waits for clients to process the task
-        rollout = await server.poll_completed_rollout(task_id, timeout=30)
+        rollout = await server.poll_completed_rollout(task_id, timeout=60)
         assert rollout, "Expected a completed rollout from the client."
         print(f"[Algo] Received Result: {rollout}")
         reward = rollout.final_reward