rename source code

Author: binary-husky

ajet/backbone/trainer_verl.py

@@ -144,6 +144,35 @@ def union_gen_batch_via_task_id(tasks, batch: DataProto, gen_batch_output: DataP
         logger.info(f'task_id_counter: {task_id_counter}')
         return gen_batch_output
 
+def import_or_export_data_proto(batch: DataProto, direction: str = "export", file: str = "./tmp.pkl") -> DataProto:
+    """Import or export a DataProto batch to/from a pickle file.
+
+    Args:
+        batch: The DataProto batch object. Used when direction is "export";
+               ignored (can be None) when direction is "import".
+        direction: "import" to load a batch from file, "export" to save the batch to file.
+        file: Path to the pickle file. Defaults to "./tmp.pkl".
+
+    Returns:
+        The DataProto batch — either the one just loaded (import) or the one just saved (export).
+
+    Raises:
+        ValueError: If direction is not "import" or "export".
+        FileNotFoundError: If direction is "import" and the file does not exist.
+    """
+    import pickle
+    if direction == "export":
+        with open(file, "wb") as f:
+            pickle.dump(batch, f)
+        logger.info(f"[import_or_export_data_proto] Exported batch to {file}")
+        return batch
+    elif direction == "import":
+        with open(file, "rb") as f:
+            batch = pickle.load(f)
+        logger.info(f"[import_or_export_data_proto] Imported batch from {file}")
+        return batch
+    else:
+        raise ValueError(f"direction must be 'import' or 'export', got '{direction}'")
 
 def compute_advantage(
     data: DataProto,

ajet/copilot/job.py

@@ -13,7 +13,7 @@
 
 from typing import Any, Callable, Union, cast
 from loguru import logger
-from ajet.default_config.ajet_default import Config
+from ajet.default_config.ajet_config_schema import Config
 from ajet.utils.config_utils import (
     expand_ajet_hierarchical_config,
     read_ajet_hierarchical_config,
@@ -42,7 +42,7 @@ class AgentJetJob:
     """Programmatic interface for configuring and launching AgentJet training jobs.
 
     Args:
-        base_yaml_config: Path to base YAML configuration file. If None, uses default config (at ./ajet/default_config/ajet_ts_default.yaml).
+        base_yaml_config: Path to base YAML configuration file. If None, uses default config (at ./ajet/default_config/ajet_swarm_default.yaml).
         experiment_dir: Directory where experiment outputs will be saved.
         project_name: Name of the project for organizing experiments.
         experiment_name: Unique name for this specific experiment run.
@@ -86,7 +86,7 @@ def __init__(
     ) -> None:
 
         if base_yaml_config is None:
-            base_yaml_config = os.path.abspath(os.path.join(os.path.dirname(__file__), '..', "default_config/ajet_ts_default.yaml"))
+            base_yaml_config = os.path.abspath(os.path.join(os.path.dirname(__file__), '..', "default_config/ajet_swarm_default.yaml"))
         else:
             logger.warning(f"Reading config from {base_yaml_config}.")
             time.sleep(1)
@@ -121,7 +121,7 @@ def __init__(
         self.max_model_len: int = cast(int, max_model_len)
         self.mini_batch_num: int = cast(int, mini_batch_num)
 
-        # see `ajet/default_config/ajet_ts_default.yaml`
+        # see `ajet/default_config/ajet_swarm_default.yaml`
         overrides = {
             # left: [yaml key navigation]                  right: [AgentJetJob self attr]
             "ajet.experiment_dir":                          "experiment_dir",

ajet/copilot/train-complex-blackbox/SKILL.md

@@ -54,7 +54,7 @@ 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_config_schema import AjetTaskReader, HuggingfaceDatRepo
 from ajet.tuner_lib.experimental.swarm_client import SwarmClient
 
 # python -m tutorial.example_math_swarm.math

ajet/copilot/write-swarm-client/SKILL.md

@@ -365,7 +365,7 @@ Below are some reference materials.
     ```python
     from ajet.copilot.job import AgentJetJob
     from ajet.tuner_lib.experimental.swarm_client import SwarmClient, run_episodes_until_all_complete
-    from ajet.default_config.ajet_default import AjetTaskReader, HuggingfaceDatRepo
+    from ajet.default_config.ajet_config_schema import AjetTaskReader, HuggingfaceDatRepo
     from ajet.task_reader import RouterTaskReader
     from tutorial.example_academic_trans_swarm.trans import execute_agent
 

ajet/default_config/ajet_default.py ajet/default_config/ajet_config_schema.pyajet/default_config/ajet_default.py renamed to ajet/default_config/ajet_config_schema.py

@@ -207,7 +207,7 @@ def main():
     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"
+                os.path.dirname(__file__), "default_config/ajet_swarm_default.yaml"
             )
         )
         assert os.path.exists(args.conf), (

ajet/default_config/ajet_ts_default.yaml …t/default_config/ajet_swarm_default.yamlajet/default_config/ajet_ts_default.yaml renamed to ajet/default_config/ajet_swarm_default.yaml ajet/default_config/ajet_ts_default.yaml renamed to ajet/default_config/ajet_swarm_default.yaml

@@ -46,7 +46,7 @@ def cmd_start(args):
     if not args.conf:
         args.conf = os.path.abspath(
             os.path.join(
-                os.path.dirname(__file__), "default_config/ajet_ts_default.yaml"
+                os.path.dirname(__file__), "default_config/ajet_swarm_default.yaml"
             )
         )
         assert os.path.exists(args.conf), (

ajet/launcher.py

@@ -133,7 +133,7 @@ Hint: you do not have to use `run_episodes_until_all_complete`, you are free to
 ```python
 from ajet.copilot.job import AgentJetJob
 from ajet.tuner_lib.experimental.swarm_client import SwarmClient, run_episodes_until_all_complete
-from ajet.default_config.ajet_default import AjetTaskReader, HuggingfaceDatRepo
+from ajet.default_config.ajet_config_schema import AjetTaskReader, HuggingfaceDatRepo
 from ajet.task_reader import RouterTaskReader
 from tutorial.example_academic_trans_swarm.trans import execute_agent
 

ajet/swarm_cli.py

@@ -496,7 +496,7 @@ Create your client script. The client reads the dataset, runs the agent workflow
     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_config_schema import AjetTaskReader, HuggingfaceDatRepo
     from ajet.tuner_lib.experimental.swarm_client import SwarmClient
 
     # Configuration
@@ -649,7 +649,7 @@ The server handles gradient computation and model updates automatically.
     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_config_schema import AjetTaskReader, HuggingfaceDatRepo
     from ajet.tuner_lib.experimental.swarm_client import SwarmClient
 
     GRPO_N = 4  # grpo group size