allow user to use any port as swarm port

Author: binary-husky

ajet/swarm_cli.py

@@ -29,8 +29,7 @@ def start_swarm_server(env, config, port):
     )
 
     # Set the port in the config
-    if hasattr(config.ajet, 'experimental_interchange_server_port'):
-        config.ajet.experimental_interchange_server_port = port
+    config.ajet.interchange_server.interchange_server_port = port
 
     from ajet.tuner_lib.experimental.as_oai_model_server import (
         start_interchange_server,
@@ -62,7 +61,7 @@ def cmd_start(args):
         exp_name,
         exp_config,
     ) = prepare_experiment_config(
-        yaml_path, exp_dir, args.backbone, storage=False
+        yaml_path, exp_dir, "verl", storage=False
     )
 
     # Setup environment variables
@@ -75,7 +74,7 @@ def __init__(self, conf, backbone, exp_dir):
             self.swarm_overwatch = ""
             self.debug = ""
 
-    swarm_args = SwarmArgs(args.conf, args.backbone, args.exp_dir)
+    swarm_args = SwarmArgs(args.conf, "verl", args.exp_dir)
     env, exp_config = setup_environment_vars(swarm_args, exp_config, main_yaml_fp)
 
     # Start swarm server
@@ -86,8 +85,8 @@ def cmd_overwatch(args):
     """Handle the 'overwatch' subcommand."""
     from ajet.utils.swarm_overwatch import start_overwatch
 
-    logger.info(f"Starting Swarm Overwatch for server: {args.swarm_port}")
-    start_overwatch(args.swarm_port, refresh_interval=args.refresh_interval)
+    logger.info(f"Starting Swarm Overwatch for server: {args.swarm_url}")
+    start_overwatch(args.swarm_url, refresh_interval=args.refresh_interval)
 
 
 def main():
@@ -117,19 +116,13 @@ def main():
         required=False,
         help="Path to experiment directory",
     )
-    parser_start.add_argument(
-        "--backbone",
-        type=str,
-        default="verl",
-        required=False,
-        help="verl or trinity or debug",
-    )
+
     parser_start.set_defaults(func=cmd_start)
 
     # Subcommand: overwatch
     parser_overwatch = subparsers.add_parser("overwatch", help="Monitor the swarm server")
     parser_overwatch.add_argument(
-        "--swarm-port",
+        "--swarm-url",
         type=str,
         default="http://localhost:10086",
         required=False,

ajet/tuner_lib/experimental/as_swarm_client.py

@@ -469,7 +469,8 @@ def print_rollout_stat(self):
                 "Completed episodes: (current) / (required)": f"{stat.get('completed_episodes', 0)} / {stat.get('completed_episode_target', 0)}",
                 "Average episodes per task: (current) / (expected)": f"{stat.get('average_episodes_per_task', 0):.2f} / {stat.get('task_expected_num_repeat', 0)}",
                 "Completed num-dummy tasks: (current) / (required)": f"{stat.get('completed_non_dummy_tasks', 0)} / {stat.get('completed_task_target', 0)}",
-                "Tasks (Number of episodes completed for each task)": task_buffer
+                "Tasks (Number of episodes completed for each task)": task_buffer,
+                "Hint": f"Please run `ajet-swarm overwatch --swarm-url={self.server_url}` to get more details."
             }
             print_dict(stat, mod="console", header="Current Swarm Rollout Pool Information")
         except:

ajet/tuner_lib/experimental/as_swarm_server.py

@@ -625,8 +625,8 @@ async def abort_episode(req: EndEpisodeRequest):
         if VERBOSE:
             logger.info(f"Running [{episode_uuid}]: /abort_episode")
 
-        assert "task_id" in workflow_output.metadata, "workflow_output.metadata must contain task_id"
-        assert workflow_output.metadata["task_id"] == task_id, "workflow_output.metadata.task_id must match req.task_id"
+        # assert "task_id" in workflow_output.metadata, "workflow_output.metadata must contain task_id"
+        # assert workflow_output.metadata["task_id"] == task_id, "workflow_output.metadata.task_id must match req.task_id"
 
         if "episodes" not in shared_mem_dict:
             logger.error(f"[server] No episodes registered yet.")
@@ -715,6 +715,7 @@ async def get_current_batch_rollout_pool_information():
                         running_episode_details[es.episode_uuid] = {
                             "episode_status": es.episode_status,
                             "time_since_last_activity": f"{time_since_last_activity:.1f}s",
+                            "discard_episode_timeout": f"{es.discard_episode_timeout:.1f}s",
                         }
             pool_info.running_episode_details = running_episode_details if running_episode_details else None
 

ajet/tuner_lib/experimental/swarm_overwatch_utils.py

