feat: enhance TinkerScript integration with improved engine status handling and configuration updates

Author: binary-husky

ajet/backbone/main_verl.py

@@ -22,7 +22,6 @@
 import hydra
 import ray
 from beast_logger import print_dict
-from loguru import logger
 from omegaconf import OmegaConf
 from verl.trainer.ppo.reward import load_reward_manager
 from verl.utils.device import is_cuda_available
@@ -110,6 +109,7 @@ def run(self, config):
         # Print the initial configuration. `resolve=True` will evaluate symbolic values.
         from pprint import pprint
 
+        from loguru import logger
         from omegaconf import OmegaConf
         from verl.utils.fs import copy_to_local
 

ajet/backbone/main_vllm.py

@@ -191,7 +191,7 @@ def main(config):
         start_interchange_server(config)
         if config.ajet.enable_tinkerscript_mode:
             from ajet.tuner_lib.weight_tuner.experimental.interchange_utils import http_change_engine_status
-            http_change_engine_status(config, "ROLLING")
+            http_change_engine_status(config, "ENGINE.ROLLING")
 
     def companion_launch():
         import torch

ajet/backbone/trainer_verl.py

@@ -559,7 +559,7 @@ def fit(self):  # noqa: C901
                         assert self.async_rollout_mode
                         logger.info("=== wake up begin ===")
                         self.async_rollout_manager.wake_up()
-                        self._update_interchange_server_status_flag("ROLLING")
+                        self._update_interchange_server_status_flag("ENGINE.ROLLING")
                         logger.info("=== wake up end ===")
                         tasks: List[Task] = [
                             dict_to_ajet_task(dict(
@@ -585,7 +585,7 @@ def fit(self):  # noqa: C901
                             tasks, mode="sample", epoch=f"train.{epoch}"
                         )
                         logger.info("=" * 10 + "end fit rollout" + "=" * 10)
-                        self._update_interchange_server_status_flag("UPDATE_WEIGHT")
+                        self._update_interchange_server_status_flag("ENGINE.WEIGHT_SYNCING")
                         logger.info("begin to convert context_tracker_arr to dataproto")
                         gen_batch_output = self.parallel_env.to_dataproto(context_tracker_arr)
                         logger.info("end convertion")

ajet/copilot/job.py

@@ -42,16 +42,23 @@ def __init__(
         n_gpu: int = 8,
         algorithm: str = "grpo",
         n_gpu_for_infer: int | None = None, # only for trinity backbone
+        grpo_n: int = 8,
+        tinkerscript_mode: bool = True,
         *kwargs,
     ) -> None:
         self.backbone = backbone
-        self.config_as_dict: dict = self.build_job_from_yaml(None)
+        if tinkerscript_mode:
+            default_yaml = os.path.abspath(os.path.join(os.path.dirname(__file__), '..', "default_config/ajet_ts_default.yaml"))
+        else:
+            default_yaml = None
+        self.config_as_dict: dict = self.build_job_from_yaml(default_yaml)
         self.config = Config.update_from_dict_recursive(Config(), self.config_as_dict)
 
         self.config.ajet.backbone = backbone
         self.config.ajet.model.path = model
         self.config.ajet.trainer_common.n_gpus_per_node = n_gpu
         self.config.ajet.trainer_common.algorithm.adv_estimator = algorithm
+        self.config.ajet.rollout.num_repeat = grpo_n
         if n_gpu_for_infer is None and backbone == "trinity":
             raise ValueError("Please specify `n_gpu_for_infer` (n_gpu_for_infer < n_gpu) for trinity backbone.")
         if (n_gpu_for_infer is not None) and backbone == "verl":

ajet/default_config/ajet_default.py

@@ -30,6 +30,7 @@ class AjetRollout:
     user_workflow: str = "tutorial.example_appworld.appworld->ExampleAgentScopeWorkflow"
     n_vllm_engine: int = 1
     tensor_model_parallel_size: int = 1
+    num_repeat: int = 8
 
 
 @dataclass

ajet/default_config/ajet_default.yaml

@@ -292,6 +292,7 @@ ajet:
     num_fastapi_process: 2  # 1, 2 or 4 is fine
     max_fastapi_threads: 128  # 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`
 
 
   task_runner: