From 3a5f4516bdd1522ef5098cb36b965acaa540d97e Mon Sep 17 00:00:00 2001 From: fzyzcjy Date: Wed, 14 Jan 2026 15:45:18 +0800 Subject: [PATCH 01/46] more --- miles/rollout/modular_rollout/__init__.py | 0 miles/rollout/modular_rollout/compatibility.py | 0 tests/integration/__init__.py | 0 tests/unit/__init__.py | 0 4 files changed, 0 insertions(+), 0 deletions(-) create mode 100644 miles/rollout/modular_rollout/__init__.py create mode 100644 miles/rollout/modular_rollout/compatibility.py create mode 100644 tests/integration/__init__.py create mode 100644 tests/unit/__init__.py diff --git a/miles/rollout/modular_rollout/__init__.py b/miles/rollout/modular_rollout/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/miles/rollout/modular_rollout/compatibility.py b/miles/rollout/modular_rollout/compatibility.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/integration/__init__.py b/tests/integration/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/__init__.py b/tests/unit/__init__.py new file mode 100644 index 00000000000..e69de29bb2d From 1f0903394d1311f58c6c40490161a0aa190d9b55 Mon Sep 17 00:00:00 2001 From: fzyzcjy Date: Wed, 14 Jan 2026 15:46:49 +0800 Subject: [PATCH 02/46] more --- miles/rollout/base_types.py | 25 +++++++++++++++++++++++++ 1 file changed, 25 insertions(+) diff --git a/miles/rollout/base_types.py b/miles/rollout/base_types.py index faa85c72697..f5cc07cb90b 100644 --- a/miles/rollout/base_types.py +++ b/miles/rollout/base_types.py @@ -1,8 +1,33 @@ +from argparse import Namespace from dataclasses import dataclass from typing import Any from miles.utils.types import Sample +@dataclass +class RolloutFnBaseInput: + args: Namespace + rollout_id: int + data_source: Any + + @property + def evaluation(self): + raise NotImplementedError + + +@dataclass +class RolloutFnTrainInput(RolloutFnBaseInput): + @property + def evaluation(self): + return False + + +@dataclass +class RolloutFnEvalInput(RolloutFnBaseInput): + @property + def evaluation(self): + return True + @dataclass class RolloutFnTrainOutput: From 717c3835d557f9135585764a7479ddf29094bb29 Mon Sep 17 00:00:00 2001 From: fzyzcjy Date: Wed, 14 Jan 2026 15:47:15 +0800 Subject: [PATCH 03/46] fmt --- miles/rollout/base_types.py | 1 + 1 file changed, 1 insertion(+) diff --git a/miles/rollout/base_types.py b/miles/rollout/base_types.py index f5cc07cb90b..6f2b216df83 100644 --- a/miles/rollout/base_types.py +++ b/miles/rollout/base_types.py @@ -4,6 +4,7 @@ from miles.utils.types import Sample + @dataclass class RolloutFnBaseInput: args: Namespace From b03bfb6e63a35315a7e6d35b457229fc5f2e5620 Mon Sep 17 00:00:00 2001 From: fzyzcjy Date: Wed, 14 Jan 2026 15:47:23 +0800 Subject: [PATCH 04/46] more --- miles/rollout/base_types.py | 1 - 1 file changed, 1 deletion(-) diff --git a/miles/rollout/base_types.py b/miles/rollout/base_types.py index 6f2b216df83..81e9270efc9 100644 --- a/miles/rollout/base_types.py +++ b/miles/rollout/base_types.py @@ -9,7 +9,6 @@ class RolloutFnBaseInput: args: Namespace rollout_id: int - data_source: Any @property def evaluation(self): From 7860261be8c2c372f5dc7abceccfb121f8ac589b Mon Sep 17 00:00:00 2001 From: fzyzcjy Date: Wed, 14 Jan 2026 15:47:57 +0800 Subject: [PATCH 05/46] more --- miles/rollout/base_types.py | 1 + 1 file changed, 1 insertion(+) diff --git a/miles/rollout/base_types.py b/miles/rollout/base_types.py index 81e9270efc9..2780752e291 100644 --- a/miles/rollout/base_types.py +++ b/miles/rollout/base_types.py @@ -41,6 +41,7 @@ class RolloutFnEvalOutput: metrics: dict[str, Any] = None +# TODO move / refactor def call_rollout_fn(fn, *args, evaluation: bool, **kwargs): output = fn(*args, **kwargs, evaluation=evaluation) From a0c4035d8669f52f1edbef03d00c63a96618fef0 Mon Sep 17 00:00:00 2001 From: fzyzcjy Date: Wed, 14 Jan 2026 15:48:27 +0800 Subject: [PATCH 06/46] more --- miles/rollout/base_types.py | 1 + 1 file changed, 1 insertion(+) diff --git a/miles/rollout/base_types.py b/miles/rollout/base_types.py index 2780752e291..17e72a57aed 100644 --- a/miles/rollout/base_types.py +++ b/miles/rollout/base_types.py @@ -15,6 +15,7 @@ def evaluation(self): raise NotImplementedError +# subclassing for different data in the future @dataclass class RolloutFnTrainInput(RolloutFnBaseInput): @property From ead54942b57a8bde7cd547682428b37d4699011c Mon Sep 17 00:00:00 2001 From: fzyzcjy Date: Wed, 14 Jan 2026 15:50:21 +0800 Subject: [PATCH 07/46] more --- miles/rollout/base_types.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/miles/rollout/base_types.py b/miles/rollout/base_types.py index 17e72a57aed..7cc54dd53b0 100644 --- a/miles/rollout/base_types.py +++ b/miles/rollout/base_types.py @@ -5,7 +5,7 @@ from miles.utils.types import Sample -@dataclass +@dataclass(frozen=True) class RolloutFnBaseInput: args: Namespace rollout_id: int @@ -16,14 +16,14 @@ def evaluation(self): # subclassing for different data in the future -@dataclass +@dataclass(frozen=True) class RolloutFnTrainInput(RolloutFnBaseInput): @property def evaluation(self): return False -@dataclass +@dataclass(frozen=True) class RolloutFnEvalInput(RolloutFnBaseInput): @property def evaluation(self): From 788d848baade2fb8bc63b081df86fcc8dde2f2d1 Mon Sep 17 00:00:00 2001 From: fzyzcjy Date: Wed, 14 Jan 2026 15:50:33 +0800 Subject: [PATCH 08/46] more --- miles/rollout/base_types.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/miles/rollout/base_types.py b/miles/rollout/base_types.py index 7cc54dd53b0..c90ecfaad72 100644 --- a/miles/rollout/base_types.py +++ b/miles/rollout/base_types.py @@ -1,4 +1,3 @@ -from argparse import Namespace from dataclasses import dataclass from typing import Any @@ -7,7 +6,6 @@ @dataclass(frozen=True) class RolloutFnBaseInput: - args: Namespace rollout_id: int @property From 54982392693ade9358777b620365c9e971931250 Mon Sep 17 00:00:00 2001 From: fzyzcjy Date: Wed, 14 Jan 2026 15:51:06 +0800 Subject: [PATCH 09/46] more --- miles/rollout/base_types.py | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/miles/rollout/base_types.py b/miles/rollout/base_types.py index c90ecfaad72..cffa2e75951 100644 --- a/miles/rollout/base_types.py +++ b/miles/rollout/base_types.py @@ -1,9 +1,16 @@ +from argparse import Namespace from dataclasses import dataclass from typing import Any from miles.utils.types import Sample +@dataclass(frozen=True) +class RolloutFnConstructorInput: + args: Namespace + data_source: Any + + @dataclass(frozen=True) class RolloutFnBaseInput: rollout_id: int From c9c1ed1d8ca181d5120ec7369475c65286a348ee Mon Sep 17 00:00:00 2001 From: fzyzcjy Date: Wed, 14 Jan 2026 15:51:55 +0800 Subject: [PATCH 10/46] more --- miles/rollout/base_types.py | 10 +++++++++- 1 file changed, 9 insertions(+), 1 deletion(-) diff --git a/miles/rollout/base_types.py b/miles/rollout/base_types.py index cffa2e75951..7cdfa36fd80 100644 --- a/miles/rollout/base_types.py +++ b/miles/rollout/base_types.py @@ -1,6 +1,6 @@ from argparse import Namespace from dataclasses import dataclass -from typing import Any +from typing import Any, Protocol from miles.utils.types import Sample @@ -47,6 +47,14 @@ class RolloutFnEvalOutput: metrics: dict[str, Any] = None +class RolloutFnProtocol(Protocol): + def __init__(self, input: RolloutFnConstructorInput): + ... + + def __call__(self, input: RolloutFnTrainInput | RolloutFnEvalInput) -> RolloutFnTrainOutput | RolloutFnEvalOutput: + ... + + # TODO move / refactor def call_rollout_fn(fn, *args, evaluation: bool, **kwargs): output = fn(*args, **kwargs, evaluation=evaluation) From db7dbf29e4bf1967ed1852f9df616c1e106ad5be Mon Sep 17 00:00:00 2001 From: fzyzcjy Date: Wed, 14 Jan 2026 15:53:07 +0800 Subject: [PATCH 11/46] more --- miles/rollout/base_types.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/miles/rollout/base_types.py b/miles/rollout/base_types.py index 7cdfa36fd80..d0c1446e800 100644 --- a/miles/rollout/base_types.py +++ b/miles/rollout/base_types.py @@ -1,4 +1,4 @@ -from argparse import Namespace +from argparse import Namespace, ArgumentParser from dataclasses import dataclass from typing import Any, Protocol @@ -51,6 +51,10 @@ class RolloutFnProtocol(Protocol): def __init__(self, input: RolloutFnConstructorInput): ... + @classmethod + def add_arguments(cls, parser: ArgumentParser): + ... + def __call__(self, input: RolloutFnTrainInput | RolloutFnEvalInput) -> RolloutFnTrainOutput | RolloutFnEvalOutput: ... From 0de80bbd3da5ce0a0ff649d3dea62493add33661 Mon Sep 17 00:00:00 2001 From: fzyzcjy Date: Wed, 14 Jan 2026 15:54:19 +0800 Subject: [PATCH 12/46] more --- miles/rollout/base_types.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/miles/rollout/base_types.py b/miles/rollout/base_types.py index d0c1446e800..e3e8c6cb10d 100644 --- a/miles/rollout/base_types.py +++ b/miles/rollout/base_types.py @@ -2,13 +2,14 @@ from dataclasses import dataclass from typing import Any, Protocol +from miles.rollout.data_source import DataSource from miles.utils.types import Sample @dataclass(frozen=True) class RolloutFnConstructorInput: args: Namespace - data_source: Any + data_source: DataSource @dataclass(frozen=True) From 2c368f6da7ca1098b65bd664ca9a8da63e86dca1 Mon Sep 17 00:00:00 2001 From: fzyzcjy Date: Wed, 14 Jan 2026 15:54:43 +0800 Subject: [PATCH 13/46] more --- miles/rollout/base_types.py | 1 + 1 file changed, 1 insertion(+) diff --git a/miles/rollout/base_types.py b/miles/rollout/base_types.py index e3e8c6cb10d..f6f79f53a18 100644 --- a/miles/rollout/base_types.py +++ b/miles/rollout/base_types.py @@ -9,6 +9,7 @@ @dataclass(frozen=True) class RolloutFnConstructorInput: args: Namespace + # TODO may refactor DataSource API data_source: DataSource From 4a504b138e4389beb7d558a257c3cac218e145a8 Mon Sep 17 00:00:00 2001 From: fzyzcjy Date: Wed, 14 Jan 2026 15:56:16 +0800 Subject: [PATCH 14/46] more --- miles/rollout/base_types.py | 1 + 1 file changed, 1 insertion(+) diff --git a/miles/rollout/base_types.py b/miles/rollout/base_types.py index f6f79f53a18..3c0d6fca79a 100644 --- a/miles/rollout/base_types.py +++ b/miles/rollout/base_types.py @@ -49,6 +49,7 @@ class RolloutFnEvalOutput: metrics: dict[str, Any] = None +# Duck typing, users do not need to extend this class class RolloutFnProtocol(Protocol): def __init__(self, input: RolloutFnConstructorInput): ... From f474d84745151b7c5f429020cf26e92aa10031c5 Mon Sep 17 00:00:00 2001 From: fzyzcjy Date: Wed, 14 Jan 2026 15:56:24 +0800 Subject: [PATCH 15/46] fmt --- miles/rollout/base_types.py | 13 ++++++------- 1 file changed, 6 insertions(+), 7 deletions(-) diff --git a/miles/rollout/base_types.py b/miles/rollout/base_types.py index 3c0d6fca79a..f30b66171ff 100644 --- a/miles/rollout/base_types.py +++ b/miles/rollout/base_types.py @@ -1,4 +1,4 @@ -from argparse import Namespace, ArgumentParser +from argparse import ArgumentParser, Namespace from dataclasses import dataclass from typing import Any, Protocol @@ -51,15 +51,14 @@ class RolloutFnEvalOutput: # Duck typing, users do not need to extend this class class RolloutFnProtocol(Protocol): - def __init__(self, input: RolloutFnConstructorInput): - ... + def __init__(self, input: RolloutFnConstructorInput): ... @classmethod - def add_arguments(cls, parser: ArgumentParser): - ... + def add_arguments(cls, parser: ArgumentParser): ... - def __call__(self, input: RolloutFnTrainInput | RolloutFnEvalInput) -> RolloutFnTrainOutput | RolloutFnEvalOutput: - ... + def __call__( + self, input: RolloutFnTrainInput | RolloutFnEvalInput + ) -> RolloutFnTrainOutput | RolloutFnEvalOutput: ... # TODO move / refactor From 041cbfc4258b655bb93a4bcf8aae497597191d0f Mon Sep 17 00:00:00 2001 From: fzyzcjy Date: Wed, 14 Jan 2026 15:57:37 +0800 Subject: [PATCH 16/46] more --- miles/rollout/modular_rollout/compatibility.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/miles/rollout/modular_rollout/compatibility.py b/miles/rollout/modular_rollout/compatibility.py index e69de29bb2d..5306526e5b8 100644 --- a/miles/rollout/modular_rollout/compatibility.py +++ b/miles/rollout/modular_rollout/compatibility.py @@ -0,0 +1,2 @@ +class LegacyRolloutFnAdapter: + TODO From 5228f5460543c75f7798d12a858c59af81981824 Mon Sep 17 00:00:00 2001 From: fzyzcjy Date: Wed, 14 Jan 2026 15:58:27 +0800 Subject: [PATCH 17/46] more --- miles/rollout/base_types.py | 8 +++++--- miles/rollout/modular_rollout/compatibility.py | 6 +++++- 2 files changed, 10 insertions(+), 4 deletions(-) diff --git a/miles/rollout/base_types.py b/miles/rollout/base_types.py index f30b66171ff..e79acf99f1f 100644 --- a/miles/rollout/base_types.py +++ b/miles/rollout/base_types.py @@ -49,6 +49,10 @@ class RolloutFnEvalOutput: metrics: dict[str, Any] = None +RolloutFnInput = RolloutFnTrainInput | RolloutFnEvalInput +RolloutFnOutput = RolloutFnTrainOutput | RolloutFnEvalOutput + + # Duck typing, users do not need to extend this class class RolloutFnProtocol(Protocol): def __init__(self, input: RolloutFnConstructorInput): ... @@ -56,9 +60,7 @@ def __init__(self, input: RolloutFnConstructorInput): ... @classmethod def add_arguments(cls, parser: ArgumentParser): ... - def __call__( - self, input: RolloutFnTrainInput | RolloutFnEvalInput - ) -> RolloutFnTrainOutput | RolloutFnEvalOutput: ... + def __call__(self, input: RolloutFnInput) -> RolloutFnOutput: ... # TODO move / refactor diff --git a/miles/rollout/modular_rollout/compatibility.py b/miles/rollout/modular_rollout/compatibility.py index 5306526e5b8..4e270aec9d9 100644 --- a/miles/rollout/modular_rollout/compatibility.py +++ b/miles/rollout/modular_rollout/compatibility.py @@ -1,2 +1,6 @@ +from miles.rollout.base_types import RolloutFnInput, RolloutFnOutput + + class LegacyRolloutFnAdapter: - TODO + def __call__(self, input: RolloutFnInput) -> RolloutFnOutput: + TODO From b2772d7ef2e74dd7b6a1314e99f3113ae23c0510 Mon Sep 17 00:00:00 2001 From: fzyzcjy Date: Wed, 14 Jan 2026 16:00:06 +0800 Subject: [PATCH 18/46] more --- miles/rollout/modular_rollout/compatibility.py | 16 ++++++++++++++-- 1 file changed, 14 insertions(+), 2 deletions(-) diff --git a/miles/rollout/modular_rollout/compatibility.py b/miles/rollout/modular_rollout/compatibility.py index 4e270aec9d9..5a5573f0858 100644 --- a/miles/rollout/modular_rollout/compatibility.py +++ b/miles/rollout/modular_rollout/compatibility.py @@ -1,6 +1,18 @@ -from miles.rollout.base_types import RolloutFnInput, RolloutFnOutput +from miles.rollout.base_types import RolloutFnInput, RolloutFnOutput, RolloutFnConstructorInput, RolloutFnTrainOutput, \ + RolloutFnEvalOutput class LegacyRolloutFnAdapter: + def __init__(self, input: RolloutFnConstructorInput): + self.args = input.args + self.data_source = input.data_source + self.fn = TODO + def __call__(self, input: RolloutFnInput) -> RolloutFnOutput: - TODO + 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 From 85123b1340bcf8bb3b81bca685ee97d1c204d060 Mon Sep 17 00:00:00 2001 From: fzyzcjy Date: Wed, 14 Jan 2026 16:00:19 +0800 Subject: [PATCH 19/46] fmt --- miles/rollout/modular_rollout/compatibility.py | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/miles/rollout/modular_rollout/compatibility.py b/miles/rollout/modular_rollout/compatibility.py index 5a5573f0858..f19f4ce4b98 100644 --- a/miles/rollout/modular_rollout/compatibility.py +++ b/miles/rollout/modular_rollout/compatibility.py @@ -1,5 +1,10 @@ -from miles.rollout.base_types import RolloutFnInput, RolloutFnOutput, RolloutFnConstructorInput, RolloutFnTrainOutput, \ - RolloutFnEvalOutput +from miles.rollout.base_types import ( + RolloutFnConstructorInput, + RolloutFnEvalOutput, + RolloutFnInput, + RolloutFnOutput, + RolloutFnTrainOutput, +) class LegacyRolloutFnAdapter: From 4ab723dc9c27de9d305a02783d5b6c2ce23e00e3 Mon Sep 17 00:00:00 2001 From: fzyzcjy Date: Wed, 14 Jan 2026 16:00:34 +0800 Subject: [PATCH 20/46] more --- miles/rollout/modular_rollout/compatibility.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/miles/rollout/modular_rollout/compatibility.py b/miles/rollout/modular_rollout/compatibility.py index f19f4ce4b98..cf2c9bde0af 100644 --- a/miles/rollout/modular_rollout/compatibility.py +++ b/miles/rollout/modular_rollout/compatibility.py @@ -1,3 +1,5 @@ +from typing import Callable + from miles.rollout.base_types import ( RolloutFnConstructorInput, RolloutFnEvalOutput, @@ -8,10 +10,10 @@ class LegacyRolloutFnAdapter: - def __init__(self, input: RolloutFnConstructorInput): + def __init__(self, input: RolloutFnConstructorInput, fn: Callable): self.args = input.args self.data_source = input.data_source - self.fn = TODO + self.fn = fn def __call__(self, input: RolloutFnInput) -> RolloutFnOutput: output = self.fn(self.args, input.rollout_id, self.data_source, evaluation=input.evaluation) From 34113016795bc0e6487e522c4e4814a649ded1b4 Mon Sep 17 00:00:00 2001 From: fzyzcjy Date: Wed, 14 Jan 2026 16:02:26 +0800 Subject: [PATCH 21/46] more --- miles/ray/rollout.py | 6 +++--- miles/rollout/base_types.py | 11 ----------- 2 files changed, 3 insertions(+), 14 deletions(-) diff --git a/miles/ray/rollout.py b/miles/ray/rollout.py index 79c6649be04..973232220e6 100644 --- a/miles/ray/rollout.py +++ b/miles/ray/rollout.py @@ -13,7 +13,7 @@ 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 call_rollout_fn, RolloutFnEvalInput, RolloutFnTrainInput 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 @@ -142,7 +142,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 = 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 +224,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 = self.generate_rollout(RolloutFnTrainInput(rollout_id=rollout_id)) metrics = data.metrics data = data.samples # flatten the data if it is a list of lists diff --git a/miles/rollout/base_types.py b/miles/rollout/base_types.py index e79acf99f1f..ffc81869b3c 100644 --- a/miles/rollout/base_types.py +++ b/miles/rollout/base_types.py @@ -61,14 +61,3 @@ def __init__(self, input: RolloutFnConstructorInput): ... def add_arguments(cls, parser: ArgumentParser): ... def __call__(self, input: RolloutFnInput) -> RolloutFnOutput: ... - - -# TODO move / refactor -def call_rollout_fn(fn, *args, evaluation: bool, **kwargs): - output = fn(*args, **kwargs, evaluation=evaluation) - - # compatibility for legacy version - if not isinstance(output, (RolloutFnTrainOutput, RolloutFnEvalOutput)): - output = RolloutFnEvalOutput(data=output) if evaluation else RolloutFnTrainOutput(samples=output) - - return output From 11e49cd39def75ee37b41fbbd934c0341cdc648c Mon Sep 17 00:00:00 2001 From: fzyzcjy Date: Wed, 14 Jan 2026 16:03:34 +0800 Subject: [PATCH 22/46] more --- miles/ray/rollout.py | 5 +++-- miles/rollout/modular_rollout/compatibility.py | 6 ++++++ 2 files changed, 9 insertions(+), 2 deletions(-) diff --git a/miles/ray/rollout.py b/miles/ray/rollout.py index 973232220e6..dd4d5a75576 100644 --- a/miles/ray/rollout.py +++ b/miles/ray/rollout.py @@ -26,6 +26,7 @@ from miles.utils.seqlen_balancing import get_seqlen_balanced_partitions from miles.utils.tracking_utils import init_tracking from miles.utils.types import Sample +from miles.rollout.modular_rollout.compatibility import load_rollout_function from ..utils.metric_utils import has_repetition from .utils import NOSET_VISIBLE_DEVICES_ENV_VARS_LIST, Lock @@ -53,8 +54,8 @@ 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) + self.generate_rollout = load_rollout_function(self.args.rollout_function_path) + self.eval_generate_rollout = load_rollout_function(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) diff --git a/miles/rollout/modular_rollout/compatibility.py b/miles/rollout/modular_rollout/compatibility.py index cf2c9bde0af..fbe1a9f4063 100644 --- a/miles/rollout/modular_rollout/compatibility.py +++ b/miles/rollout/modular_rollout/compatibility.py @@ -7,6 +7,7 @@ RolloutFnOutput, RolloutFnTrainOutput, ) +from miles.utils.misc import load_function class LegacyRolloutFnAdapter: @@ -23,3 +24,8 @@ def __call__(self, input: RolloutFnInput) -> RolloutFnOutput: output = RolloutFnEvalOutput(data=output) if input.evaluation else RolloutFnTrainOutput(samples=output) return output + + +def load_rollout_function(path): + fn = load_function(path) + return TODO From 32caeaac3b7a715869560883f3e52c8f3c18c46a Mon Sep 17 00:00:00 2001 From: fzyzcjy Date: Wed, 14 Jan 2026 16:03:41 +0800 Subject: [PATCH 23/46] more --- miles/ray/rollout.py | 4 ++-- miles/rollout/modular_rollout/compatibility.py | 2 +- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/miles/ray/rollout.py b/miles/ray/rollout.py index dd4d5a75576..fe802eda38e 100644 --- a/miles/ray/rollout.py +++ b/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, RolloutFnEvalInput, RolloutFnTrainInput +from miles.rollout.base_types import RolloutFnEvalInput, RolloutFnTrainInput +from miles.rollout.modular_rollout.compatibility import 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 @@ -26,7 +27,6 @@ from miles.utils.seqlen_balancing import get_seqlen_balanced_partitions from miles.utils.tracking_utils import init_tracking from miles.utils.types import Sample -from miles.rollout.modular_rollout.compatibility import load_rollout_function from ..utils.metric_utils import has_repetition from .utils import NOSET_VISIBLE_DEVICES_ENV_VARS_LIST, Lock diff --git a/miles/rollout/modular_rollout/compatibility.py b/miles/rollout/modular_rollout/compatibility.py index fbe1a9f4063..be3e1129a7e 100644 --- a/miles/rollout/modular_rollout/compatibility.py +++ b/miles/rollout/modular_rollout/compatibility.py @@ -1,4 +1,4 @@ -from typing import Callable +from collections.abc import Callable from miles.rollout.base_types import ( RolloutFnConstructorInput, From 874c3aea1a2c1b03dab28bf336d22f059500005d Mon Sep 17 00:00:00 2001 From: fzyzcjy Date: Wed, 14 Jan 2026 16:06:02 +0800 Subject: [PATCH 24/46] more --- miles/ray/rollout.py | 7 ++++--- miles/rollout/modular_rollout/compatibility.py | 9 +++++++-- 2 files changed, 11 insertions(+), 5 deletions(-) diff --git a/miles/ray/rollout.py b/miles/ray/rollout.py index fe802eda38e..b4df5d8fb83 100644 --- a/miles/ray/rollout.py +++ b/miles/ray/rollout.py @@ -13,7 +13,7 @@ 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 RolloutFnEvalInput, RolloutFnTrainInput +from miles.rollout.base_types import RolloutFnEvalInput, RolloutFnTrainInput, RolloutFnConstructorInput from miles.rollout.modular_rollout.compatibility import load_rollout_function from miles.utils import tracking_utils from miles.utils.health_monitor import RolloutHealthMonitor @@ -54,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_rollout_function(self.args.rollout_function_path) - self.eval_generate_rollout = load_rollout_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) diff --git a/miles/rollout/modular_rollout/compatibility.py b/miles/rollout/modular_rollout/compatibility.py index be3e1129a7e..a454eb8507d 100644 --- a/miles/rollout/modular_rollout/compatibility.py +++ b/miles/rollout/modular_rollout/compatibility.py @@ -1,4 +1,5 @@ from collections.abc import Callable +import inspect from miles.rollout.base_types import ( RolloutFnConstructorInput, @@ -26,6 +27,10 @@ def __call__(self, input: RolloutFnInput) -> RolloutFnOutput: return output -def load_rollout_function(path): +def load_rollout_function(input: RolloutFnConstructorInput, path: str): fn = load_function(path) - return TODO + + if not inspect.isclass(fn): + fn = LegacyRolloutFnAdapter(input, fn) + + return fn From 6c5fb233a1d33907893097fb3ecddf5d104de278 Mon Sep 17 00:00:00 2001 From: fzyzcjy Date: Wed, 14 Jan 2026 16:06:53 +0800 Subject: [PATCH 25/46] more --- tests/integration/__init__.py | 0 tests/unit/__init__.py | 0 2 files changed, 0 insertions(+), 0 deletions(-) delete mode 100644 tests/integration/__init__.py delete mode 100644 tests/unit/__init__.py diff --git a/tests/integration/__init__.py b/tests/integration/__init__.py deleted file mode 100644 index e69de29bb2d..00000000000 diff --git a/tests/unit/__init__.py b/tests/unit/__init__.py deleted file mode 100644 index e69de29bb2d..00000000000 From 23295c6d56ede382a0134899f018fe9dc813e009 Mon Sep 17 00:00:00 2001 From: fzyzcjy Date: Wed, 14 Jan 2026 16:07:08 +0800 Subject: [PATCH 26/46] more --- tests/rollout/__init__.py | 0 tests/rollout/modular_rollout/__init__.py | 0 tests/rollout/modular_rollout/test_compatibility.py | 0 3 files changed, 0 insertions(+), 0 deletions(-) create mode 100644 tests/rollout/__init__.py create mode 100644 tests/rollout/modular_rollout/__init__.py create mode 100644 tests/rollout/modular_rollout/test_compatibility.py diff --git a/tests/rollout/__init__.py b/tests/rollout/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/rollout/modular_rollout/__init__.py b/tests/rollout/modular_rollout/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/rollout/modular_rollout/test_compatibility.py b/tests/rollout/modular_rollout/test_compatibility.py new file mode 100644 index 00000000000..e69de29bb2d From d6e8ce48f7ff4a974e24776e54c24b96b317bd6d Mon Sep 17 00:00:00 2001 From: fzyzcjy Date: Wed, 14 Jan 2026 16:08:24 +0800 Subject: [PATCH 27/46] more --- miles/rollout/base_types.py | 1 + 1 file changed, 1 insertion(+) diff --git a/miles/rollout/base_types.py b/miles/rollout/base_types.py index ffc81869b3c..e7842c977d7 100644 --- a/miles/rollout/base_types.py +++ b/miles/rollout/base_types.py @@ -53,6 +53,7 @@ class RolloutFnEvalOutput: RolloutFnOutput = RolloutFnTrainOutput | RolloutFnEvalOutput +# TODO: may add save/load if need it to be stateful # Duck typing, users do not need to extend this class class RolloutFnProtocol(Protocol): def __init__(self, input: RolloutFnConstructorInput): ... From b7e8e31c5ccc8937f327d455d867f29f14ebbf05 Mon Sep 17 00:00:00 2001 From: fzyzcjy Date: Wed, 14 Jan 2026 16:09:45 +0800 Subject: [PATCH 28/46] fmt --- miles/ray/rollout.py | 2 +- miles/rollout/modular_rollout/compatibility.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/miles/ray/rollout.py b/miles/ray/rollout.py index b4df5d8fb83..1a1c40e5436 100644 --- a/miles/ray/rollout.py +++ b/miles/ray/rollout.py @@ -13,7 +13,7 @@ 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 RolloutFnEvalInput, RolloutFnTrainInput, RolloutFnConstructorInput +from miles.rollout.base_types import RolloutFnConstructorInput, RolloutFnEvalInput, RolloutFnTrainInput from miles.rollout.modular_rollout.compatibility import load_rollout_function from miles.utils import tracking_utils from miles.utils.health_monitor import RolloutHealthMonitor diff --git a/miles/rollout/modular_rollout/compatibility.py b/miles/rollout/modular_rollout/compatibility.py index a454eb8507d..21239590926 100644 --- a/miles/rollout/modular_rollout/compatibility.py +++ b/miles/rollout/modular_rollout/compatibility.py @@ -1,5 +1,5 @@ -from collections.abc import Callable import inspect +from collections.abc import Callable from miles.rollout.base_types import ( RolloutFnConstructorInput, From 9dfbe87939e3e221859443c8b90bfcf6a175b676 Mon Sep 17 00:00:00 2001 From: fzyzcjy Date: Wed, 14 Jan 2026 16:10:22 +0800 Subject: [PATCH 29/46] more --- .../modular_rollout/test_compatibility.py | 85 +++++++++++++++++++ 1 file changed, 85 insertions(+) diff --git a/tests/rollout/modular_rollout/test_compatibility.py b/tests/rollout/modular_rollout/test_compatibility.py index e69de29bb2d..60c06ed743c 100644 --- a/tests/rollout/modular_rollout/test_compatibility.py +++ b/tests/rollout/modular_rollout/test_compatibility.py @@ -0,0 +1,85 @@ +from unittest.mock import MagicMock, patch + +import pytest + +from miles.rollout.base_types import ( + RolloutFnConstructorInput, + RolloutFnEvalInput, + RolloutFnEvalOutput, + RolloutFnTrainInput, + RolloutFnTrainOutput, +) +from miles.rollout.modular_rollout.compatibility import ( + LegacyRolloutFnAdapter, + load_rollout_function, +) + + +@pytest.fixture +def constructor_input(): + return RolloutFnConstructorInput(args="dummy_args", data_source="dummy_data_source") + + +class TestLoadRolloutFunction: + def test_load_class(self, constructor_input): + class MockRolloutClass: + pass + + with patch("miles.rollout.modular_rollout.compatibility.load_function", return_value=MockRolloutClass): + result = load_rollout_function(constructor_input, "some.module.MockRolloutClass") + + assert result is MockRolloutClass + + def test_load_function_returns_adapter(self, constructor_input): + def mock_fn(): + pass + + with patch("miles.rollout.modular_rollout.compatibility.load_function", return_value=mock_fn): + result = load_rollout_function(constructor_input, "some.module.mock_fn") + + assert isinstance(result, LegacyRolloutFnAdapter) + assert result.fn is mock_fn + assert result.args == "dummy_args" + assert result.data_source == "dummy_data_source" + + +class TestLegacyRolloutFnAdapter: + def test_call_with_train_input_wraps_output(self, constructor_input): + mock_samples = [[{"text": "sample"}]] + mock_fn = MagicMock(return_value=mock_samples) + adapter = LegacyRolloutFnAdapter(constructor_input, mock_fn) + + result = adapter(RolloutFnTrainInput(rollout_id=1)) + + mock_fn.assert_called_once_with("dummy_args", 1, "dummy_data_source", evaluation=False) + assert isinstance(result, RolloutFnTrainOutput) + assert result.samples == mock_samples + + def test_call_with_eval_input_wraps_output(self, constructor_input): + mock_data = {"metric": {"accuracy": 0.9}} + mock_fn = MagicMock(return_value=mock_data) + adapter = LegacyRolloutFnAdapter(constructor_input, mock_fn) + + result = adapter(RolloutFnEvalInput(rollout_id=2)) + + mock_fn.assert_called_once_with("dummy_args", 2, "dummy_data_source", evaluation=True) + assert isinstance(result, RolloutFnEvalOutput) + assert result.data == mock_data + + def test_passthrough_train_output(self, constructor_input): + expected_output = RolloutFnTrainOutput(samples=[[]]) + mock_fn = MagicMock(return_value=expected_output) + adapter = LegacyRolloutFnAdapter(constructor_input, mock_fn) + + result = adapter(RolloutFnTrainInput(rollout_id=0)) + + assert result is expected_output + + def test_passthrough_eval_output(self, constructor_input): + expected_output = RolloutFnEvalOutput(data={}) + mock_fn = MagicMock(return_value=expected_output) + adapter = LegacyRolloutFnAdapter(constructor_input, mock_fn) + + result = adapter(RolloutFnEvalInput(rollout_id=0)) + + assert result is expected_output From 84dd548b7c3074ec888eb83706b97b36495775c0 Mon Sep 17 00:00:00 2001 From: fzyzcjy Date: Wed, 14 Jan 2026 16:12:09 +0800 Subject: [PATCH 30/46] more --- .../modular_rollout/test_compatibility.py | 92 ++++++++++--------- 1 file changed, 47 insertions(+), 45 deletions(-) diff --git a/tests/rollout/modular_rollout/test_compatibility.py b/tests/rollout/modular_rollout/test_compatibility.py index 60c06ed743c..0c27efe6c56 100644 --- a/tests/rollout/modular_rollout/test_compatibility.py +++ b/tests/rollout/modular_rollout/test_compatibility.py @@ -20,66 +20,68 @@ def constructor_input(): return RolloutFnConstructorInput(args="dummy_args", data_source="dummy_data_source") -class TestLoadRolloutFunction: - def test_load_class(self, constructor_input): - class MockRolloutClass: - pass +def test_load_class(constructor_input): + class MockRolloutClass: + pass - with patch("miles.rollout.modular_rollout.compatibility.load_function", return_value=MockRolloutClass): - result = load_rollout_function(constructor_input, "some.module.MockRolloutClass") + with patch("miles.rollout.modular_rollout.compatibility.load_function", return_value=MockRolloutClass): + result = load_rollout_function(constructor_input, "some.module.MockRolloutClass") - assert result is MockRolloutClass + assert result is MockRolloutClass - def test_load_function_returns_adapter(self, constructor_input): - def mock_fn(): - pass - with patch("miles.rollout.modular_rollout.compatibility.load_function", return_value=mock_fn): - result = load_rollout_function(constructor_input, "some.module.mock_fn") +def test_load_function_returns_adapter(constructor_input): + def mock_fn(): + pass - assert isinstance(result, LegacyRolloutFnAdapter) - assert result.fn is mock_fn - assert result.args == "dummy_args" - assert result.data_source == "dummy_data_source" + with patch("miles.rollout.modular_rollout.compatibility.load_function", return_value=mock_fn): + result = load_rollout_function(constructor_input, "some.module.mock_fn") + assert isinstance(result, LegacyRolloutFnAdapter) + assert result.fn is mock_fn + assert result.args == "dummy_args" + assert result.data_source == "dummy_data_source" -class TestLegacyRolloutFnAdapter: - def test_call_with_train_input_wraps_output(self, constructor_input): - mock_samples = [[{"text": "sample"}]] - mock_fn = MagicMock(return_value=mock_samples) - adapter = LegacyRolloutFnAdapter(constructor_input, mock_fn) - result = adapter(RolloutFnTrainInput(rollout_id=1)) +def test_adapter_call_with_train_input_wraps_output(constructor_input): + mock_samples = [[{"text": "sample"}]] + mock_fn = MagicMock(return_value=mock_samples) + adapter = LegacyRolloutFnAdapter(constructor_input, mock_fn) - mock_fn.assert_called_once_with("dummy_args", 1, "dummy_data_source", evaluation=False) - assert isinstance(result, RolloutFnTrainOutput) - assert result.samples == mock_samples + result = adapter(RolloutFnTrainInput(rollout_id=1)) - def test_call_with_eval_input_wraps_output(self, constructor_input): - mock_data = {"metric": {"accuracy": 0.9}} - mock_fn = MagicMock(return_value=mock_data) - adapter = LegacyRolloutFnAdapter(constructor_input, mock_fn) + mock_fn.assert_called_once_with("dummy_args", 1, "dummy_data_source", evaluation=False) + assert isinstance(result, RolloutFnTrainOutput) + assert result.samples == mock_samples - result = adapter(RolloutFnEvalInput(rollout_id=2)) - mock_fn.assert_called_once_with("dummy_args", 2, "dummy_data_source", evaluation=True) - assert isinstance(result, RolloutFnEvalOutput) - assert result.data == mock_data +def test_adapter_call_with_eval_input_wraps_output(constructor_input): + mock_data = {"metric": {"accuracy": 0.9}} + mock_fn = MagicMock(return_value=mock_data) + adapter = LegacyRolloutFnAdapter(constructor_input, mock_fn) - def test_passthrough_train_output(self, constructor_input): - expected_output = RolloutFnTrainOutput(samples=[[]]) - mock_fn = MagicMock(return_value=expected_output) - adapter = LegacyRolloutFnAdapter(constructor_input, mock_fn) + result = adapter(RolloutFnEvalInput(rollout_id=2)) - result = adapter(RolloutFnTrainInput(rollout_id=0)) + mock_fn.assert_called_once_with("dummy_args", 2, "dummy_data_source", evaluation=True) + assert isinstance(result, RolloutFnEvalOutput) + assert result.data == mock_data - assert result is expected_output - def test_passthrough_eval_output(self, constructor_input): - expected_output = RolloutFnEvalOutput(data={}) - mock_fn = MagicMock(return_value=expected_output) - adapter = LegacyRolloutFnAdapter(constructor_input, mock_fn) +def test_adapter_passthrough_train_output(constructor_input): + expected_output = RolloutFnTrainOutput(samples=[[]]) + mock_fn = MagicMock(return_value=expected_output) + adapter = LegacyRolloutFnAdapter(constructor_input, mock_fn) - result = adapter(RolloutFnEvalInput(rollout_id=0)) + result = adapter(RolloutFnTrainInput(rollout_id=0)) - assert result is expected_output + assert result is expected_output + + +def test_adapter_passthrough_eval_output(constructor_input): + expected_output = RolloutFnEvalOutput(data={}) + mock_fn = MagicMock(return_value=expected_output) + adapter = LegacyRolloutFnAdapter(constructor_input, mock_fn) + + result = adapter(RolloutFnEvalInput(rollout_id=0)) + + assert result is expected_output From 6836e6044b8d6cfda6bd19d524f0b6f9da1c288c Mon Sep 17 00:00:00 2001 From: fzyzcjy Date: Wed, 14 Jan 2026 16:16:25 +0800 Subject: [PATCH 31/46] more --- .../modular_rollout/test_compatibility.py | 92 +++++++++---------- 1 file changed, 45 insertions(+), 47 deletions(-) diff --git a/tests/rollout/modular_rollout/test_compatibility.py b/tests/rollout/modular_rollout/test_compatibility.py index 0c27efe6c56..60c06ed743c 100644 --- a/tests/rollout/modular_rollout/test_compatibility.py +++ b/tests/rollout/modular_rollout/test_compatibility.py @@ -20,68 +20,66 @@ def constructor_input(): return RolloutFnConstructorInput(args="dummy_args", data_source="dummy_data_source") -def test_load_class(constructor_input): - class MockRolloutClass: - pass +class TestLoadRolloutFunction: + def test_load_class(self, constructor_input): + class MockRolloutClass: + pass - with patch("miles.rollout.modular_rollout.compatibility.load_function", return_value=MockRolloutClass): - result = load_rollout_function(constructor_input, "some.module.MockRolloutClass") + with patch("miles.rollout.modular_rollout.compatibility.load_function", return_value=MockRolloutClass): + result = load_rollout_function(constructor_input, "some.module.MockRolloutClass") - assert result is MockRolloutClass + assert result is MockRolloutClass + def test_load_function_returns_adapter(self, constructor_input): + def mock_fn(): + pass -def test_load_function_returns_adapter(constructor_input): - def mock_fn(): - pass + with patch("miles.rollout.modular_rollout.compatibility.load_function", return_value=mock_fn): + result = load_rollout_function(constructor_input, "some.module.mock_fn") - with patch("miles.rollout.modular_rollout.compatibility.load_function", return_value=mock_fn): - result = load_rollout_function(constructor_input, "some.module.mock_fn") + assert isinstance(result, LegacyRolloutFnAdapter) + assert result.fn is mock_fn + assert result.args == "dummy_args" + assert result.data_source == "dummy_data_source" - assert isinstance(result, LegacyRolloutFnAdapter) - assert result.fn is mock_fn - assert result.args == "dummy_args" - assert result.data_source == "dummy_data_source" +class TestLegacyRolloutFnAdapter: + def test_call_with_train_input_wraps_output(self, constructor_input): + mock_samples = [[{"text": "sample"}]] + mock_fn = MagicMock(return_value=mock_samples) + adapter = LegacyRolloutFnAdapter(constructor_input, mock_fn) -def test_adapter_call_with_train_input_wraps_output(constructor_input): - mock_samples = [[{"text": "sample"}]] - mock_fn = MagicMock(return_value=mock_samples) - adapter = LegacyRolloutFnAdapter(constructor_input, mock_fn) + result = adapter(RolloutFnTrainInput(rollout_id=1)) - result = adapter(RolloutFnTrainInput(rollout_id=1)) + mock_fn.assert_called_once_with("dummy_args", 1, "dummy_data_source", evaluation=False) + assert isinstance(result, RolloutFnTrainOutput) + assert result.samples == mock_samples - mock_fn.assert_called_once_with("dummy_args", 1, "dummy_data_source", evaluation=False) - assert isinstance(result, RolloutFnTrainOutput) - assert result.samples == mock_samples + def test_call_with_eval_input_wraps_output(self, constructor_input): + mock_data = {"metric": {"accuracy": 0.9}} + mock_fn = MagicMock(return_value=mock_data) + adapter = LegacyRolloutFnAdapter(constructor_input, mock_fn) + result = adapter(RolloutFnEvalInput(rollout_id=2)) -def test_adapter_call_with_eval_input_wraps_output(constructor_input): - mock_data = {"metric": {"accuracy": 0.9}} - mock_fn = MagicMock(return_value=mock_data) - adapter = LegacyRolloutFnAdapter(constructor_input, mock_fn) + mock_fn.assert_called_once_with("dummy_args", 2, "dummy_data_source", evaluation=True) + assert isinstance(result, RolloutFnEvalOutput) + assert result.data == mock_data - result = adapter(RolloutFnEvalInput(rollout_id=2)) + def test_passthrough_train_output(self, constructor_input): + expected_output = RolloutFnTrainOutput(samples=[[]]) + mock_fn = MagicMock(return_value=expected_output) + adapter = LegacyRolloutFnAdapter(constructor_input, mock_fn) - mock_fn.assert_called_once_with("dummy_args", 2, "dummy_data_source", evaluation=True) - assert isinstance(result, RolloutFnEvalOutput) - assert result.data == mock_data + result = adapter(RolloutFnTrainInput(rollout_id=0)) + assert result is expected_output -def test_adapter_passthrough_train_output(constructor_input): - expected_output = RolloutFnTrainOutput(samples=[[]]) - mock_fn = MagicMock(return_value=expected_output) - adapter = LegacyRolloutFnAdapter(constructor_input, mock_fn) + def test_passthrough_eval_output(self, constructor_input): + expected_output = RolloutFnEvalOutput(data={}) + mock_fn = MagicMock(return_value=expected_output) + adapter = LegacyRolloutFnAdapter(constructor_input, mock_fn) - result = adapter(RolloutFnTrainInput(rollout_id=0)) + result = adapter(RolloutFnEvalInput(rollout_id=0)) - assert result is expected_output - - -def test_adapter_passthrough_eval_output(constructor_input): - expected_output = RolloutFnEvalOutput(data={}) - mock_fn = MagicMock(return_value=expected_output) - adapter = LegacyRolloutFnAdapter(constructor_input, mock_fn) - - result = adapter(RolloutFnEvalInput(rollout_id=0)) - - assert result is expected_output + assert result is expected_output From a2a74e84dddf1af0905235f2b18e4815235d79cc Mon Sep 17 00:00:00 2001 From: fzyzcjy Date: Wed, 14 Jan 2026 16:17:22 +0800 Subject: [PATCH 32/46] fmt --- tests/rollout/modular_rollout/test_compatibility.py | 5 +---- 1 file changed, 1 insertion(+), 4 deletions(-) diff --git a/tests/rollout/modular_rollout/test_compatibility.py b/tests/rollout/modular_rollout/test_compatibility.py index 60c06ed743c..65edfce7fde 100644 --- a/tests/rollout/modular_rollout/test_compatibility.py +++ b/tests/rollout/modular_rollout/test_compatibility.py @@ -9,10 +9,7 @@ RolloutFnTrainInput, RolloutFnTrainOutput, ) -from miles.rollout.modular_rollout.compatibility import ( - LegacyRolloutFnAdapter, - load_rollout_function, -) +from miles.rollout.modular_rollout.compatibility import LegacyRolloutFnAdapter, load_rollout_function @pytest.fixture From 7adffa1d8bee15eb3a21c17ae469651edd88a48b Mon Sep 17 00:00:00 2001 From: fzyzcjy Date: Wed, 14 Jan 2026 16:19:39 +0800 Subject: [PATCH 33/46] more --- miles/rollout/base_types.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/miles/rollout/base_types.py b/miles/rollout/base_types.py index e7842c977d7..83c7c41e94d 100644 --- a/miles/rollout/base_types.py +++ b/miles/rollout/base_types.py @@ -53,12 +53,10 @@ class RolloutFnEvalOutput: RolloutFnOutput = RolloutFnTrainOutput | RolloutFnEvalOutput +# 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 class RolloutFnProtocol(Protocol): def __init__(self, input: RolloutFnConstructorInput): ... - @classmethod - def add_arguments(cls, parser: ArgumentParser): ... - def __call__(self, input: RolloutFnInput) -> RolloutFnOutput: ... From 886355cbe5ba18d082a11516ccd4674487d3864b Mon Sep 17 00:00:00 2001 From: fzyzcjy Date: Wed, 14 Jan 2026 16:20:48 +0800 Subject: [PATCH 34/46] more --- miles/rollout/base_types.py | 4 +--- miles/rollout/modular_rollout/compatibility.py | 4 ++++ 2 files changed, 5 insertions(+), 3 deletions(-) diff --git a/miles/rollout/base_types.py b/miles/rollout/base_types.py index 83c7c41e94d..503f0c14c12 100644 --- a/miles/rollout/base_types.py +++ b/miles/rollout/base_types.py @@ -1,4 +1,4 @@ -from argparse import ArgumentParser, Namespace +from argparse import Namespace from dataclasses import dataclass from typing import Any, Protocol @@ -57,6 +57,4 @@ class RolloutFnEvalOutput: # TODO: may add save/load if need it to be stateful # Duck typing, users do not need to extend this class class RolloutFnProtocol(Protocol): - def __init__(self, input: RolloutFnConstructorInput): ... - def __call__(self, input: RolloutFnInput) -> RolloutFnOutput: ... diff --git a/miles/rollout/modular_rollout/compatibility.py b/miles/rollout/modular_rollout/compatibility.py index 21239590926..209d13fa06a 100644 --- a/miles/rollout/modular_rollout/compatibility.py +++ b/miles/rollout/modular_rollout/compatibility.py @@ -6,6 +6,7 @@ RolloutFnEvalOutput, RolloutFnInput, RolloutFnOutput, + RolloutFnProtocol, RolloutFnTrainOutput, ) from miles.utils.misc import load_function @@ -27,6 +28,9 @@ def __call__(self, input: RolloutFnInput) -> RolloutFnOutput: return output +assert isinstance(LegacyRolloutFnAdapter, RolloutFnProtocol) + + def load_rollout_function(input: RolloutFnConstructorInput, path: str): fn = load_function(path) From 5e8e2e5e091aace18ff90a0f4eb56344ea3726f1 Mon Sep 17 00:00:00 2001 From: fzyzcjy Date: Wed, 14 Jan 2026 16:21:03 +0800 Subject: [PATCH 35/46] more --- miles/rollout/base_types.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/miles/rollout/base_types.py b/miles/rollout/base_types.py index 503f0c14c12..d133364d3d1 100644 --- a/miles/rollout/base_types.py +++ b/miles/rollout/base_types.py @@ -1,6 +1,6 @@ from argparse import Namespace from dataclasses import dataclass -from typing import Any, Protocol +from typing import Any, Protocol, runtime_checkable from miles.rollout.data_source import DataSource from miles.utils.types import Sample @@ -56,5 +56,6 @@ class RolloutFnEvalOutput: # 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: ... From 426010ba4a303ec334695ae3a0efe7c73e097773 Mon Sep 17 00:00:00 2001 From: fzyzcjy Date: Wed, 14 Jan 2026 16:21:11 +0800 Subject: [PATCH 36/46] more --- miles/rollout/modular_rollout/compatibility.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/miles/rollout/modular_rollout/compatibility.py b/miles/rollout/modular_rollout/compatibility.py index 209d13fa06a..b1e1d56eb97 100644 --- a/miles/rollout/modular_rollout/compatibility.py +++ b/miles/rollout/modular_rollout/compatibility.py @@ -28,7 +28,7 @@ def __call__(self, input: RolloutFnInput) -> RolloutFnOutput: return output -assert isinstance(LegacyRolloutFnAdapter, RolloutFnProtocol) +assert issubclass(LegacyRolloutFnAdapter, RolloutFnProtocol) def load_rollout_function(input: RolloutFnConstructorInput, path: str): From e8c20e99e0ca7b2ed0a669dffeb07caba6aa03e8 Mon Sep 17 00:00:00 2001 From: fzyzcjy Date: Wed, 14 Jan 2026 16:43:12 +0800 Subject: [PATCH 37/46] more (cherry picked from commit 22edda1266375bf25a1b330bbb4999711b6b966b) --- miles/rollout/modular_rollout/compatibility.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/miles/rollout/modular_rollout/compatibility.py b/miles/rollout/modular_rollout/compatibility.py index b1e1d56eb97..17728c9ebab 100644 --- a/miles/rollout/modular_rollout/compatibility.py +++ b/miles/rollout/modular_rollout/compatibility.py @@ -34,7 +34,7 @@ def __call__(self, input: RolloutFnInput) -> RolloutFnOutput: def load_rollout_function(input: RolloutFnConstructorInput, path: str): fn = load_function(path) - if not inspect.isclass(fn): - fn = LegacyRolloutFnAdapter(input, fn) - - return fn + if inspect.isclass(fn): + return fn(input) + else: + return LegacyRolloutFnAdapter(input, fn) From fd4d78cac2c1ca6f3c7e790b3b3d8868c91d461e Mon Sep 17 00:00:00 2001 From: fzyzcjy Date: Wed, 14 Jan 2026 16:46:36 +0800 Subject: [PATCH 38/46] more --- tests/rollout/modular_rollout/test_compatibility.py | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/tests/rollout/modular_rollout/test_compatibility.py b/tests/rollout/modular_rollout/test_compatibility.py index 65edfce7fde..55748c68b56 100644 --- a/tests/rollout/modular_rollout/test_compatibility.py +++ b/tests/rollout/modular_rollout/test_compatibility.py @@ -18,14 +18,16 @@ def constructor_input(): class TestLoadRolloutFunction: - def test_load_class(self, constructor_input): + def test_load_class_returns_instance(self, constructor_input): class MockRolloutClass: - pass + def __init__(self, input): + self.input = input with patch("miles.rollout.modular_rollout.compatibility.load_function", return_value=MockRolloutClass): result = load_rollout_function(constructor_input, "some.module.MockRolloutClass") - assert result is MockRolloutClass + assert isinstance(result, MockRolloutClass) + assert result.input is constructor_input def test_load_function_returns_adapter(self, constructor_input): def mock_fn(): From 9a2d3e262bab47be12ab898cfae216df9491be9e Mon Sep 17 00:00:00 2001 From: fzyzcjy Date: Wed, 14 Jan 2026 16:48:51 +0800 Subject: [PATCH 39/46] more --- miles/ray/rollout.py | 6 +++--- miles/rollout/base_types.py | 4 ++-- miles/rollout/modular_rollout/compatibility.py | 10 ++++++++++ 3 files changed, 15 insertions(+), 5 deletions(-) diff --git a/miles/ray/rollout.py b/miles/ray/rollout.py index 1a1c40e5436..3867765ee1a 100644 --- a/miles/ray/rollout.py +++ b/miles/ray/rollout.py @@ -14,7 +14,7 @@ from miles.backends.sglang_utils.sglang_engine import SGLangEngine from miles.rollout.base_types import RolloutFnConstructorInput, RolloutFnEvalInput, RolloutFnTrainInput -from miles.rollout.modular_rollout.compatibility import load_rollout_function +from miles.rollout.modular_rollout.compatibility import load_rollout_function, call_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 @@ -144,7 +144,7 @@ def eval(self, rollout_id): return self.health_monitoring_resume() - result = self.eval_generate_rollout(RolloutFnEvalInput(rollout_id=rollout_id)) + 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) @@ -226,7 +226,7 @@ def _get_rollout_data(self, rollout_id): ) metrics = None else: - data = self.generate_rollout(RolloutFnTrainInput(rollout_id=rollout_id)) + 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 diff --git a/miles/rollout/base_types.py b/miles/rollout/base_types.py index d133364d3d1..0d75b726d15 100644 --- a/miles/rollout/base_types.py +++ b/miles/rollout/base_types.py @@ -1,6 +1,6 @@ from argparse import Namespace from dataclasses import dataclass -from typing import Any, Protocol, runtime_checkable +from typing import Any, Protocol, runtime_checkable, Awaitable from miles.rollout.data_source import DataSource from miles.utils.types import Sample @@ -58,4 +58,4 @@ class RolloutFnEvalOutput: # Duck typing, users do not need to extend this class @runtime_checkable class RolloutFnProtocol(Protocol): - def __call__(self, input: RolloutFnInput) -> RolloutFnOutput: ... + def __call__(self, input: RolloutFnInput) -> RolloutFnOutput | Awaitable[RolloutFnOutput]: ... diff --git a/miles/rollout/modular_rollout/compatibility.py b/miles/rollout/modular_rollout/compatibility.py index 17728c9ebab..2b980c89b0b 100644 --- a/miles/rollout/modular_rollout/compatibility.py +++ b/miles/rollout/modular_rollout/compatibility.py @@ -10,6 +10,7 @@ RolloutFnTrainOutput, ) from miles.utils.misc import load_function +from miles.utils.async_utils import run class LegacyRolloutFnAdapter: @@ -38,3 +39,12 @@ def load_rollout_function(input: RolloutFnConstructorInput, path: str): 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 From e9d05949bed20950bc941a0240bdc3f3696d2740 Mon Sep 17 00:00:00 2001 From: fzyzcjy Date: Wed, 14 Jan 2026 16:49:04 +0800 Subject: [PATCH 40/46] fmt --- miles/ray/rollout.py | 2 +- miles/rollout/base_types.py | 3 ++- miles/rollout/modular_rollout/compatibility.py | 2 +- 3 files changed, 4 insertions(+), 3 deletions(-) diff --git a/miles/ray/rollout.py b/miles/ray/rollout.py index 3867765ee1a..1cba8b7e00b 100644 --- a/miles/ray/rollout.py +++ b/miles/ray/rollout.py @@ -14,7 +14,7 @@ from miles.backends.sglang_utils.sglang_engine import SGLangEngine from miles.rollout.base_types import RolloutFnConstructorInput, RolloutFnEvalInput, RolloutFnTrainInput -from miles.rollout.modular_rollout.compatibility import load_rollout_function, call_rollout_function +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 diff --git a/miles/rollout/base_types.py b/miles/rollout/base_types.py index 0d75b726d15..d6eb1e8f0d2 100644 --- a/miles/rollout/base_types.py +++ b/miles/rollout/base_types.py @@ -1,6 +1,7 @@ from argparse import Namespace +from collections.abc import Awaitable from dataclasses import dataclass -from typing import Any, Protocol, runtime_checkable, Awaitable +from typing import Any, Protocol, runtime_checkable from miles.rollout.data_source import DataSource from miles.utils.types import Sample diff --git a/miles/rollout/modular_rollout/compatibility.py b/miles/rollout/modular_rollout/compatibility.py index 2b980c89b0b..7d1a70e79c6 100644 --- a/miles/rollout/modular_rollout/compatibility.py +++ b/miles/rollout/modular_rollout/compatibility.py @@ -9,8 +9,8 @@ RolloutFnProtocol, RolloutFnTrainOutput, ) -from miles.utils.misc import load_function from miles.utils.async_utils import run +from miles.utils.misc import load_function class LegacyRolloutFnAdapter: From 7752e2f0004ffb5e39c390d2bcd7af80031f24e5 Mon Sep 17 00:00:00 2001 From: fzyzcjy Date: Wed, 14 Jan 2026 16:54:30 +0800 Subject: [PATCH 41/46] more --- .../modular_rollout/test_compatibility.py | 53 +++++++++++++++++-- 1 file changed, 48 insertions(+), 5 deletions(-) diff --git a/tests/rollout/modular_rollout/test_compatibility.py b/tests/rollout/modular_rollout/test_compatibility.py index 55748c68b56..d394faca663 100644 --- a/tests/rollout/modular_rollout/test_compatibility.py +++ b/tests/rollout/modular_rollout/test_compatibility.py @@ -1,3 +1,4 @@ +import asyncio from unittest.mock import MagicMock, patch import pytest @@ -9,7 +10,11 @@ RolloutFnTrainInput, RolloutFnTrainOutput, ) -from miles.rollout.modular_rollout.compatibility import LegacyRolloutFnAdapter, load_rollout_function +from miles.rollout.modular_rollout.compatibility import ( + LegacyRolloutFnAdapter, + call_rollout_function, + load_rollout_function, +) @pytest.fixture @@ -48,7 +53,7 @@ def test_call_with_train_input_wraps_output(self, constructor_input): mock_fn = MagicMock(return_value=mock_samples) adapter = LegacyRolloutFnAdapter(constructor_input, mock_fn) - result = adapter(RolloutFnTrainInput(rollout_id=1)) + result = call_rollout_function(adapter, RolloutFnTrainInput(rollout_id=1)) mock_fn.assert_called_once_with("dummy_args", 1, "dummy_data_source", evaluation=False) assert isinstance(result, RolloutFnTrainOutput) @@ -59,7 +64,7 @@ def test_call_with_eval_input_wraps_output(self, constructor_input): mock_fn = MagicMock(return_value=mock_data) adapter = LegacyRolloutFnAdapter(constructor_input, mock_fn) - result = adapter(RolloutFnEvalInput(rollout_id=2)) + result = call_rollout_function(adapter, RolloutFnEvalInput(rollout_id=2)) mock_fn.assert_called_once_with("dummy_args", 2, "dummy_data_source", evaluation=True) assert isinstance(result, RolloutFnEvalOutput) @@ -70,7 +75,7 @@ def test_passthrough_train_output(self, constructor_input): mock_fn = MagicMock(return_value=expected_output) adapter = LegacyRolloutFnAdapter(constructor_input, mock_fn) - result = adapter(RolloutFnTrainInput(rollout_id=0)) + result = call_rollout_function(adapter, RolloutFnTrainInput(rollout_id=0)) assert result is expected_output @@ -79,6 +84,44 @@ def test_passthrough_eval_output(self, constructor_input): mock_fn = MagicMock(return_value=expected_output) adapter = LegacyRolloutFnAdapter(constructor_input, mock_fn) - result = adapter(RolloutFnEvalInput(rollout_id=0)) + result = call_rollout_function(adapter, RolloutFnEvalInput(rollout_id=0)) assert result is expected_output + + +async def async_mock_fn_train(args, rollout_id, data_source, evaluation): + await asyncio.sleep(0.01) + return RolloutFnTrainOutput(samples=[[{"text": "async_sample"}]]) + + +async def async_mock_fn_eval(args, rollout_id, data_source, evaluation): + await asyncio.sleep(0.01) + return RolloutFnEvalOutput(data={"metric": {"accuracy": 0.95}}) + + +class TestCallRolloutFunction: + def test_sync_function(self, constructor_input): + mock_samples = [[{"text": "sample"}]] + mock_fn = MagicMock(return_value=mock_samples) + adapter = LegacyRolloutFnAdapter(constructor_input, mock_fn) + + result = call_rollout_function(adapter, RolloutFnTrainInput(rollout_id=1)) + + assert isinstance(result, RolloutFnTrainOutput) + assert result.samples == mock_samples + + def test_async_function_train(self, constructor_input): + adapter = LegacyRolloutFnAdapter(constructor_input, async_mock_fn_train) + + result = call_rollout_function(adapter, RolloutFnTrainInput(rollout_id=1)) + + assert isinstance(result, RolloutFnTrainOutput) + assert result.samples == [[{"text": "async_sample"}]] + + def test_async_function_eval(self, constructor_input): + adapter = LegacyRolloutFnAdapter(constructor_input, async_mock_fn_eval) + + result = call_rollout_function(adapter, RolloutFnEvalInput(rollout_id=2)) + + assert isinstance(result, RolloutFnEvalOutput) + assert result.data == {"metric": {"accuracy": 0.95}} From d99bfb4438cf1a305ced924c80f48da32025fcf0 Mon Sep 17 00:00:00 2001 From: fzyzcjy Date: Wed, 14 Jan 2026 16:54:59 +0800 Subject: [PATCH 42/46] more --- .../modular_rollout/test_compatibility.py | 66 ++++++++++++++----- 1 file changed, 51 insertions(+), 15 deletions(-) diff --git a/tests/rollout/modular_rollout/test_compatibility.py b/tests/rollout/modular_rollout/test_compatibility.py index d394faca663..265f3b5d981 100644 --- a/tests/rollout/modular_rollout/test_compatibility.py +++ b/tests/rollout/modular_rollout/test_compatibility.py @@ -89,18 +89,46 @@ def test_passthrough_eval_output(self, constructor_input): assert result is expected_output -async def async_mock_fn_train(args, rollout_id, data_source, evaluation): - await asyncio.sleep(0.01) - return RolloutFnTrainOutput(samples=[[{"text": "async_sample"}]]) +class MockSyncRolloutClass: + def __init__(self, input): + self.input = input + def __call__(self, input): + return RolloutFnTrainOutput(samples=[[{"text": "sync_class"}]]) + + @classmethod + def add_arguments(cls, parser): + pass -async def async_mock_fn_eval(args, rollout_id, data_source, evaluation): - await asyncio.sleep(0.01) - return RolloutFnEvalOutput(data={"metric": {"accuracy": 0.95}}) + +class MockAsyncRolloutClass: + def __init__(self, input): + self.input = input + + async def __call__(self, input): + await asyncio.sleep(0.01) + return RolloutFnTrainOutput(samples=[[{"text": "async_class"}]]) + + @classmethod + def add_arguments(cls, parser): + pass + + +class MockAsyncRolloutClassEval: + def __init__(self, input): + self.input = input + + async def __call__(self, input): + await asyncio.sleep(0.01) + return RolloutFnEvalOutput(data={"metric": {"accuracy": 0.98}}) + + @classmethod + def add_arguments(cls, parser): + pass class TestCallRolloutFunction: - def test_sync_function(self, constructor_input): + def test_sync_adapter(self, constructor_input): mock_samples = [[{"text": "sample"}]] mock_fn = MagicMock(return_value=mock_samples) adapter = LegacyRolloutFnAdapter(constructor_input, mock_fn) @@ -110,18 +138,26 @@ def test_sync_function(self, constructor_input): assert isinstance(result, RolloutFnTrainOutput) assert result.samples == mock_samples - def test_async_function_train(self, constructor_input): - adapter = LegacyRolloutFnAdapter(constructor_input, async_mock_fn_train) + def test_sync_class(self, constructor_input): + instance = MockSyncRolloutClass(constructor_input) - result = call_rollout_function(adapter, RolloutFnTrainInput(rollout_id=1)) + result = call_rollout_function(instance, RolloutFnTrainInput(rollout_id=1)) assert isinstance(result, RolloutFnTrainOutput) - assert result.samples == [[{"text": "async_sample"}]] + assert result.samples == [[{"text": "sync_class"}]] - def test_async_function_eval(self, constructor_input): - adapter = LegacyRolloutFnAdapter(constructor_input, async_mock_fn_eval) + def test_async_class(self, constructor_input): + instance = MockAsyncRolloutClass(constructor_input) - result = call_rollout_function(adapter, RolloutFnEvalInput(rollout_id=2)) + result = call_rollout_function(instance, RolloutFnTrainInput(rollout_id=1)) + + assert isinstance(result, RolloutFnTrainOutput) + assert result.samples == [[{"text": "async_class"}]] + + def test_async_class_eval(self, constructor_input): + instance = MockAsyncRolloutClassEval(constructor_input) + + result = call_rollout_function(instance, RolloutFnEvalInput(rollout_id=2)) assert isinstance(result, RolloutFnEvalOutput) - assert result.data == {"metric": {"accuracy": 0.95}} + assert result.data == {"metric": {"accuracy": 0.98}} From 47f80cf8901a89cd07e8c93f598e90f0336c4f40 Mon Sep 17 00:00:00 2001 From: fzyzcjy Date: Wed, 14 Jan 2026 16:55:09 +0800 Subject: [PATCH 43/46] more --- tests/rollout/modular_rollout/test_compatibility.py | 12 ------------ 1 file changed, 12 deletions(-) diff --git a/tests/rollout/modular_rollout/test_compatibility.py b/tests/rollout/modular_rollout/test_compatibility.py index 265f3b5d981..c406fc8b7dd 100644 --- a/tests/rollout/modular_rollout/test_compatibility.py +++ b/tests/rollout/modular_rollout/test_compatibility.py @@ -96,10 +96,6 @@ def __init__(self, input): def __call__(self, input): return RolloutFnTrainOutput(samples=[[{"text": "sync_class"}]]) - @classmethod - def add_arguments(cls, parser): - pass - class MockAsyncRolloutClass: def __init__(self, input): @@ -109,10 +105,6 @@ async def __call__(self, input): await asyncio.sleep(0.01) return RolloutFnTrainOutput(samples=[[{"text": "async_class"}]]) - @classmethod - def add_arguments(cls, parser): - pass - class MockAsyncRolloutClassEval: def __init__(self, input): @@ -122,10 +114,6 @@ async def __call__(self, input): await asyncio.sleep(0.01) return RolloutFnEvalOutput(data={"metric": {"accuracy": 0.98}}) - @classmethod - def add_arguments(cls, parser): - pass - class TestCallRolloutFunction: def test_sync_adapter(self, constructor_input): From b8ff00c82e0bac5d65d3bf6c114710193204dcad Mon Sep 17 00:00:00 2001 From: fzyzcjy Date: Wed, 14 Jan 2026 16:55:36 +0800 Subject: [PATCH 44/46] fmt --- tests/rollout/modular_rollout/test_compatibility.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/tests/rollout/modular_rollout/test_compatibility.py b/tests/rollout/modular_rollout/test_compatibility.py index c406fc8b7dd..191a835d9b8 100644 --- a/tests/rollout/modular_rollout/test_compatibility.py +++ b/tests/rollout/modular_rollout/test_compatibility.py @@ -95,7 +95,7 @@ def __init__(self, input): def __call__(self, input): return RolloutFnTrainOutput(samples=[[{"text": "sync_class"}]]) - + class MockAsyncRolloutClass: def __init__(self, input): @@ -104,7 +104,7 @@ def __init__(self, input): async def __call__(self, input): await asyncio.sleep(0.01) return RolloutFnTrainOutput(samples=[[{"text": "async_class"}]]) - + class MockAsyncRolloutClassEval: def __init__(self, input): @@ -113,7 +113,7 @@ def __init__(self, input): async def __call__(self, input): await asyncio.sleep(0.01) return RolloutFnEvalOutput(data={"metric": {"accuracy": 0.98}}) - + class TestCallRolloutFunction: def test_sync_adapter(self, constructor_input): From cccc61069a289c0b5c458c23962c341788eff75e Mon Sep 17 00:00:00 2001 From: fzyzcjy Date: Wed, 14 Jan 2026 20:57:01 +0800 Subject: [PATCH 45/46] more --- .../modular_rollout/test_compatibility.py | 233 ++++++++---------- 1 file changed, 105 insertions(+), 128 deletions(-) diff --git a/tests/rollout/modular_rollout/test_compatibility.py b/tests/rollout/modular_rollout/test_compatibility.py index 191a835d9b8..6ead5f436e3 100644 --- a/tests/rollout/modular_rollout/test_compatibility.py +++ b/tests/rollout/modular_rollout/test_compatibility.py @@ -1,5 +1,5 @@ import asyncio -from unittest.mock import MagicMock, patch +from unittest.mock import patch import pytest @@ -22,130 +22,107 @@ def constructor_input(): return RolloutFnConstructorInput(args="dummy_args", data_source="dummy_data_source") -class TestLoadRolloutFunction: - def test_load_class_returns_instance(self, constructor_input): - class MockRolloutClass: - def __init__(self, input): - self.input = input - - with patch("miles.rollout.modular_rollout.compatibility.load_function", return_value=MockRolloutClass): - result = load_rollout_function(constructor_input, "some.module.MockRolloutClass") - - assert isinstance(result, MockRolloutClass) - assert result.input is constructor_input - - def test_load_function_returns_adapter(self, constructor_input): - def mock_fn(): - pass - - with patch("miles.rollout.modular_rollout.compatibility.load_function", return_value=mock_fn): - result = load_rollout_function(constructor_input, "some.module.mock_fn") - - assert isinstance(result, LegacyRolloutFnAdapter) - assert result.fn is mock_fn - assert result.args == "dummy_args" - assert result.data_source == "dummy_data_source" - - -class TestLegacyRolloutFnAdapter: - def test_call_with_train_input_wraps_output(self, constructor_input): - mock_samples = [[{"text": "sample"}]] - mock_fn = MagicMock(return_value=mock_samples) - adapter = LegacyRolloutFnAdapter(constructor_input, mock_fn) - - result = call_rollout_function(adapter, RolloutFnTrainInput(rollout_id=1)) - - mock_fn.assert_called_once_with("dummy_args", 1, "dummy_data_source", evaluation=False) - assert isinstance(result, RolloutFnTrainOutput) - assert result.samples == mock_samples - - def test_call_with_eval_input_wraps_output(self, constructor_input): - mock_data = {"metric": {"accuracy": 0.9}} - mock_fn = MagicMock(return_value=mock_data) - adapter = LegacyRolloutFnAdapter(constructor_input, mock_fn) - - result = call_rollout_function(adapter, RolloutFnEvalInput(rollout_id=2)) - - mock_fn.assert_called_once_with("dummy_args", 2, "dummy_data_source", evaluation=True) - assert isinstance(result, RolloutFnEvalOutput) - assert result.data == mock_data - - def test_passthrough_train_output(self, constructor_input): - expected_output = RolloutFnTrainOutput(samples=[[]]) - mock_fn = MagicMock(return_value=expected_output) - adapter = LegacyRolloutFnAdapter(constructor_input, mock_fn) - - result = call_rollout_function(adapter, RolloutFnTrainInput(rollout_id=0)) - - assert result is expected_output - - def test_passthrough_eval_output(self, constructor_input): - expected_output = RolloutFnEvalOutput(data={}) - mock_fn = MagicMock(return_value=expected_output) - adapter = LegacyRolloutFnAdapter(constructor_input, mock_fn) - - result = call_rollout_function(adapter, RolloutFnEvalInput(rollout_id=0)) - - assert result is expected_output - - -class MockSyncRolloutClass: - def __init__(self, input): - self.input = input - - def __call__(self, input): - return RolloutFnTrainOutput(samples=[[{"text": "sync_class"}]]) - - -class MockAsyncRolloutClass: - def __init__(self, input): - self.input = input - - async def __call__(self, input): - await asyncio.sleep(0.01) - return RolloutFnTrainOutput(samples=[[{"text": "async_class"}]]) - - -class MockAsyncRolloutClassEval: - def __init__(self, input): - self.input = input - - async def __call__(self, input): - await asyncio.sleep(0.01) - return RolloutFnEvalOutput(data={"metric": {"accuracy": 0.98}}) - - -class TestCallRolloutFunction: - def test_sync_adapter(self, constructor_input): - mock_samples = [[{"text": "sample"}]] - mock_fn = MagicMock(return_value=mock_samples) - adapter = LegacyRolloutFnAdapter(constructor_input, mock_fn) - - result = call_rollout_function(adapter, RolloutFnTrainInput(rollout_id=1)) - - assert isinstance(result, RolloutFnTrainOutput) - assert result.samples == mock_samples - - def test_sync_class(self, constructor_input): - instance = MockSyncRolloutClass(constructor_input) - - result = call_rollout_function(instance, RolloutFnTrainInput(rollout_id=1)) - - assert isinstance(result, RolloutFnTrainOutput) - assert result.samples == [[{"text": "sync_class"}]] - - def test_async_class(self, constructor_input): - instance = MockAsyncRolloutClass(constructor_input) - - result = call_rollout_function(instance, RolloutFnTrainInput(rollout_id=1)) - - assert isinstance(result, RolloutFnTrainOutput) - assert result.samples == [[{"text": "async_class"}]] - - def test_async_class_eval(self, constructor_input): - instance = MockAsyncRolloutClassEval(constructor_input) - - result = call_rollout_function(instance, RolloutFnEvalInput(rollout_id=2)) - - assert isinstance(result, RolloutFnEvalOutput) - assert result.data == {"metric": {"accuracy": 0.98}} +class TestSupportedRolloutFormats: + """ + Supported rollout function formats: + + Format 1: Legacy function returning raw data + def fn(args, rollout_id, data_source, evaluation=False) -> list | dict + + Format 2: Legacy function returning typed output + def fn(args, rollout_id, data_source, evaluation=False) -> RolloutFnTrainOutput | RolloutFnEvalOutput + + Format 3: Sync class + class Fn: + def __init__(self, input: RolloutFnConstructorInput): ... + def __call__(self, input: RolloutFnInput) -> RolloutFnOutput: ... + + Format 4: Async class + class Fn: + def __init__(self, input: RolloutFnConstructorInput): ... + async def __call__(self, input: RolloutFnInput) -> RolloutFnOutput: ... + """ + + @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) From ac23b134c19a801cdd33dc3aed7c423b16931b93 Mon Sep 17 00:00:00 2001 From: fzyzcjy Date: Wed, 14 Jan 2026 20:57:25 +0800 Subject: [PATCH 46/46] more --- .../modular_rollout/test_compatibility.py | 18 +----------------- 1 file changed, 1 insertion(+), 17 deletions(-) diff --git a/tests/rollout/modular_rollout/test_compatibility.py b/tests/rollout/modular_rollout/test_compatibility.py index 6ead5f436e3..596fa76270f 100644 --- a/tests/rollout/modular_rollout/test_compatibility.py +++ b/tests/rollout/modular_rollout/test_compatibility.py @@ -24,23 +24,7 @@ def constructor_input(): class TestSupportedRolloutFormats: """ - Supported rollout function formats: - - Format 1: Legacy function returning raw data - def fn(args, rollout_id, data_source, evaluation=False) -> list | dict - - Format 2: Legacy function returning typed output - def fn(args, rollout_id, data_source, evaluation=False) -> RolloutFnTrainOutput | RolloutFnEvalOutput - - Format 3: Sync class - class Fn: - def __init__(self, input: RolloutFnConstructorInput): ... - def __call__(self, input: RolloutFnInput) -> RolloutFnOutput: ... - - Format 4: Async class - class Fn: - def __init__(self, input: RolloutFnConstructorInput): ... - async def __call__(self, input: RolloutFnInput) -> RolloutFnOutput: ... + Documentation test to show various supported rollout function formats """ @pytest.mark.parametrize("evaluation", [False, True])