@@ -12,7 +12,7 @@ class CurrentBatchRolloutPoolInformation(BaseModel):
     completed_non_dummy_task_target: int = 0
     task_expected_num_repeat: int = 0
     completed_tasks_details: Dict[str, List[str]] = {}  # task_id -> list of episode_uuids
-    running_episode_details: Dict[str, Dict[str, str]] | None = None # episode_uuid -> { "episode_status": ..., "time_since_last_activity": ...}
+    running_episode_details: Dict[str, Dict[str, str]] | None = None # episode_uuid -> { "episode_status": ..., "time_since_last_activity": ..., "discard_episode_timeout": ...}
     engine_status: str | None = None
     global_step: int | None = None
     booting_start_time: float | None = None  # timestamp when ENGINE.BOOTING started

ajet/utils/swarm_overwatch.py

@@ -113,7 +113,7 @@ def create_summary_table(self, info: CurrentBatchRolloutPoolInformation) -> Tabl
         title_prefix = "" if is_active else "[WAITING ENGINE.ROLLING] "
 
         table = Table(
-            title=f"{title_prefix}Rollout Pool Summary (Progress to Hit Next Weight Update)",
+            title=f"{title_prefix}Completed Episode Pool Summary (Progress to Hit Next Weight Update)",
             show_header=True,
             header_style="bold magenta",
             border_style=border_style,
@@ -140,13 +140,13 @@ def create_summary_table(self, info: CurrentBatchRolloutPoolInformation) -> Tabl
 
         # Episodes
         ep_cur, ep_tgt, ep_pct = self.create_progress_bar(
-            info.completed_episodes, info.completed_episode_target, "Episodes"
+            info.completed_episodes, info.completed_episode_target, "Completed Episodes"
         )
         ep_bar = self._create_text_bar(ep_pct)
         ep_metric = (
-            "*Episodes (chosen)*"
+            "-> *Completed Episodes (chosen)*"
             if highlight_episodes
-            else "Episodes"
+            else "Completed Episodes"
         )
         ep_style = "bold green" if highlight_episodes else None
         table.add_row(
@@ -166,7 +166,7 @@ def create_summary_table(self, info: CurrentBatchRolloutPoolInformation) -> Tabl
         )
         task_bar = self._create_text_bar(task_pct)
         task_metric = (
-            "*Completed Tasks (chosen)*" if highlight_tasks else "Completed Tasks"
+            "-> *Completed Tasks (chosen)*" if highlight_tasks else "Completed Tasks"
         )
         task_style = "bold green" if highlight_tasks else None
         table.add_row(
@@ -188,7 +188,7 @@ def create_summary_table(self, info: CurrentBatchRolloutPoolInformation) -> Tabl
         )
         nd_bar = self._create_text_bar(nd_pct)
         nd_metric = (
-            "*Completed Non-Dummy Tasks (chosen)*"
+            "-> *Completed Non-Dummy Tasks (chosen)*"
             if highlight_non_dummy
             else "Completed Non-Dummy Tasks"
         )
@@ -245,7 +245,7 @@ def create_running_episodes_table(
 
         table.add_column("Episode UUID", style="cyan", no_wrap=True, width=20, overflow="ellipsis")
         table.add_column("Status", style="green", width=15)
-        table.add_column("Time Since Last Activity", style="yellow", width=30)
+        table.add_column("Last Req / Patience", style="yellow", width=30)
 
         if not info.running_episode_details:
             table.add_row("[dim]No running episodes[/dim]", "", "")
@@ -259,10 +259,12 @@ def create_running_episodes_table(
         )
 
         for episode_uuid, details in sorted_episodes[:30]:
+            last_req = details["time_since_last_activity"]
+            patience = details.get("discard_episode_timeout", "N/A")
             table.add_row(
                 episode_uuid[:40] if len(episode_uuid) > 40 else episode_uuid,
                 details["episode_status"],
-                details["time_since_last_activity"],
+                f"{last_req} / {patience}",
             )
 
         if len(sorted_episodes) > 30:

tutorial/example_academic_trans/trans_roll.py

@@ -13,7 +13,7 @@
 LOCAL_NUM_EPOCH = 1
 LOCAL_MAX_PARALLEL = 32
 LOCAL_DATASET_PATH = "/mnt/data_cpfs/qingxu.fu/agentjet/agentjet/tmp/arxiv_papers/train.parquet"
-REMOTE_SWARM_URL = "http://localhost:10086" # Change to your swarm remote url
+REMOTE_SWARM_URL = "http://localhost:10099" # Change to your swarm remote url
 
 # --------- configurations that take effect remotely -------------
 REMOTE_BATCH_SIZE = 8