Add new API with extensibility and compatibility adapters

miles/ray/rollout.py

@@ -13,7 +13,8 @@
 from sglang.srt.constants import GPU_MEMORY_TYPE_CUDA_GRAPH, GPU_MEMORY_TYPE_KV_CACHE, GPU_MEMORY_TYPE_WEIGHTS
 
 from miles.backends.sglang_utils.sglang_engine import SGLangEngine
-from miles.rollout.base_types import call_rollout_fn
+from miles.rollout.base_types import RolloutFnConstructorInput, RolloutFnEvalInput, RolloutFnTrainInput
+from miles.rollout.modular_rollout.compatibility import call_rollout_function, load_rollout_function
 from miles.utils import tracking_utils
 from miles.utils.health_monitor import RolloutHealthMonitor
 from miles.utils.http_utils import _wrap_ipv6, find_available_port, get_host_info, init_http_client
@@ -53,8 +54,9 @@ def __init__(self, args, pg):
         data_source_cls = load_function(self.args.data_source_path)
         self.data_source = data_source_cls(args)
 
-        self.generate_rollout = load_function(self.args.rollout_function_path)
-        self.eval_generate_rollout = load_function(self.args.eval_function_path)
+        input = RolloutFnConstructorInput(args=args, data_source=self.data_source)
+        self.generate_rollout = load_rollout_function(input, self.args.rollout_function_path)
+        self.eval_generate_rollout = load_rollout_function(input, self.args.eval_function_path)
         self.custom_reward_post_process_func = None
         if self.args.custom_reward_post_process_path is not None:
             self.custom_reward_post_process_func = load_function(self.args.custom_reward_post_process_path)
@@ -142,7 +144,7 @@ def eval(self, rollout_id):
             return
         self.health_monitoring_resume()
 
-        result = call_rollout_fn(self.eval_generate_rollout, self.args, rollout_id, self.data_source, evaluation=True)
+        result = call_rollout_function(self.eval_generate_rollout, RolloutFnEvalInput(rollout_id=rollout_id))
         data = result.data
         self._save_debug_rollout_data(data, rollout_id=rollout_id, evaluation=True)
         metrics = _log_eval_rollout_data(rollout_id, self.args, data, result.metrics)
@@ -224,7 +226,7 @@ def _get_rollout_data(self, rollout_id):
                 )
             metrics = None
         else:
-            data = call_rollout_fn(self.generate_rollout, self.args, rollout_id, self.data_source, evaluation=False)
+            data = call_rollout_function(self.generate_rollout, RolloutFnTrainInput(rollout_id=rollout_id))
             metrics = data.metrics
             data = data.samples
             # flatten the data if it is a list of lists

miles/rollout/base_types.py

@@ -1,9 +1,43 @@
+from argparse import Namespace
+from collections.abc import Awaitable
 from dataclasses import dataclass
-from typing import Any
+from typing import Any, Protocol, runtime_checkable
 
+from miles.rollout.data_source import DataSource
 from miles.utils.types import Sample
 
 
+@dataclass(frozen=True)
+class RolloutFnConstructorInput:
+    args: Namespace
+    # TODO may refactor DataSource API
+    data_source: DataSource
+
+
+@dataclass(frozen=True)
+class RolloutFnBaseInput:
+    rollout_id: int
+
+    @property
+    def evaluation(self):
+        raise NotImplementedError
+
+
+# subclassing for different data in the future
+@dataclass(frozen=True)
+class RolloutFnTrainInput(RolloutFnBaseInput):
+    @property
+    def evaluation(self):
+        return False
+
+
+@dataclass(frozen=True)
+class RolloutFnEvalInput(RolloutFnBaseInput):
+    @property
+    def evaluation(self):
+        return True
+
+
 @dataclass
 class RolloutFnTrainOutput:
     samples: list[list[Sample]]
@@ -16,11 +50,13 @@ class RolloutFnEvalOutput:
     metrics: dict[str, Any] = None
 
 
-def call_rollout_fn(fn, *args, evaluation: bool, **kwargs):
-    output = fn(*args, **kwargs, evaluation=evaluation)
+RolloutFnInput = RolloutFnTrainInput | RolloutFnEvalInput
+RolloutFnOutput = RolloutFnTrainOutput | RolloutFnEvalOutput
 
-    # compatibility for legacy version
-    if not isinstance(output, (RolloutFnTrainOutput, RolloutFnEvalOutput)):
-        output = RolloutFnEvalOutput(data=output) if evaluation else RolloutFnTrainOutput(samples=output)
 
-    return output
+# TODO: may add add_arguments
+# TODO: may add save/load if need it to be stateful
+# Duck typing, users do not need to extend this class
+@runtime_checkable
+class RolloutFnProtocol(Protocol):
+    def __call__(self, input: RolloutFnInput) -> RolloutFnOutput | Awaitable[RolloutFnOutput]: ...

miles/rollout/modular_rollout/compatibility.py

@@ -0,0 +1,50 @@
+import inspect
+from collections.abc import Callable
+
+from miles.rollout.base_types import (
+    RolloutFnConstructorInput,
+    RolloutFnEvalOutput,
+    RolloutFnInput,
+    RolloutFnOutput,
+    RolloutFnProtocol,
+    RolloutFnTrainOutput,
+)
+from miles.utils.async_utils import run
+from miles.utils.misc import load_function
+
+
+class LegacyRolloutFnAdapter:
+    def __init__(self, input: RolloutFnConstructorInput, fn: Callable):
+        self.args = input.args
+        self.data_source = input.data_source
+        self.fn = fn
+
+    def __call__(self, input: RolloutFnInput) -> RolloutFnOutput:
+        output = self.fn(self.args, input.rollout_id, self.data_source, evaluation=input.evaluation)
+
+        # compatibility for legacy version
+        if not isinstance(output, (RolloutFnTrainOutput, RolloutFnEvalOutput)):
+            output = RolloutFnEvalOutput(data=output) if input.evaluation else RolloutFnTrainOutput(samples=output)
+
+        return output
+
+
+assert issubclass(LegacyRolloutFnAdapter, RolloutFnProtocol)
+
+
+def load_rollout_function(input: RolloutFnConstructorInput, path: str):
+    fn = load_function(path)
+
+    if inspect.isclass(fn):
+        return fn(input)
+    else:
+        return LegacyRolloutFnAdapter(input, fn)
+
+
+def call_rollout_function(fn: RolloutFnProtocol, input: RolloutFnInput) -> RolloutFnOutput:
+    output = fn(input)
+
+    if inspect.iscoroutine(output):
+        output = run(output)
+
+    return output

tests/rollout/modular_rollout/test_compatibility.py

