stage

Author: binary-husky

.pre-commit-config.yaml

@@ -11,12 +11,12 @@ repos:
       - id: check-merge-conflict
       - id: detect-private-key
 
-  # - repo: https://github.com/psf/black
-  #   rev: 23.7.0
-  #   hooks:
-  #   - id: black
-  #     language_version: python3.10
-  #     args: [--line-length=100]
+  - repo: https://github.com/psf/black
+    rev: 23.7.0
+    hooks:
+    - id: black
+      language_version: python3.10
+      args: [--line-length=999999]
 
   # - repo: https://github.com/pycqa/isort
   #   rev: 5.12.0

ajet/__init__.py

@@ -4,13 +4,6 @@
 from ajet.workflow import Workflow
 from ajet.utils.vsdb import vscode_conditional_breakpoint as bp
 
-__all__ = [
-    "Workflow",
-    "WorkflowTask",
-    "WorkflowOutput",
-    "AjetTuner",
-    "AgentJetJob",
-    "bp"
-]
+__all__ = ["Workflow", "WorkflowTask", "WorkflowOutput", "AjetTuner", "AgentJetJob", "bp"]
 
 __version__ = "0.1.0"

ajet/backbone/main_trinity.py

@@ -54,6 +54,7 @@ def patched_trainer_get_actor(cls, config: Config):
 
     if ajet_config.ajet.enable_experimental_interchange_server:
         from ajet.tuner_lib.weight_tuner.experimental.as_oai_model_server import start_interchange_server
+
         start_interchange_server(ajet_config)
 
 

ajet/backbone/main_verl.py

@@ -66,16 +66,10 @@ def run_ppo(config) -> None:
 
     # Create a remote instance of the TaskRunner class, and
     # Execute the `run` method of the TaskRunner instance remotely and wait for it to complete
-    if (
-        is_cuda_available
-        and config.trainer.get("profile_steps") is not None
-        and len(config.trainer.get("profile_steps", [])) > 0
-    ):
+    if is_cuda_available and config.trainer.get("profile_steps") is not None and len(config.trainer.get("profile_steps", [])) > 0:
         from verl.utils.import_utils import is_nvtx_available
 
-        assert (
-            is_nvtx_available()
-        ), "nvtx is not available in CUDA platform. Please 'pip3 install nvtx'"
+        assert is_nvtx_available(), "nvtx is not available in CUDA platform. Please 'pip3 install nvtx'"
         nsight_options = OmegaConf.to_container(config.trainer.controller_nsight_options)
         runner = TaskRunner.options(runtime_env={"nsight": nsight_options}).remote()
     else:
@@ -223,9 +217,7 @@ def run(self, config):
             num_examine=1,
             **config.reward_model.get("reward_kwargs", {}),
         )
-        resource_pool_manager = ResourcePoolManager(
-            resource_pool_spec=resource_pool_spec, mapping=mapping
-        )
+        resource_pool_manager = ResourcePoolManager(resource_pool_spec=resource_pool_spec, mapping=mapping)
 
         from verl.utils.dataset.rl_dataset import collate_fn
 
@@ -248,6 +240,7 @@ def run(self, 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)
 
         # Initialize the PPO trainer.

ajet/backbone/main_vllm.py

@@ -82,9 +82,7 @@ def submit_chat_completions(self, messages, sampling_params, request_id, tools=[
                 "request_id": completion.id,
                 "content": message["content"],
                 "tool_calls": message.get("tool_calls", None),
-                "tokens": [
-                    TokenAndProbVllmDebug(t) for t in completion.choices[0].logprobs.content # type: ignore
-                ],
+                "tokens": [TokenAndProbVllmDebug(t) for t in completion.choices[0].logprobs.content],  # type: ignore
             }
         )
         return messages
@@ -130,13 +128,12 @@ async def submit_chat_completions_async(self, messages, sampling_params, request
                 "request_id": completion.id,
                 "content": message["content"],
                 "tool_calls": message.get("tool_calls", None),
-                "tokens": [
-                    TokenAndProbVllmDebug(t) for t in completion.choices[0].logprobs.content # type: ignore
-                ],
+                "tokens": [TokenAndProbVllmDebug(t) for t in completion.choices[0].logprobs.content],  # type: ignore
             }
         )
         return messages
 
+
 def run(config):
     from ajet.task_reader import RouterTaskReader
 
@@ -147,9 +144,7 @@ def run(config):
     vllm_port = config.ajet.debug.debug_vllm_port
 
     # --------- init ---------
-    async_rollout_manager = ChatCompletionScheduler(
-        config=config, url=f"http://localhost:{vllm_port}/v1"
-    )
+    async_rollout_manager = ChatCompletionScheduler(config=config, url=f"http://localhost:{vllm_port}/v1")
     parallel_env = VerlRolloutManager(
         config=config,
         async_rollout_manager=async_rollout_manager,
@@ -159,16 +154,13 @@ def run(config):
         tokenizer=async_rollout_manager.tokenizer,
     )
 
-
     task_reader = RouterTaskReader(
         config.ajet.task_reader.type,
         config.ajet.task_reader,
     )
     tasks = task_reader.get_validation_tasks()
     logger.info(tasks[:n_task])
-    ctx_tracker = parallel_env.rollout(
-        tasks=tasks[:n_task], mode="sample", epoch="1"
-    )  # "sample" or "validate"
+    ctx_tracker = parallel_env.rollout(tasks=tasks[:n_task], mode="sample", epoch="1")  # "sample" or "validate"
     _ = parallel_env.to_dataproto(ctx_tracker)
 
 
@@ -179,13 +171,15 @@ def run(config):
 )
 def main(config):
     from omegaconf import OmegaConf
+
     OmegaConf.resolve(config)
     runtime_env = get_runtime_env(config)
     os.environ.update(runtime_env["env_vars"])
     # atexit.register(lambda: print("Process exiting, performing cleanup..."))
 
     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)
 
     def companion_launch():
@@ -198,9 +192,7 @@ def companion_launch():
         tensor_parallel_size = config.ajet.debug.debug_tensor_parallel_size
         n_avail_gpus = torch.cuda.device_count()
         if tensor_parallel_size > n_avail_gpus:
-            logger.info(
-                f"Warning: tensor_parallel_size {tensor_parallel_size} is greater than available GPUs {n_avail_gpus}. Setting tensor_parallel_size to {n_avail_gpus}."
-            )
+            logger.info(f"Warning: tensor_parallel_size {tensor_parallel_size} is greater than available GPUs {n_avail_gpus}. Setting tensor_parallel_size to {n_avail_gpus}.")
             tensor_parallel_size = n_avail_gpus
         gpu_memory_utilization = config.actor_rollout_ref.rollout.gpu_memory_utilization
         max_num_seqs = config.actor_rollout_ref.rollout.max_num_seqs

ajet/backbone/trainer_trinity.py

@@ -66,6 +66,7 @@ def __init__(
 
     def convert_task(self, task: TrinityTask):
         from ajet.schema.task import Task
+
         assert isinstance(task.raw_task, dict)
         return dict_to_ajet_task(task.raw_task)
 
@@ -150,16 +151,10 @@ async def run_async(self):
                 "madness": tracker.reward_structure.madness,
             }
 
