rose-opd-implementation / frameworks /slime /tests /plugin_contracts /test_plugin_rollout_contracts.py
| from __future__ import annotations | |
| import inspect | |
| import pytest | |
| try: | |
| from ._shared import get_contract_path, install_paths, install_stubs, run_contract_test_for_file | |
| except ImportError: | |
| try: | |
| from plugin_contracts._shared import ( | |
| get_contract_path, | |
| install_paths, | |
| install_stubs, | |
| run_contract_test_for_file, | |
| ) | |
| except ImportError: | |
| from _shared import get_contract_path, install_paths, install_stubs, run_contract_test_for_file | |
| install_paths() | |
| install_stubs(with_sglang_router=True, with_transformers=True) | |
| NUM_GPUS = 0 | |
| DEFAULT_ROLLOUT_FUNCTION_PATH = "slime.rollout.sglang_rollout.generate_rollout" | |
| REFERENCE_ROLLOUT_FUNCTION_PATH = "plugin_contracts.test_plugin_rollout_contracts.valid_rollout_function" | |
| from slime.rollout.base_types import RolloutFnEvalOutput, RolloutFnTrainOutput, call_rollout_fn | |
| from slime.rollout.sglang_rollout import generate_rollout as default_generate_rollout | |
| from slime.utils.misc import load_function | |
| from slime.utils.types import Sample | |
| def run_contract_test_file() -> None: | |
| run_contract_test_for_file(__file__, path_args=["rollout-function-path"]) | |
| def make_sample(index: int, reward: float = 1.0) -> Sample: | |
| tokens = [1000 + index, 2000 + index] | |
| return Sample( | |
| index=index, | |
| prompt=f"prompt-{index}", | |
| response=f"response-{index}", | |
| tokens=tokens, | |
| response_length=len(tokens), | |
| reward=reward, | |
| status=Sample.Status.COMPLETED, | |
| metadata={}, | |
| ) | |
| class ContractDataSource: | |
| def __init__(self, n_samples_per_prompt: int = 2): | |
| self.n_samples_per_prompt = n_samples_per_prompt | |
| self.next_index = 0 | |
| def get_samples(self, num_samples: int) -> list[list[Sample]]: | |
| groups = [] | |
| for _ in range(num_samples): | |
| group = [make_sample(self.next_index + i, reward=0.0) for i in range(self.n_samples_per_prompt)] | |
| self.next_index += self.n_samples_per_prompt | |
| groups.append(group) | |
| return groups | |
| def valid_rollout_function(args, rollout_id, data_source, evaluation=False): | |
| if evaluation: | |
| sample = make_sample(0, reward=0.75) | |
| return RolloutFnEvalOutput( | |
| data={"contract_eval": {"rewards": [sample.reward], "truncated": [False], "samples": [sample]}}, | |
| metrics={"source": "contract"}, | |
| ) | |
| groups = data_source.get_samples(2) | |
| for group_index, group in enumerate(groups): | |
| for sample_index, sample in enumerate(group): | |
| sample.tokens = [group_index, sample_index, rollout_id] | |
| sample.response = f"group-{group_index}-sample-{sample_index}" | |
| sample.response_length = len(sample.tokens) | |
| sample.reward = float(group_index + sample_index) | |
| sample.status = Sample.Status.COMPLETED | |
| return RolloutFnTrainOutput(samples=groups, metrics={"source": "contract"}) | |
| def invalid_rollout_function(args, rollout_id, data_source, evaluation=False): | |
| sample = make_sample(0) | |
| sample.reward = None | |
| return RolloutFnTrainOutput(samples=[[sample]]) | |
| def legacy_rollout_function(args, rollout_id, data_source, evaluation=False): | |
| if evaluation: | |
| sample = make_sample(1, reward=0.5) | |
| return {"legacy_eval": {"rewards": [sample.reward], "truncated": [False], "samples": [sample]}} | |
| return [[make_sample(1)]] | |
| def assert_sample_contract(sample: Sample) -> None: | |
| assert isinstance(sample, Sample) | |
| assert isinstance(sample.tokens, list) | |
| assert all(isinstance(token, int) for token in sample.tokens) | |
| assert isinstance(sample.response, str) | |
| assert isinstance(sample.response_length, int) | |
| assert sample.reward is not None | |
| assert isinstance(sample.status, Sample.Status) | |
| def assert_train_rollout_contract(output: RolloutFnTrainOutput, n_samples_per_prompt: int) -> None: | |
| assert isinstance(output, RolloutFnTrainOutput) | |
| assert output.samples | |
| for group in output.samples: | |
| assert len(group) == n_samples_per_prompt | |
| for sample in group: | |
| assert_sample_contract(sample) | |
| def assert_eval_rollout_contract(output: RolloutFnEvalOutput) -> None: | |
| assert isinstance(output, RolloutFnEvalOutput) | |
| assert output.data | |
| for dataset_data in output.data.values(): | |
| assert set(dataset_data) >= {"rewards", "truncated", "samples"} | |
| assert len(dataset_data["rewards"]) == len(dataset_data["truncated"]) == len(dataset_data["samples"]) | |
| for sample in dataset_data["samples"]: | |
| assert_sample_contract(sample) | |
| def assert_rollout_function_signature_matches_default(fn) -> None: | |
| default_sig = inspect.signature(default_generate_rollout) | |
| candidate_sig = inspect.signature(fn) | |
| assert tuple(candidate_sig.parameters) == tuple(default_sig.parameters) | |
| for name, default_param in default_sig.parameters.items(): | |
| candidate_param = candidate_sig.parameters[name] | |
| assert candidate_param.kind == default_param.kind | |
| assert candidate_param.default == default_param.default | |
| def assert_rollout_function_matches_default_contract(fn) -> None: | |
| assert_rollout_function_signature_matches_default(fn) | |
| data_source = ContractDataSource() | |
| train_output = call_rollout_fn(fn, None, 2, data_source, evaluation=False) | |
| eval_output = call_rollout_fn(fn, None, 2, data_source, evaluation=True) | |
| assert_train_rollout_contract(train_output, n_samples_per_prompt=2) | |
| assert_eval_rollout_contract(eval_output) | |
| def test_rollout_function_path_contract_supports_user_override(): | |
| rollout_path = get_contract_path("ROLLOUT_FUNCTION_PATH", DEFAULT_ROLLOUT_FUNCTION_PATH) | |
| rollout_fn = load_function(rollout_path) | |
| assert_rollout_function_signature_matches_default(rollout_fn) | |
| if rollout_path != DEFAULT_ROLLOUT_FUNCTION_PATH: | |
| assert_rollout_function_matches_default_contract(rollout_fn) | |
| def test_load_function_can_load_local_rollout_plugin(): | |
| func = load_function(REFERENCE_ROLLOUT_FUNCTION_PATH) | |
| assert func.__name__ == valid_rollout_function.__name__ | |
| assert func.__module__ == "plugin_contracts.test_plugin_rollout_contracts" | |
| def test_default_rollout_signature_is_stable(): | |
| default_sig = inspect.signature(default_generate_rollout) | |
| assert tuple(default_sig.parameters) == ("args", "rollout_id", "data_source", "evaluation") | |
| assert default_sig.parameters["evaluation"].default is False | |
| def test_default_rollout_compat_wrapper_stability(): | |
| data_source = ContractDataSource(n_samples_per_prompt=1) | |
| train_output = call_rollout_fn(legacy_rollout_function, None, 1, data_source, evaluation=False) | |
| eval_output = call_rollout_fn(legacy_rollout_function, None, 1, data_source, evaluation=True) | |
| assert_train_rollout_contract(train_output, n_samples_per_prompt=1) | |
| assert_eval_rollout_contract(eval_output) | |
| def test_local_rollout_plugin_aligns_with_default_input_output_format(): | |
| assert_rollout_function_matches_default_contract(valid_rollout_function) | |
| def test_misaligned_rollout_plugin_is_rejected(): | |
| with pytest.raises(AssertionError): | |
| assert_rollout_function_matches_default_contract(invalid_rollout_function) | |
| if __name__ == "__main__": | |
| run_contract_test_file() | |