patch save dir bug

Author: binary-husky

ajet/context_tracker/multiagent_tracking.py

@@ -319,6 +319,7 @@ def save_llm_interaction_timeline(self, tools, llm_ext_msg, timeline):
             )
         ):
             logger.bind(exception=True).info(f"General Warning: merge failure discovered.\n")
+            # from ajet import bp; bp("SWARM")
         return
 
 

ajet/launcher.py

@@ -214,7 +214,7 @@ def main():
             "Please provide a valid config file for swarm server mode."
         )
     if args.conf:
-        exp_dir = args.exp_dir or DEFAULT_DIR
+        exp_base_dir = args.exp_dir or DEFAULT_DIR
         yaml_path = args.conf
         (
             main_yaml_fp,
@@ -223,7 +223,7 @@ def main():
             exp_config,
         ) = prepare_experiment_config(
             yaml_path=yaml_path,
-            exp_base_dir=exp_dir,
+            exp_base_dir=exp_base_dir,
             backbone=args.backbone,
             storage=(not args.swarm_server)
         )

ajet/swarm_cli.py

@@ -42,7 +42,7 @@ def start_swarm_server(env, config, port):
 def cmd_start(args):
     """Handle the 'start' subcommand."""
     # Use default config if not provided
-    exp_dir = args.exp_dir or DEFAULT_DIR
+    exp_base_dir = args.exp_dir or DEFAULT_DIR
     if not args.conf:
         args.conf = os.path.abspath(
             os.path.join(
@@ -62,7 +62,7 @@ def cmd_start(args):
         exp_config,
     ) = prepare_experiment_config(
         yaml_path=yaml_path,
-        exp_base_dir=exp_dir,
+        exp_base_dir=exp_base_dir,
         backbone="verl",
         storage=False
     )

ajet/task_rollout/native_parallel_worker.py

@@ -173,7 +173,7 @@ def rollout_swarm(  # noqa: C901
         Build a pool of threads to run context trackers in parallel,
         each thread re-spawn after complete, until reaching conditions to stop.
         """
-        # from ajet import bp; bp("SWARM")
+
         tracker_array: List[SingleAgentContextTracker] = []
         rollout_n = self.rollout_n
         n_batch_task = len(tasks)

ajet/tuner_lib/experimental/as_swarm_server.py

@@ -334,9 +334,11 @@ async def start_engine():
             yaml_str = shared_mem_dict["train_config_yaml"]
             config_dict = yaml_module.safe_load(yaml_str)
             backbone = config_dict.get("ajet", {}).get("backbone", "verl")
-            exp_base_dir = os.path.dirname(
-                config_dict.get("ajet", {}).get("experiment_dir", "saved_experiments")
-            )
+            DEFAULT_DIR = "saved_experiments"
+            experiment_dir = config_dict.get("ajet", {}).get("experiment_dir", DEFAULT_DIR)
+            if experiment_dir == "auto":
+                exp_base_dir = DEFAULT_DIR
+            exp_base_dir = os.path.dirname(os.path.abspath(experiment_dir))
 
             # Save YAML to temporary file
             with tempfile.NamedTemporaryFile(mode="w", delete=False, suffix=".yaml") as temp_file:

tutorial/example_werewolves_swarm/agent_roll.py

@@ -1,7 +1,7 @@
 # -*- coding: utf-8 -*-
 
 import os
-from ajet.schema.task import Task
+from ajet.schema.task import Task, WorkflowTask
 from ajet.copilot.job import AgentJetJob
 from ajet.task_reader import RouterTaskReader
 from ajet.utils.thread_executors import PeriodicDrainThreadPoolExecutor
@@ -24,37 +24,36 @@ def main():
         base_yaml_config="tutorial/example_werewolves_swarm/werewolves.yaml",
         algorithm="grpo",
         experiment_name="werewolves_swarm",
+        max_env_worker=128,
     )
 
     # Hand shake with remote swarm server
     swarm_worker = SwarmClient(AJET_SWARM_URL)
     swarm_worker.auto_sync_train_config_and_start_engine(
         ajet_job,
-        force_restart=False,
+        # force_restart=True,
     )
 
     GRPO_N = ajet_job.num_repeat
     REMOTE_BATCH_SIZE = ajet_job.batch_size
 
     def rollout(task):
-        try:
-            # begin episode
-            episode_uuid, api_baseurl_key = swarm_worker.begin_episode(discard_episode_timeout=60)
-            # execute agent ( base_url = api_baseurl_key.base_url, api_key = api_baseurl_key.api_key )
-            workflow_output = execute_agent(task, api_baseurl_key)  # reward is in `workflow_output`
-            # report output back to swarm remote
-            swarm_worker.end_episode(task, episode_uuid, workflow_output)
-            return
-        except:
-            pass
+        # begin episode
+        episode_uuid, api_baseurl_key = swarm_worker.begin_episode(discard_episode_timeout=240)
+        # execute agent ( base_url = api_baseurl_key.base_url, api_key = api_baseurl_key.api_key )
+        workflow_output = execute_agent(task, api_baseurl_key)  # reward is in `workflow_output`
+        # report output back to swarm remote
+        swarm_worker.end_episode(task, episode_uuid, workflow_output)
+        return
 
-    executor = PeriodicDrainThreadPoolExecutor(workers=GRPO_N * REMOTE_BATCH_SIZE, auto_retry=True)
+
+    executor = PeriodicDrainThreadPoolExecutor(workers=1, max_parallel=64, auto_retry=True, block_first_run=True)
     for _ in range(NUM_EPOCH):
         for _, task in enumerate(dataset.generate_training_tasks()):
             for _ in range(GRPO_N):
                 executor.submit_with_periodic_drain(fn=rollout, task=task)
 
-    return None
+    return
 
 
 def execute_agent(task: Task, api_baseurl_key: OpenaiBaseUrlAndApiKey):
@@ -63,9 +62,9 @@ def execute_agent(task: Task, api_baseurl_key: OpenaiBaseUrlAndApiKey):
     game = ExampleWerewolves(
         trainable_targets=["werewolf"],
         big_external_opponent_llm_name="Qwen/Qwen3-235B-A22B-Instruct-2507",
-        big_external_opponent_llm_url="http://22.14.116.243/v1",
+        big_external_opponent_llm_url="http://22.14.116.243:2888/v1",
     )
-    res = asyncio.run(game.execute(task, api_baseurl_key))
+    res = asyncio.run(game.execute(WorkflowTask(task=task), api_baseurl_key))
     return res