stage werewolve example

Author: binary-husky

ajet/copilot/job.py

@@ -132,13 +132,9 @@ def __init__(
 
 
     def build_job_from_yaml(self, yaml_path: str | None) -> dict:
-        assert self.experiment_dir is not None, "experiment_dir must be provided either in constructor or in yaml config."
         self.config_as_dict = read_ajet_hierarchical_config(
             yaml_path,
-            exp_name=self.experiment_name,
-            backbone=self.backbone,
             write_to=None,
-            exp_dir=self.experiment_dir,
         )
         self.config_as_dict = expand_ajet_hierarchical_config(self.config_as_dict, write_to=None)
         logger.info(f"Built AgentJet job config: {yaml_path}")

ajet/tuner_lib/experimental/as_swarm_client.py

@@ -52,6 +52,8 @@ def raise_for_status_with_detail(resp):
             raise RuntimeError(f"SwarmClient error {resp.status_code} with non-JSON response: {response_text}") from e
 
 
+class SwarmServerOfflineError(Exception): ...
+
 
 class SwarmClient(object):
 
@@ -437,7 +439,7 @@ def start_engine(self):
         self._wait_until_status_change_to(desired_status="ENGINE.ROLLING")
         logger.success("Training engine is now ROLLING and ready.")
 
-    def _wait_until_status_change_to(self, desired_status="ENGINE.ROLLING", verbose=True):
+    def _wait_until_status_change_to(self, desired_status="ENGINE.ROLLING", verbose=True, timeout=1800):
         """
         Poll engine status until it reaches desired_status.
         Reports status every 5 seconds while waiting.
@@ -446,12 +448,20 @@ def _wait_until_status_change_to(self, desired_status="ENGINE.ROLLING", verbose=
             self.logger_info(f"Polling engine status until {desired_status}...")
         last_report_time = time.time()
         init_poll_time = last_report_time
+        initial_status, _ = self.get_engine_status()
 
         while True:
             try:
                 current_status, _ = self.get_engine_status()
                 current_time = time.time()
 
+                # Check if timeout has been reached
+                if current_time - init_poll_time >= timeout:
+                    raise TimeoutError(f"Timeout reached while waiting for engine status to change to {desired_status}")
+
+                if (initial_status == "ENGINE.OFFLINE") and (current_status == "ENGINE.OFFLINE"):
+                    raise SwarmServerOfflineError(f"Engine status changed from {initial_status} to OFFLINE while waiting for {desired_status}. This may indicate an error in the engine. Please check the swarm server logs for details.")
+
                 # Report status every 5 seconds
                 if current_time - last_report_time >= 30:
                     if verbose:
@@ -467,6 +477,9 @@ def _wait_until_status_change_to(self, desired_status="ENGINE.ROLLING", verbose=
                 # Wait a bit before next poll
                 time.sleep(5)
 
+            except SwarmServerOfflineError as e:
+                raise e
+
             except Exception as e:
                 logger.error(f"Error polling engine status: {e}")
                 time.sleep(5)

ajet/tuner_lib/experimental/as_swarm_server.py

@@ -393,6 +393,7 @@ def override_param_callback(config):
                     main_yaml_fp,
                     env,
                     exp_config,
+                    True,  # is_swarm_server
                 ),
             )
             p.daemon = True

ajet/tuner_lib/experimental/interchange_utils.py

@@ -111,8 +111,12 @@ class UpdateEngineStatusRequest(BaseModel):
 
 def get_interchange_server_url(config):
     port = os.getenv("AJET_DAT_INTERCHANGE_PORT")
-    if config.ajet.interchange_server.interchange_server_port != 'auto':
-        port = str(int(config.ajet.interchange_server.interchange_server_port))
+    if isinstance(config, dict):
+        interchange_server_port = config.get("ajet", {}).get("interchange_server", {}).get("interchange_server_port", "auto")
+    else:
+        interchange_server_port = config.ajet.interchange_server.interchange_server_port
+    if interchange_server_port != 'auto':
+        port = str(int(interchange_server_port))
     assert port is not None, "AJET_DAT_INTERCHANGE_PORT env var must be set"
     master_node_ip = os.getenv("MASTER_NODE_IP", "localhost")
     base_url = f"http://{master_node_ip}:{port}"

ajet/utils/config_utils.py

@@ -171,7 +171,7 @@ def config_safe_guard(config: dict, backbone: str) -> dict:
 
 
 def read_ajet_hierarchical_config(
-    yaml_fp, exp_name, backbone, write_to=None, exp_dir=DEFAULT_DIR, override_param_callback=None
+    yaml_fp, exp_name=None, backbone=None, write_to=None, exp_dir=None, override_param_callback=None
 ):
     if yaml_fp is None:
         config = {
@@ -193,9 +193,12 @@ def read_ajet_hierarchical_config(
     else:
         with open(yaml_fp, "r", encoding="utf-8") as file:
             config = yaml.safe_load(file)
-    config["ajet"]["experiment_name"] = exp_name
-    config["ajet"]["experiment_dir"] = os.path.join(exp_dir, exp_name)
-    config["ajet"]["backbone"] = backbone
+    if exp_name is not None:
+        config["ajet"]["experiment_name"] = exp_name
+    if (exp_dir is not None) and (exp_name is not None):
+        config["ajet"]["experiment_dir"] = os.path.join(exp_dir, exp_name)
+    if backbone is not None:
+        config["ajet"]["backbone"] = backbone
 
     # remove extra config of verl for trinity
     if backbone == "debug":
@@ -324,7 +327,7 @@ def prepare_experiment_config(yaml_path, exp_dir, backbone, override_param_callb
 
     ## 4. edit new yaml
     config = read_ajet_hierarchical_config(
-        yaml_backup_dst, exp_name, backbone, write_to=yaml_backup_dst, exp_dir=exp_dir, override_param_callback=override_param_callback
+        yaml_backup_dst, exp_name=exp_name, backbone=backbone, write_to=yaml_backup_dst, exp_dir=exp_dir, override_param_callback=override_param_callback
     )
     config_final = expand_ajet_hierarchical_config(config, write_to=yaml_backup_dst)
 

ajet/utils/launch_utils.py

@@ -319,6 +319,7 @@ def execute_training_process(
     exe_yaml_path,
     env,
     exp_config,
+    is_swarm_server=False,
 ):
     """
     Execute the training process based on the specified backbone and configuration.
@@ -403,7 +404,13 @@ def execute_training_process(
         subprocess.run(cmd, check=True, cwd=os.path.abspath("./"), env=env)
     except subprocess.CalledProcessError as e:
         logger.error(f"Error running subprocess: {e}")
+        if is_swarm_server:
+            from ajet.tuner_lib.experimental.interchange_utils import http_change_engine_status
+            http_change_engine_status(exp_config, "ENGINE.OFFLINE", global_step=0)
         sys.exit(1)
     except Exception as e:
         logger.error(f"Unexpected error: {e}")
+        if is_swarm_server:
+            from ajet.tuner_lib.experimental.interchange_utils import http_change_engine_status
+            http_change_engine_status(exp_config, "ENGINE.OFFLINE", global_step=0)
         sys.exit(1)

docs/en/ajet-swarm-docker.md

@@ -24,6 +24,7 @@ docker run --rm -it \
   -v ./swarmlog:/workspace/log \
   -v ./swarmexp:/workspace/saved_experiments \
   -p 10086:10086 \
+  -e SWANLAB_API_KEY=$SWANLAB_API_KEY \
   --gpus=all \
   --shm-size=32GB \
   ghcr.io/modelscope/agentjet:main \
@@ -89,6 +90,7 @@ docker run --rm -it \
   -v ./swarmlog:/workspace/log \
   -v ./swarmexp:/workspace/saved_experiments \
   -p 10086:10086 \
+  -e SWANLAB_API_KEY=$SWANLAB_API_KEY \
   --gpus=all \
   --shm-size=32GB \
   ghcr.io/modelscope/agentjet:main \