apply review suggestions

Author: ji-huazhong

verl/single_controller/base/decorator.py

@@ -18,16 +18,7 @@
 
 from verl.protocol import DataProtoFuture, _padding_size_key
 from verl.utils.py_functional import DynamicEnum
-
-# TODO: Use a hacky workaround for ImportError since
-# transfer_queue isn't a default verl dependency.
-try:
-    from transfer_queue import BatchMeta
-except ImportError:
-
-    class BatchMeta:
-        pass
-
+from verl.utils.transferqueue_utils import BatchMeta
 
 # here we add a magic number of avoid user-defined function already have this attribute
 MAGIC_ATTR = "attrs_3141562937"

verl/trainer/main_ppo.py

@@ -67,7 +67,9 @@ def run_ppo(config, task_runner_class=None) -> None:
         ray.init(**OmegaConf.to_container(ray_init_kwargs))
 
     if task_runner_class is None:
-        task_runner_class = TaskRunner
+        task_runner_class = ray.remote(TaskRunner).options(
+            num_cpus=1
+        )  # please make sure main_task is not scheduled on head
 
     # Create a remote instance of the TaskRunner class, and
     # Execute the `run` method of the TaskRunner instance remotely and wait for it to complete
@@ -95,7 +97,6 @@ def run_ppo(config, task_runner_class=None) -> None:
         ray.timeline(filename=timeline_json_file)
 
 
-@ray.remote(num_cpus=1)  # please make sure main_task is not scheduled on head
 class TaskRunner:
     """Ray remote class for executing distributed PPO training tasks.
 

verl/utils/transferqueue_utils.py

@@ -31,6 +31,12 @@
 except ImportError:
     HAS_TQ = False
 
+    # TODO: Use a hacky workaround for ImportError since
+    # transfer_queue isn't a default verl dependency.
+    class BatchMeta:
+        pass
+
+
 from verl.protocol import DataProto
 
 _TRANSFER_QUEUE_CLIENT = None
@@ -135,6 +141,26 @@ def _update_batchmeta_with_output(output: DataProto, batchmeta: "BatchMeta") ->
 
 
 def tqbridge(put_data: bool = True):
+    """ "Creates a decorator for bridging BatchMeta and DataProto.
+
+    This decorator automatically handles conversions between `BatchMeta` and
+    `DataProto` in function parameters, and decides whether to sync function
+    output back to `BatchMeta` based on configuration(`put_data`). It supports
+    both synchronous and asynchronous functions (async def), and can control
+    whether to enable enhanced logic via the global `HAS_TQ` variable (when disabled,
+    simply calls the original function as-is).
+
+    Args:
+        put_data: Whether put the DataProto into Storage after func return.
+                  If True, after function execution, the output result will be
+                  updated to `BatchMeta` and `BatchMeta` will be returned;
+                  If False, the function output result will be returned directly.
+                  Defaults to True.
+
+    Returns:
+        A decorator function used to decorate target functions (synchronous or asynchronous).
+    """
+
     def decorator(func):
         @wraps(func)
         def inner(*args, **kwargs):