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()