-            if (
-                len(response_ids) + len(prompt_ids) == len(input_ids)
-                and len(logprobs) == len(response_ids)
-                and len(logprobs) > 0
-            ):
+            if len(response_ids) + len(prompt_ids) == len(input_ids) and len(logprobs) == len(response_ids) and len(logprobs) > 0:
                 exp = Experience(
                     tokens=input_ids,  # [seq_length] prompt + response
-                    prompt_length=len(
-                        prompt_ids
-                    ),  # Length of the prompt in tokens, used for generating attention masks
+                    prompt_length=len(prompt_ids),  # Length of the prompt in tokens, used for generating attention masks
                     logprobs=logprobs,  # [resp_length]
                     reward=reward,  #
                     # advantages=None,
@@ -211,19 +206,11 @@ def __init__(self, config):
             if "train" in self.split:
                 dataset_segments.append(task_to_standard_dataset(task_reader.get_training_tasks()))
             if "val" in self.split:
-                dataset_segments.append(
-                    task_to_standard_dataset(task_reader.get_validation_tasks())
-                )
+                dataset_segments.append(task_to_standard_dataset(task_reader.get_validation_tasks()))
             if not dataset_segments:
-                raise ValueError(
-                    f"Unsupported split '{self.split}'. Expected to contain 'train' or 'val'."
-                )
+                raise ValueError(f"Unsupported split '{self.split}'. Expected to contain 'train' or 'val'.")
 
-            concatenated_dataset = (
-                dataset_segments[0]
-                if len(dataset_segments) == 1
-                else datasets.concatenate_datasets(dataset_segments)
-            )
+            concatenated_dataset = dataset_segments[0] if len(dataset_segments) == 1 else datasets.concatenate_datasets(dataset_segments)
 
             self.dataset = _HFBatchReader(
                 concatenated_dataset,
@@ -271,15 +258,9 @@ class SwanlabMonitor(Monitor):
     """
 
     def __init__(self, project: str, group: str, name: str, role: str, config) -> None:
-        assert (
-            swanlab is not None
-        ), "swanlab is not installed. Please install it to use SwanlabMonitor."
-
-        monitor_args = (
-            (config.monitor.monitor_args or {})
-            if config and getattr(config, "monitor", None)
-            else {}
-        )
+        assert swanlab is not None, "swanlab is not installed. Please install it to use SwanlabMonitor."
+
+        monitor_args = (config.monitor.monitor_args or {}) if config and getattr(config, "monitor", None) else {}
 
         # Optional API login via code if provided; otherwise try environment, then rely on prior `swanlab login`.
         api_key = os.environ.get("SWANLAB_API_KEY")
@@ -331,9 +312,7 @@ def __init__(self, project: str, group: str, name: str, role: str, config) -> No
         self.data_dashboard_url = run_info["cloud"]["experiment_url"]
 
     def log_table(self, table_name: str, experiences_table, step: int):
-        assert (
-            swanlab is not None
-        ), "swanlab is not installed. Please install it to use SwanlabMonitor."
+        assert swanlab is not None, "swanlab is not installed. Please install it to use SwanlabMonitor."
 
         # Convert pandas DataFrame to SwanLab ECharts Table
         headers: List[str] = list(experiences_table.columns)
@@ -351,9 +330,7 @@ def log_table(self, table_name: str, experiences_table, step: int):
     def log(self, data: dict, step: int, commit: bool = False) -> None:
         """Log metrics."""
         # SwanLab doesn't use commit flag; keep signature for compatibility
-        assert (
-            swanlab is not None
-        ), "swanlab is not installed. Please install it to use SwanlabMonitor."
+        assert swanlab is not None, "swanlab is not installed. Please install it to use SwanlabMonitor."
         swanlab.log(data, step=step)
         self.console_logger.info(f"Step {step}: {data}")
 
@@ -372,9 +349,7 @@ def log(self, data: dict, step: int, commit: bool = False) -> None:
             test_robot_data = {}
             test_robot_data["step"] = step
             test_robot_data["data_dashboard_url"] = self.data_dashboard_url
-            test_robot_data["reward_for_test_robot"] = data[
-                "experience_pipeline/group_advantages/reward_mean/mean"
-            ]
+            test_robot_data["reward_for_test_robot"] = data["experience_pipeline/group_advantages/reward_mean/mean"]
             _test_if_test_mode(key="reward_probe", value=test_robot_data, config=ajet_config)
 
     def close(self) -> None: