refactor: update configuration and improve swarm client functionality

Author: binary-husky

.gitignore

@@ -170,3 +170,4 @@ modelscope_cache
 prompts
 swarmexp
 swarmlog
+werewolves_swarm

ajet/default_config/ajet_default.yaml

@@ -3,7 +3,7 @@ ajet:
   project_name: "ajet_default_project"
   experiment_name: "read_yaml_name"
   experiment_dir: "auto"  # {exp-dir}/{experiment_name}
-  backbone: debug # `debug` or `trinity` or `verl`
+  backbone: verl # `debug` or `trinity` or `verl`
 
 
   model:

ajet/tuner_lib/as_oai_baseurl_apikey.py

@@ -29,8 +29,11 @@ class OpenaiBaseUrlAndApiKey(BaseModel):
     episode_uuid: str = Field(default="episode_id", description="reserved field.")
 
     def as_agentscope_model(self, *args, **kwargs):
-        from agentscope.model import DashScopeChatModel
-        return DashScopeChatModel(model_name="AgentJet-Model", api_key=self.api_key, base_http_api_url=self.base_url)
+        from agentscope.model import OpenAIChatModel
+        return OpenAIChatModel(
+            model_name="AgentJet-Model", api_key=self.api_key,
+            client_args={"base_url": self.base_url}
+        )
 
     def as_raw_openai_sdk_client(self, *args, **kwargs):
         from openai import AsyncOpenAI

ajet/tuner_lib/experimental/as_swarm_client.py

@@ -195,7 +195,7 @@ def _should_throttle(self, throttle_policy: SwarmThrottlePolicy, pool_info: Curr
             self._remember_seen_task(throttle_policy.current_task_id, throttle_policy.expected_batch_size, throttle_policy.expected_num_repeat)
         return should_throttle
 
-    def begin_episode(self, discard_episode_timeout=600, episode_type="train", throttle_policy: SwarmThrottlePolicy|None = None) -> Tuple[str, OpenaiBaseUrlAndApiKey]:
+    def begin_episode(self, discard_episode_timeout=240, episode_type="train", throttle_policy: SwarmThrottlePolicy|None = None) -> Tuple[str, OpenaiBaseUrlAndApiKey]:
         """
         Block until an episode is claimed.
         Argument:
@@ -210,7 +210,7 @@ def begin_episode(self, discard_episode_timeout=600, episode_type="train", throt
         """
         return self._begin_episode_auto_retry(discard_episode_timeout, episode_type, throttle_policy)
 
-    def _begin_episode_auto_retry(self, discard_episode_timeout=600, episode_type="train", throttle_policy: SwarmThrottlePolicy|None = None) -> Tuple[str, OpenaiBaseUrlAndApiKey]:
+    def _begin_episode_auto_retry(self, discard_episode_timeout=240, episode_type="train", throttle_policy: SwarmThrottlePolicy|None = None) -> Tuple[str, OpenaiBaseUrlAndApiKey]:
         # max_episode_time: when an episode has **lasted** for more than X seconds, it will be terminated **locally** by client (call `end_episode` will be re-route to `abort_episode`)
         max_episode_time = 2*discard_episode_timeout
 

ajet/tuner_lib/experimental/as_swarm_server.py

@@ -708,7 +708,7 @@ async def get_episode_buffer():
     @app.post("/update_current_batch_rollout_pool_information", response_model=BoolResponse)
     async def update_current_batch_rollout_pool_information(req: CurrentBatchRolloutPoolInformation):
         """Update the current batch rollout pool information."""
-        if VERBOSE:
+        if DEBUG:
             logger.info(f"Running /update_current_batch_rollout_pool_information")
         try:
             with shared_mem_dict_lock:

tutorial/example_math_swarm/math.py

@@ -52,16 +52,13 @@ def main():
     )
 
     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=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
 
     executor = PeriodicDrainThreadPoolExecutor(workers=GRPO_N * REMOTE_BATCH_SIZE, auto_retry=True)
     for _ in range(NUM_EPOCH):

