[misc] single_controller: support assign blocking value in runtime

verl/single_controller/base/decorator.py

@@ -13,6 +13,7 @@
 # limitations under the License.
 
 import inspect
+import logging
 from functools import wraps
 from types import FunctionType
 from typing import Dict, List, Tuple
@@ -507,7 +508,15 @@ async def async_inner(*args, **kwargs):
             return await func(*args, **kwargs)
 
         wrapper = async_inner if inspect.iscoroutinefunction(func) else inner
-        attrs = {"dispatch_mode": dispatch_mode, "execute_mode": execute_mode, "blocking": blocking}
+
+        blocking_key = "blocking_v2"
+        sig = inspect.signature(func)
+        if "blocking" in sig.parameters:
+            logging.warning(f"method {func.__name__} hack \'blocking\' keyword in parameters, falling back to old mode")
+            blocking_key = "blocking"
+
+        attrs = {"dispatch_mode": dispatch_mode, "execute_mode": execute_mode, blocking_key: blocking}
+
         setattr(wrapper, MAGIC_ATTR, attrs)
         return wrapper
 

verl/single_controller/base/worker_group.py

@@ -162,7 +162,8 @@ def _bind_worker_method(self, user_defined_cls, func_generator):
 
                 dispatch_mode = attribute["dispatch_mode"]
                 execute_mode = attribute["execute_mode"]
-                blocking = attribute["blocking"]
+                blocking = attribute.get("blocking", None)
+                blocking_v2 = attribute.get("blocking_v2", None)
 
                 # get dispatch fn
                 if isinstance(dispatch_mode, Dispatch):
@@ -197,6 +198,7 @@ def _bind_worker_method(self, user_defined_cls, func_generator):
                     collect_fn=collect_fn,
                     execute_fn=execute_fn,
                     blocking=blocking,
+                    blocking_v2=blocking_v2
                 )
 
                 try:

verl/single_controller/ray/base.py

@@ -40,12 +40,17 @@ def get_random_string(length: int) -> str:
     return "".join(random.choice(letters_digits) for _ in range(length))
 
 
-def func_generator(self, method_name, dispatch_fn, collect_fn, execute_fn, blocking):
+def func_generator(self, method_name, dispatch_fn, collect_fn, execute_fn, blocking, blocking_v2):
+    if blocking is not None:
+        assert blocking_v2 is None, f"blocking and blocking_v2 cannot be not None at the same time"
     def func(*args, **kwargs):
         args, kwargs = dispatch_fn(self, *args, **kwargs)
         padding_count = kwargs.pop(_padding_size_key, 0)
+        blocking_effect = blocking
+        if blocking_v2 is not None:
+            blocking_effect = kwargs.pop("blocking_v2", blocking_v2)
         output = execute_fn(method_name, *args, **kwargs)
-        if blocking:
+        if blocking_effect:
             output = ray.get(output)
         output = collect_fn(self, output)
         if padding_count > 0: