Merge pull request #16 from Agent-One-Lab/agents

Several bug fixes. Update the monitor. Now support setting monitor types.

Author: Reason-Wang

.gitignore

@@ -121,8 +121,10 @@ tests/e2e/toy_examples/deepspeed/synchronous/output.txt
 *.lock
 
 # data
+data/
 *.parquet
 agentfly/agents/data/*
+test_cache/
 
 # local logs
 logs
@@ -133,6 +135,10 @@ data/
 test_cache/
 /*.jpg
 /*.png
+slurm/
+*.err
+*.out
+*.log
 
 # Notebooks
 agentfly/tests/*.ipynb
@@ -146,3 +152,11 @@ test_outputs/
 agentfly/data/
 *.ipynb
 
+# training scripts
+training_scripts/
+verl/training_scripts/
+
+# training scripts
+training_scripts/
+verl/training_scripts/
+

agentfly/agents/agent_base.py

@@ -1,9 +1,10 @@
 from abc import ABC, abstractmethod
 from collections import defaultdict
+from datetime import datetime
 import json
 from .utils.messages import MessagesList
 from ..templates.templates import get_template
-from ..__init__ import AGENT_DATA_DIR
+from .. import AGENT_DATA_DIR
 from .llm_backends import (
     AsyncVLLMBackend,
     AsyncVerlBackend,
@@ -23,6 +24,7 @@
 import logging
 from .chain.streaming_observer import ConsoleStreamObserver, StreamingManager
 from .utils.tokenizer import create_processor, create_tokenizer
+from ..utils.monitor import JsonlSink, Monitor, WandbSink
 try:
     from verl.protocol import DataProto
 except ImportError:
@@ -51,10 +53,12 @@ def __init__(
         backend_config: Any = None,
         reward_fn: Callable = None,
         log_file: str = "agent",
-        project_name: str = None,
-        run_name: str = None,
         streaming: str = "console",
         debug: bool = False,
+        monitors: List[str] = [],
+        wandb_project_name: str = None,
+        wandb_run_name: str = None,
+        local_cache_dir: str = None,
         **kwargs # To pass other unused arguments
     ):
         """
@@ -94,7 +98,6 @@ def __init__(
 
         # Create appropriate tokenizer for trajectory processing
         self.tokenizer = create_tokenizer(model_name_or_path)
-
         self.processor = create_processor(model_name_or_path)
 
         self._reward_fn = reward_fn
@@ -104,8 +107,12 @@ def __init__(
         else:
             self.jinja_template = get_template(self.template).jinja_template()
 
-        self.project_name = project_name
-        self.run_name = run_name
+        self.wandb_project_name = wandb_project_name
+        self.wandb_run_name = wandb_run_name
+        self.local_cache_dir = local_cache_dir
+        self.local_run_cache_dir = None
+        self._initialize_monitor(monitors)
+
         self.streaming_manager = StreamingManager()
         if streaming == "console":
             self.streaming_manager.add_observer(ConsoleStreamObserver())
@@ -177,6 +184,17 @@ def _preprocess_messages(self, messages: List[Dict]):
 
         return messages_list.to_list()
 
+    def _initialize_monitor(self, monitors: List[str]) -> None:
+        for monitor in monitors:
+            if monitor == "local":
+                assert self.local_cache_dir is not None, "local_cache_dir must be set when using local monitor."
+                self.local_run_cache_dir = f"{os.path.join(self.local_cache_dir, os.path.basename(self.model_name_or_path), datetime.now().strftime('%Y%m%d_%H%M%S'))}"
+                Monitor.add_sink("jsonl", JsonlSink(f"{self.local_run_cache_dir}/"))
+            elif monitor == "wandb":
+                Monitor.add_sink("wandb", WandbSink(project=self.wandb_project_name, run_name=self.wandb_run_name))
+            else:
+                raise ValueError(f"Monitor {monitor} is not supported.")
+
     async def run(self,
         messages: Union[List[dict], np.ndarray, Dict],
         max_turns: int,
@@ -392,4 +410,4 @@ def get_verl_data_proto(self):
         batch = DataProto.from_single_dict(inputs, meta_info={"use_agent": True})
 
         return batch
- 
+ 

agentfly/agents/chain/chain_base.py

@@ -137,7 +137,6 @@ def __init__(self):
         self.terminal_status = ["terminal", "finish"]
         self.global_step = 0
         self.finished_chains_count = 0
-        self.initialize_monitor()
         self.monitor_info = defaultdict(list)
 
     def reset(self) -> None:
@@ -333,7 +332,7 @@ async def _run_single_chain(self,
         await done_queue.put((chain_id, chain, current_node))
 
         self.finished_chains_count += 1
-        self.monitor_chain()
+        self.monitor_chain(trajectory=current_node.messages.messages)
 
     async def _generate_response(self, current_node, tools, depth, chain_id, enable_streaming):
         """Generate response with optional streaming support."""
@@ -485,7 +484,6 @@ async def _finalize_chain(self, chain_id, chain, current_node, depth):
 
         await self.release_resources(chain_id)
 
-
     async def release_resources(self, id: str) -> None:
         for tool in self.tools:
             if isinstance(tool, Tool):
@@ -498,10 +496,6 @@ async def set_tools(self, id: str, env_args: Dict[str, Any]) -> None:
             if isinstance(tool, Tool):
                 await tool.set_env(id, env_args)
 
-    def initialize_monitor(self) -> None:
-        Monitor.add_sink("jsonl", JsonlSink(f"{AGENT_DATA_DIR}/demo_metrics.jsonl"))
-        Monitor.add_sink("wandb", WandbSink(project=self.project_name, run_name=self.run_name))
-
     def monitor_step(self) -> None:
         messages = self.get_messages()
         avg_turns = 0
@@ -589,9 +583,19 @@ def monitor_step(self) -> None:
                 emit(evt)
 
 
-    def monitor_chain(self) -> None:
+    def monitor_chain(self, trajectory) -> None:
         self.monitor_info['Agent/chains'].append(self.finished_chains_count)
         for tool in self.tools:
             if tool.is_stateful and tool.pool_size > 0:
                 self.monitor_info[f"Agent/Tool/{tool.name}/used_env_size"].append(tool.used_env_size)
 
+         # We only log the trajectory to local jsonl file, for wandb much bandwidth is needed
+        evt = MetricEvent(
+            sinks=["jsonl"],
+            kind="text",
+            name="Agent/rollout/trajectory",
+            value=json.dumps(serialize_for_json(trajectory), indent=2),
+            x=self.global_step,
+            x_name="Agent/rollout/step"
+        )
+        emit(evt)