tutorial/example_werewolves/start.py

@@ -12,7 +12,7 @@
 
 from agentscope.agent import ReActAgent
 from agentscope.formatter import DashScopeMultiAgentFormatter, OpenAIMultiAgentFormatter
-from agentscope.model import OpenAIChatModel
+from agentscope.model import OpenAIChatModel, DashScopeChatModel
 from loguru import logger
 from pydantic import Field
 
@@ -81,8 +81,8 @@ def get_official_agent_prompt(name) -> str:
 
 class ExampleWerewolves(Workflow):
     trainable_targets: List[str] | None = Field(default=["werewolf"], description="List of agents to be fine-tuned.")
-    big_external_opponent_llm_url = "http://22.17.52.4:2888/v1"
-    big_external_opponent_llm_name = "/mnt/data_cpfs/model_cache/modelscope/hub/Qwen/Qwen/Qwen3-235B-A22B-Instruct-2507/"
+    big_external_opponent_llm_url: str = Field(default="http://22.17.52.4:2888/v1", description="The URL of the big external opponent LLM. You can replace it with any OpenAI-compatible LLM API URL.")
+    big_external_opponent_llm_name: str = Field(default="/mnt/data_cpfs/model_cache/modelscope/hub/Qwen/Qwen/Qwen3-235B-A22B-Instruct-2507/", description="The model name of the big external opponent LLM. You can replace it with any OpenAI-compatible LLM name.")
 
     async def execute(self, workflow_task: WorkflowTask, tuner: AjetTuner) -> WorkflowOutput:
 
@@ -121,9 +121,7 @@ async def execute(self, workflow_task: WorkflowTask, tuner: AjetTuner) -> Workfl
                 name=f"Player{i + 1}",
                 sys_prompt=get_official_agent_prompt(f"Player{i + 1}"),
                 model=model_for_this_agent,
-                formatter=DashScopeMultiAgentFormatter()
-                     if role in self.trainable_targets
-                     else OpenAIMultiAgentFormatter(),
+                formatter=DashScopeMultiAgentFormatter() if isinstance(model_for_this_agent, DashScopeChatModel) else OpenAIMultiAgentFormatter(),
                 max_iters=3 if role in self.trainable_targets else 5,
             )
             # agent.set_console_output_enabled(False)

tutorial/example_werewolves_swarm/agent_roll.py

@@ -1,15 +1,12 @@
 # -*- coding: utf-8 -*-
 
 import os
-import re
-import requests
-from textwrap import dedent
-from ajet.schema.task import Task, WorkflowOutput
+from ajet.schema.task import Task
 from ajet.copilot.job import AgentJetJob
 from ajet.task_reader import RouterTaskReader
 from ajet.utils.thread_executors import PeriodicDrainThreadPoolExecutor
 from ajet.tuner_lib.as_oai_baseurl_apikey import OpenaiBaseUrlAndApiKey
-from ajet.default_config.ajet_default import AjetTaskReader, HuggingfaceDatRepo
+from ajet.default_config.ajet_default import AjetTaskReader
 from ajet.tuner_lib.experimental.as_swarm_client import SwarmClient
 
 NUM_EPOCH = 10000
@@ -33,7 +30,7 @@ def main():
     swarm_worker = SwarmClient(AJET_SWARM_URL)
     swarm_worker.auto_sync_train_config_and_start_engine(
         ajet_job,
-        force_restart=True,
+        force_restart=False,
     )
 
     GRPO_N = ajet_job.num_repeat
@@ -65,6 +62,8 @@ def execute_agent(task: Task, api_baseurl_key: OpenaiBaseUrlAndApiKey):
     from tutorial.example_werewolves.start import ExampleWerewolves
     game = ExampleWerewolves(
         trainable_targets=["werewolf"],
+        big_external_opponent_llm_name="Qwen/Qwen3-235B-A22B-Instruct-2507",
+        big_external_opponent_llm_url="http://22.16.90.187/v1",
     )
     res = asyncio.run(game.execute(task, api_baseurl_key))
     return res