@@ -0,0 +1,112 @@
+import asyncio
+from unittest.mock import patch
+
+import pytest
+
+from miles.rollout.base_types import (
+    RolloutFnConstructorInput,
+    RolloutFnEvalInput,
+    RolloutFnEvalOutput,
+    RolloutFnTrainInput,
+    RolloutFnTrainOutput,
+)
+from miles.rollout.modular_rollout.compatibility import (
+    LegacyRolloutFnAdapter,
+    call_rollout_function,
+    load_rollout_function,
+)
+
+
+@pytest.fixture
+def constructor_input():
+    return RolloutFnConstructorInput(args="dummy_args", data_source="dummy_data_source")
+
+
+class TestSupportedRolloutFormats:
+    """
+    Documentation test to show various supported rollout function formats
+    """
+
+    @pytest.mark.parametrize("evaluation", [False, True])
+    def test_format_1_legacy_function_raw_output(self, constructor_input, evaluation):
+        def legacy_rollout_fn(args, rollout_id, data_source, evaluation=False):
+            if evaluation:
+                return {"metric": {"accuracy": 0.9}}
+            return [[{"text": "sample"}]]
+
+        with patch("miles.rollout.modular_rollout.compatibility.load_function", return_value=legacy_rollout_fn):
+            fn = load_rollout_function(constructor_input, "path.to.fn")
+
+        input_cls = RolloutFnEvalInput if evaluation else RolloutFnTrainInput
+        result = call_rollout_function(fn, input_cls(rollout_id=1))
+
+        assert isinstance(fn, LegacyRolloutFnAdapter)
+        if evaluation:
+            assert isinstance(result, RolloutFnEvalOutput)
+            assert result.data == {"metric": {"accuracy": 0.9}}
+        else:
+            assert isinstance(result, RolloutFnTrainOutput)
+            assert result.samples == [[{"text": "sample"}]]
+
+    @pytest.mark.parametrize("evaluation", [False, True])
+    def test_format_2_legacy_function_typed_output(self, constructor_input, evaluation):
+        def legacy_rollout_fn(args, rollout_id, data_source, evaluation=False):
+            if evaluation:
+                return RolloutFnEvalOutput(data={"ds": {"acc": 0.95}})
+            return RolloutFnTrainOutput(samples=[[{"text": "typed"}]])
+
+        with patch("miles.rollout.modular_rollout.compatibility.load_function", return_value=legacy_rollout_fn):
+            fn = load_rollout_function(constructor_input, "path.to.fn")
+
+        input_cls = RolloutFnEvalInput if evaluation else RolloutFnTrainInput
+        result = call_rollout_function(fn, input_cls(rollout_id=1))
+
+        if evaluation:
+            assert isinstance(result, RolloutFnEvalOutput)
+            assert result.data == {"ds": {"acc": 0.95}}
+        else:
+            assert isinstance(result, RolloutFnTrainOutput)
+            assert result.samples == [[{"text": "typed"}]]
+
+    @pytest.mark.parametrize("evaluation", [False, True])
+    def test_format_3_sync_class(self, constructor_input, evaluation):
+        class SyncRolloutFn:
+            def __init__(self, input: RolloutFnConstructorInput):
+                pass
+
+            def __call__(self, input):
+                if input.evaluation:
+                    return RolloutFnEvalOutput(data={"test": {"score": 1}})
+                return RolloutFnTrainOutput(samples=[[{"text": "sync"}]])
+
+        with patch("miles.rollout.modular_rollout.compatibility.load_function", return_value=SyncRolloutFn):
+            fn = load_rollout_function(constructor_input, "path.to.SyncRolloutFn")
+
+        input_cls = RolloutFnEvalInput if evaluation else RolloutFnTrainInput
+        result = call_rollout_function(fn, input_cls(rollout_id=1))
+
+        assert isinstance(fn, SyncRolloutFn)
+        expected_type = RolloutFnEvalOutput if evaluation else RolloutFnTrainOutput
+        assert isinstance(result, expected_type)
+
+    @pytest.mark.parametrize("evaluation", [False, True])
+    def test_format_4_async_class(self, constructor_input, evaluation):
+        class AsyncRolloutFn:
+            def __init__(self, input: RolloutFnConstructorInput):
+                pass
+
+            async def __call__(self, input):
+                await asyncio.sleep(0.001)
+                if input.evaluation:
+                    return RolloutFnEvalOutput(data={"benchmark": {"accuracy": 0.98}})
+                return RolloutFnTrainOutput(samples=[[{"text": "async"}]])
+
+        with patch("miles.rollout.modular_rollout.compatibility.load_function", return_value=AsyncRolloutFn):
+            fn = load_rollout_function(constructor_input, "path.to.AsyncRolloutFn")
+
+        input_cls = RolloutFnEvalInput if evaluation else RolloutFnTrainInput
+        result = call_rollout_function(fn, input_cls(rollout_id=1))
+
+        assert isinstance(fn, AsyncRolloutFn)
+        expected_type = RolloutFnEvalOutput if evaluation else RolloutFnTrainOutput
+        assert isinstance(result, expected_type)