rename to agentjet swarm

Author: binary-husky

ajet/backbone/main_vllm.py

@@ -144,7 +144,7 @@ def run(config):
     max_parallel = config.ajet.debug.debug_max_parallel
     n_task = config.ajet.debug.debug_first_n_tasks
     vllm_port = config.ajet.debug.debug_vllm_port
-    enable_tinkerscript_mode = config.ajet.enable_tinkerscript_mode
+    enable_swarm_mode = config.ajet.enable_swarm_mode
 
     # --------- init ---------
     async_rollout_manager = ChatCompletionScheduler(
@@ -168,7 +168,7 @@ def run(config):
     logger.info(tasks[:n_task])
     ctx_tracker = parallel_env.rollout(
         tasks=tasks[:n_task],
-        mode="sample" if not enable_tinkerscript_mode else "sample-ts",  # type: ignore
+        mode="sample" if not enable_swarm_mode else "sample-ts",  # type: ignore
         epoch="1"
     )
     _ = parallel_env.to_dataproto(ctx_tracker)
@@ -189,7 +189,7 @@ def main(config):
     if config.ajet.enable_experimental_interchange_server:
         from ajet.tuner_lib.weight_tuner.experimental.as_oai_model_server import start_interchange_server
         start_interchange_server(config)
-        if config.ajet.enable_tinkerscript_mode:
+        if config.ajet.enable_swarm_mode:
             from ajet.tuner_lib.weight_tuner.experimental.interchange_utils import http_change_engine_status
             http_change_engine_status(config, "ENGINE.ROLLING")
 

ajet/backbone/trainer_verl.py

@@ -458,7 +458,7 @@ def init_workers(self):
 
     def _update_interchange_server_status_flag(self, status: str):
         if self.config.ajet.enable_experimental_interchange_server:
-            if self.config.ajet.enable_tinkerscript_mode:
+            if self.config.ajet.enable_swarm_mode:
                 from ajet.tuner_lib.weight_tuner.experimental.interchange_utils import http_change_engine_status
                 http_change_engine_status(self.config, status)
 
@@ -493,7 +493,7 @@ def fit(self):  # noqa: C901
 
         # perform validation before training
         # currently, we only support validation using the reward_function.
-        if (self.val_reward_fn is not None) and (self.config.trainer.get("val_before_train", True)) and (not self.config.ajet.enable_tinkerscript_mode):
+        if (self.val_reward_fn is not None) and (self.config.trainer.get("val_before_train", True)) and (not self.config.ajet.enable_swarm_mode):
             val_metrics = self._validate()
             assert val_metrics, f"{val_metrics=}"
             pprint(f"Initial validation metrics: {val_metrics}")
@@ -651,7 +651,7 @@ def fit(self):  # noqa: C901
                         [str(uuid.uuid4()) for _ in range(len(batch.batch))],
                         dtype=object,
                     )
-                    discard_original_batch = self.config.ajet.enable_tinkerscript_mode
+                    discard_original_batch = self.config.ajet.enable_swarm_mode
                     batch = union_gen_batch_via_task_id(tasks, batch, gen_batch_output, discard_original_batch)
                     batch.batch["response_mask"] = compute_response_mask(batch)
 
@@ -784,7 +784,7 @@ def fit(self):  # noqa: C901
                         self.val_reward_fn is not None
                         and self.config.trainer.test_freq > 0
                         and (is_last_step or self.global_steps % self.config.trainer.test_freq == 0)
-                        and (not self.config.ajet.enable_tinkerscript_mode)
+                        and (not self.config.ajet.enable_swarm_mode)
                     ):
                         with marked_timer("testing", timing_raw, color="green"):
                             val_metrics: dict = self._validate()
@@ -958,7 +958,7 @@ def _validate(self):
                 dtype=object,
             )
             tasks = tasks[: len(main_val_dataset)]
-            discard_original_batch = self.config.ajet.enable_tinkerscript_mode
+            discard_original_batch = self.config.ajet.enable_swarm_mode
             test_batch = union_gen_batch_via_task_id(tasks, test_batch, test_output_gen_batch, discard_original_batch)
             # test_batch = test_batch.union(test_output_gen_batch)
             test_batch.meta_info["validate"] = True

ajet/context_tracker/multiagent_tracking.py

@@ -48,14 +48,14 @@ def __init__(
         self,
         tokenizer: PreTrainedTokenizer,
         config,
-        should_interrupt_fn,
+        should_interrupt_soft_fn,
         should_interrupt_hard_fn,
         generated_token_callback_fn,
         **kwargs,
     ):
         super().__init__(config, tokenizer, **kwargs)
         self.tokenizer = tokenizer
-        self.should_interrupt_fn = should_interrupt_fn
+        self.should_interrupt_soft_fn = should_interrupt_soft_fn
         self.should_interrupt_hard_fn = should_interrupt_hard_fn
         self.generated_token_callback_fn = generated_token_callback_fn
         self.context_overflow = False
@@ -601,7 +601,7 @@ def check_context_token_num_safe(
             token_overflow = False
         else:
             token_overflow = True
-        if self.should_interrupt_fn():
+        if self.should_interrupt_soft_fn():
             ret = (False, token_overflow, "externally_interrupted")
         elif self.already_mad_flag and self.config.ajet.rollout.agent_madness_termination:
             ret = (False, token_overflow, "already_mad")

ajet/copilot/job.py

@@ -45,11 +45,11 @@ def __init__(
         n_gpu_for_infer: int | None = None, # only for trinity backbone
         grpo_n: int = 8,
         batch_size: int = 32,
-        tinkerscript_mode: bool = True,
+        swarm_mode: bool = True,
         *kwargs,
     ) -> None:
         self.backbone = backbone
-        if tinkerscript_mode:
+        if swarm_mode:
             default_yaml = os.path.abspath(os.path.join(os.path.dirname(__file__), '..', "default_config/ajet_ts_default.yaml"))
         else:
             default_yaml = None

ajet/default_config/ajet_default.yaml

@@ -282,8 +282,8 @@ ajet:
 
 
   # the experimental ZeroMQ interchange server feature that allows `tuner.as_oai_baseurl_apikey` feature
-  enable_tinkerscript_mode: False
-  # both tinkerscript / oai share the same interchange server
+  enable_swarm_mode: False
+  # both swarm / oai share the same interchange server
   enable_experimental_interchange_server: False
   # interchange server configuration
   interchange_server:
@@ -292,7 +292,7 @@ ajet:
     num_fastapi_process: 2  # 1, 2 or 4 is fine
     max_fastapi_threads: 512  # 64 or 128 is fine
     max_inference_tracker_threads: 64 # recommend to be equal to `ajet.rollout.max_env_worker`
-    already_started: False # do not edit, used by `tinkerscript`
+    already_started: False # do not edit, used by `swarm`
 
 
   task_runner:

ajet/default_config/ajet_ts_default.yaml

@@ -23,15 +23,15 @@ ajet:
   # the experimental ZeroMQ interchange server feature that allows `tuner.as_oai_baseurl_apikey` feature
   enable_experimental_interchange_server: True
   # train in cloud, run episode locally
-  enable_tinkerscript_mode: True
-  # both tinkerscript / oai share the same interchange server
+  enable_swarm_mode: True
+  # both swarm / oai share the same interchange server
   interchange_server:
     interchange_method: 'ipc' # options: 'tcp' (multi-nodes) or  'ipc' (1 node)
     interchange_server_port: 10086
     num_fastapi_process: 2  # 1, 2 or 4 is fine
     max_fastapi_threads: 512  # 64 or 128 is fine
     max_inference_tracker_threads: 64 # recommend to be equal to `ajet.rollout.max_env_worker`
-    already_started: False # do not edit, used by `tinkerscript`
+    already_started: False # do not edit, used by `swarm`
 
   rollout:
     # maximum number of parallel environments / simulate workers

ajet/launcher.py

@@ -35,10 +35,10 @@ def parse_args():
         help="verl or trinity or debug",
     )
     parser.add_argument(
-        "--tinkerscript-server",
+        "--swarm-server",
         action="store_true",
         default=False,
-        help="Enable TinkerScript server mode",
+        help="Enable Swarm server mode",
     )
     parser.add_argument(
         "--conf",
@@ -146,12 +146,12 @@ def check_model_file_exists(exp_config):
         assert os.path.exists(model_path), f"Model path {model_path} does not exist. Please check your configuration."
 
 
-def start_tinkerscript_server(env, config):
+def start_swarm_server(env, config):
     config = dict_to_namespace(config)
-    assert config.ajet.enable_tinkerscript_mode, \
-        "Please enable_tinkerscript_mode in config to start tinkerscript server."
+    assert config.ajet.enable_swarm_mode, \
+        "Please enable_swarm_mode in config to start swarm server."
     assert config.ajet.enable_experimental_interchange_server, \
-        "Please enable_experimental_interchange_server in config to start tinkerscript server."
+        "Please enable_experimental_interchange_server in config to start swarm server."
     from ajet.tuner_lib.weight_tuner.experimental.as_oai_model_server import start_interchange_server
     start_interchange_server(config, blocking=True, env=env)
 
@@ -191,9 +191,9 @@ def main():
     # read configuration from yaml
     exp_config = None
     exp_dir = args.exp_dir or "saved_experiments"
-    if args.tinkerscript_server and (not args.conf):
+    if args.swarm_server and (not args.conf):
         args.conf = os.path.abspath(os.path.join(os.path.dirname(__file__), "default_config/ajet_ts_default.yaml"))
-        assert os.path.exists(args.conf), "Please provide a valid config file for tinkerscript server mode."
+        assert os.path.exists(args.conf), "Please provide a valid config file for swarm server mode."
     if args.conf:
         yaml_path = args.conf
         (
@@ -206,8 +206,8 @@ def main():
     # setup environment variables
     env, exp_config = setup_environment_vars(args, exp_config, main_yaml_fp)
 
-    if args.tinkerscript_server:
-        start_tinkerscript_server(env, exp_config)
+    if args.swarm_server:
+        start_swarm_server(env, exp_config)
         return
 
     if args.with_ray:

ajet/task_reader/document_reader/doc_reader.py

@@ -11,7 +11,7 @@
 try:
     from unstructured.partition.auto import partition
 except Exception:
-    logger.warning("Cannot import dependency `unstructured`")
+    logger.info("`unstructured` is not installed.")
 
 from ajet.schema.document import Document
 from ajet.task_reader.document_reader.document_reader_base import (

ajet/task_rollout/native_parallel_worker.py

@@ -144,7 +144,7 @@ def rollout(
         epoch: str,
     ) -> List[BaseContextTracker]:
         """Delegate to dynamic rollout when oversampling is enabled."""
-        if self.config.ajet.enable_tinkerscript_mode:
+        if self.config.ajet.enable_swarm_mode:
             return self.rollout_swarm(tasks, mode, epoch)
         elif (
             mode == "sample"

ajet/task_rollout/single_worker.py

@@ -15,7 +15,7 @@
 from ajet.task_rollout.async_llm_bridge import AsyncLlmBridge
 from ajet.task_rollout.resource_keeper import ResourceKeeper
 from ajet.task_runner.general_runner import GeneralRunner
-from ajet.task_runner.tinkerscript_runner import TinkerScriptRunner
+from ajet.task_runner.swarm_runner import SwarmRunner
 from ajet.utils.retry import retry_with_backoff
 from ajet.utils.retry import SwarmReceiveAbortException
 from ajet.utils.sample import get_sample_params
@@ -64,7 +64,7 @@ def __init__(
         assert isinstance(self.pad_token_id, int), "pad_token_id must be an integer"
         self.current_token = 0
         self.current_global_steps: int | str = "NA"
-        self.enable_tinkerscript_mode = config.ajet.enable_tinkerscript_mode
+        self.enable_swarm_mode = config.ajet.enable_swarm_mode
         self.async_llm_bridge = AsyncLlmBridge(
             config=config,
             async_rollout_manager=async_rollout_manager,
@@ -116,8 +116,8 @@ def rollout_env_worker(
         with ResourceKeeper(workflow_task, config=self.config) as resource_keeper:
             try:
                 workflow_task = resource_keeper.prepare()
-                if self.enable_tinkerscript_mode:
-                    agent_runner = TinkerScriptRunner(
+                if self.enable_swarm_mode:
+                    agent_runner = SwarmRunner(
                         llm_inference_fn=llm_inference_fn, tokenizer=self.tokenizer, config=self.config
                     )
                 else: