Add example_cocktail_rl_v2 with configurable batch ratios and dynamic scheduling

- Implemented AppWorld and AIME swarm clients for the new cocktail RL version.
- Introduced CocktailV2Config as a single source of truth for configuration values.
- Created train_appworld_as_swarm_client_0 and train_aime_as_swarm_client_1 scripts for running the respective clients.
- Added cocktail_v2_runner to manage shared functionality between clients.
- Included readme.md for setup instructions and configuration details.
- Enhanced evaluation and logging mechanisms for better performance tracking.

Author: binary-husky

.gitignore

@@ -185,3 +185,5 @@ research_*.json
 research_*.jsonc
 daemon_logs*
 paper
+val_results.md
+cocktail_vs_separate*

ajet/copilot/monitor-with-tmux/SKILL.md

@@ -180,3 +180,11 @@ $ python3 /tmp/tmux_wait.py ajet_session 240 && tmux capture-pane -t ajet_sessio
 tmux kill-session -t ajet_session
 
 ```
+
+
+## For AgentJet Swarm
+
+- You should create seperate tmux session for each agentjet swarm servers and each agentjet swarm clients
+- When debugging, please do not restart agentjet swarm servers frequently, that waste too much time
+- When you really having difficulty for clearing GPU memory, run `ajet --autokill` to automatically kill all python and ray processes (however, I still recommend using this as a last resort).
+- For AgentJet, always use tmux session name that starts with `ajet-*`

ajet/default_config/ajet_config_schema.py

@@ -40,6 +40,7 @@ class AjetModel:
 class AjetData:
     max_prompt_length: int = 3000
     max_response_length: int = 15000
+    # Note that this value is ignored when swarm_mode_sample_collection_method="rollout_until_all_clients_agree_sync_weight"
     train_batch_size: int = 32
 
 

ajet/default_config/ajet_default.yaml

@@ -16,6 +16,7 @@ ajet:
     # max number of tokens for response
     max_response_length: 15000
     # how many tasks per training batch
+    # Note that this value is ignored when swarm_mode_sample_collection_method="rollout_until_all_clients_agree_sync_weight"
     train_batch_size: 32
     # [Hint]: The final number of samples per update will be: N_{sample} = (data.train_batch_size * rollout.num_repeat * rollout.multi_turn.expected_steps)
 

ajet/default_config/ajet_swarm_default.yaml

@@ -67,6 +67,7 @@ ajet:
     # max number of tokens for response
     max_response_length: 15000
     # how many tasks per training batch
+    # Note that this value is ignored when swarm_mode_sample_collection_method="rollout_until_all_clients_agree_sync_weight"
     train_batch_size: 32
     # [Hint]: The final number of samples per update will be: N_{sample} = (data.train_batch_size * rollout.num_repeat * rollout.multi_turn.expected_steps)
 

ajet/swarm_cli.py

@@ -5,6 +5,7 @@
 from dotenv import load_dotenv
 from loguru import logger
 
+from ajet.utils.cleaner import fast_kill_by_keyword_bash
 from ajet.utils.config_utils import prepare_experiment_config
 from ajet.utils.launch_utils import (
     dict_to_namespace,
@@ -41,6 +42,21 @@ def start_swarm_server(env, config, port):
 
 def cmd_start(args):
     """Handle the 'start' subcommand."""
+    if args.autokill:
+        args.kill = "ray|vllm|VLLM|python"
+
+    if args.kill:
+        logger.info(f"Killing processes matching keywords: {args.kill}")
+        for keyword in args.kill.split("|"):
+            logger.info(f"Killing processes matching keyword: {keyword}")
+            killed_pids = fast_kill_by_keyword_bash(keyword)
+            if killed_pids:
+                logger.success(
+                    f"Successfully killed processes with PIDs: {killed_pids}"
+                )
+            else:
+                logger.warning(f"No processes found matching keyword: {keyword}")
+
     # Use default config if not provided
     exp_base_dir = args.exp_dir or DEFAULT_DIR
     if not args.conf:
@@ -126,6 +142,19 @@ def main():
         required=False,
         help="Debug tags; enables Ray post-mortem and DEBUG_TAGS env",
     )
+    parser_start.add_argument(
+        "--kill",
+        type=str,
+        default="",
+        required=False,
+        help="list of keywords for killing processes",
+    )
+    parser_start.add_argument(
+        "--autokill",
+        action="store_true",
+        default=False,
+        help="Kill system processes (ray + vllm + python) that may block the current experiment",
+    )
 
     parser_start.set_defaults(func=cmd_start)
 

ajet/task_rollout/async_llm_bridge.py

@@ -122,7 +122,7 @@ async def llm_chat_verl(
             ):
 
                 parsed_tool_calls = self.tool_parser.extract_tool_calls(decoded_text, None)  # type: ignore
-                parsed_tool_calls = parsed_tool_calls.model_dump()
+                parsed_tool_calls = parsed_tool_calls.model_dump(mode='json')
 
                 model_called = parsed_tool_calls["tools_called"]
                 if model_called:
@@ -155,7 +155,7 @@ async def llm_chat_verl(
                 "completion_tokens": len(token_array), # type: ignore
                 "total_tokens": len(prompt_token_ids) + len(token_array), # type: ignore
             }
-            # from ajet import bp; bp("DECODE")
+
             return {
                 "role": "assistant",
                 "request_id": request_id,
@@ -327,7 +327,7 @@ async def chat_completion_request(
         episode_uuid: str,
     ):
         from openai.types.chat.chat_completion import ChatCompletion
-        req_as_dict = req.model_dump()
+        req_as_dict = req.model_dump(mode='json')
 
         # infer + process with context tracker
         llm_output = await self.run_infer(

ajet/utils/env_service_client/env_client_ng.py

@@ -207,18 +207,20 @@ def call():
                 messages=action,
                 params=params,
             )
-            return resp["data"]
+            data = resp["data"]
+            while "data" in data and "state" not in data:
+                data = data["data"]
+            data["state"] = data["state"][0]
+            return data
 
-        res = retry_call(
+        return retry_call(
             call,
             max_retry=max_retry,
             fail_return=fallback,
             err_prefix="[step]",
             instance_id=instance_id,
             action_name="step",
         )
-        res["state"] = res["state"][0]
-        return res
 
     def evaluate(
         self,

appworld_swarm_results/val_results.md

@@ -1,35 +1,19 @@
 ## Run Appworld AgentScope Agent
 
-### 1. Install Appworld
+### 1. Install and Run Appworld
+
+- Install:
+```
+rm -rf /tmp/pack_all_in_one & wget https://dail-wlcb.oss-cn-wulanchabu.aliyuncs.com/astuner_archive/appworld_pack_v3.tar.gz  &&   tar   -xzf   ./appworld_pack_v3.tar.gz  -C /tmp
 ```
-def install_appworld():
-    # run:
-    # `rm -rf /tmp/pack_all_in_one & wget https://dail-wlcb.oss-cn-wulanchabu.aliyuncs.com/astuner_archive/appworld_pack_v3.tar.gz  &&   tar   -xzf   ./appworld_pack_v3.tar.gz  -C /tmp`
-    import shutil
-
-    if os.path.exists("/tmp/pack_all_in_one"):
-        shutil.rmtree("/tmp/pack_all_in_one")
-    if os.path.exists("./appworld_pack_v3.tar.gz"):
-        os.remove("./appworld_pack_v3.tar.gz")
-    subprocess.run(
-        [
-            "wget",
-            "https://dail-wlcb.oss-cn-wulanchabu.aliyuncs.com/astuner_archive/appworld_pack_v3.tar.gz",
-        ]
-    )
-    subprocess.run(
-        [
-            "tar",
-            "-xzf",
-            "./appworld_pack_v3.tar.gz",
-            "-C",
-            "/tmp",
-        ]
-    )
-    # write
-    os.environ["APPWORLD_PATH"] = "/tmp/pack_all_in_one"
-    os.environ["APPWORLD_SCRIPT"] = "bash EnvService/env_sandbox/appworld.sh"
+
+- Run:
 ```
+export APPWORLD_PATH="/tmp/pack_all_in_one"
+export APPWORLD_SCRIPT="bash EnvService/env_sandbox/appworld.sh"
+ajet --with-appworld --skip-check-avail-gpu
+```
+
 
 
 ### 2. Prepare AgentScope Workflow