File size: 6,692 Bytes
af66a85 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 | from __future__ import annotations
import asyncio
import inspect
import types
from contextlib import contextmanager
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
REFERENCE_CUSTOM_GENERATE_PATH = "plugin_contracts.test_plugin_generate_contracts.custom_generate"
REFERENCE_CUSTOM_GENERATE_WITH_EVAL_PATH = (
"plugin_contracts.test_plugin_generate_contracts.custom_generate_with_evaluation"
)
from slime.rollout.sglang_rollout import generate_and_rm
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=["custom-generate-function-path"])
def make_args(**overrides):
class Args:
partial_rollout = False
mask_offpolicy_in_partial_rollout = False
group_rm = False
custom_generate_function_path = None
sglang_enable_deterministic_inference = False
rollout_seed = 7
n_samples_per_prompt = 2
args = Args()
for key, value in overrides.items():
setattr(args, key, value)
return args
class FakeGenerateState:
def __init__(self, args) -> None:
self.args = args
self.semaphore = types.SimpleNamespace(__aenter__=None)
self.pendings = set()
self.remaining_batch_size = 0
self.aborted = False
self.group_sampling_seeds = [args.rollout_seed + i for i in range(args.n_samples_per_prompt)]
@contextmanager
def dp_rank_context(self):
yield 0
async def custom_generate(args, sample: Sample, sampling_params: dict):
sample.tokens = [11, 12, 13]
sample.response = "generated"
sample.response_length = len(sample.tokens)
sample.reward = 0.25
sample.status = Sample.Status.COMPLETED
return sample
async def custom_generate_with_evaluation(args, sample: Sample, sampling_params: dict, evaluation: bool = False):
sample.tokens = [21, 22]
sample.response = "eval-generated" if evaluation else "train-generated"
sample.response_length = len(sample.tokens)
sample.reward = 0.5 if evaluation else 0.75
sample.status = Sample.Status.COMPLETED
sample.metadata["evaluation"] = evaluation
return sample
def assert_sample_contract(sample: Sample) -> None:
assert isinstance(sample, Sample)
assert isinstance(sample.tokens, list)
assert isinstance(sample.response, str)
assert isinstance(sample.response_length, int)
assert sample.reward is not None
def assert_custom_generate_signature_matches_expected(fn) -> None:
params = tuple(inspect.signature(fn).parameters)
assert params[:3] == ("args", "sample", "sampling_params")
class _DummySemaphore:
async def __aenter__(self):
return None
async def __aexit__(self, exc_type, exc, tb):
return False
class _PatchedGenerateState(FakeGenerateState):
def __init__(self, args):
super().__init__(args)
self.semaphore = _DummySemaphore()
@pytest.fixture
def patch_generate_state(monkeypatch):
"""Patch GenerateState with a test-safe variant; returns the sglang_rollout module."""
from slime.rollout import sglang_rollout
monkeypatch.setattr(sglang_rollout, "GenerateState", _PatchedGenerateState)
return sglang_rollout
def test_generate_and_rm_default_generate_branch_is_stable(patch_generate_state, monkeypatch):
sglang_rollout = patch_generate_state
async def official_default_generate(args, sample: Sample, sampling_params: dict):
sample.tokens = [31, 32]
sample.response = "default-generate"
sample.response_length = 2
sample.reward = 1.0
sample.status = Sample.Status.COMPLETED
return sample
monkeypatch.setattr(sglang_rollout, "generate", official_default_generate)
result = asyncio.run(
generate_and_rm(
make_args(custom_generate_function_path=None),
Sample(index=0, prompt="prompt"),
sampling_params={"temperature": 0.3},
evaluation=False,
)
)
assert_sample_contract(result)
assert result.response == "default-generate"
def test_generate_and_rm_prefers_per_sample_generate_function(patch_generate_state):
args = make_args(custom_generate_function_path=REFERENCE_CUSTOM_GENERATE_PATH)
sample = Sample(index=0, prompt="prompt", generate_function_path=REFERENCE_CUSTOM_GENERATE_WITH_EVAL_PATH)
result = asyncio.run(generate_and_rm(args, sample, sampling_params={"temperature": 0.3}, evaluation=True))
assert_sample_contract(result)
assert result.metadata["evaluation"] is True
def test_custom_generate_function_path_supports_user_override(patch_generate_state):
custom_generate_path = get_contract_path(
"CUSTOM_GENERATE_FUNCTION_PATH",
REFERENCE_CUSTOM_GENERATE_PATH,
)
assert_custom_generate_signature_matches_expected(load_function(custom_generate_path))
result = asyncio.run(
generate_and_rm(
make_args(custom_generate_function_path=custom_generate_path),
Sample(index=0, prompt="prompt"),
sampling_params={"temperature": 0.3},
evaluation=False,
)
)
assert_sample_contract(result)
def test_generate_and_rm_group_rm_accepts_list_result_from_custom_generate(patch_generate_state, monkeypatch):
sglang_rollout = patch_generate_state
async def custom_generate_list(args, sample: Sample, sampling_params: dict):
sample.status = Sample.Status.COMPLETED
sibling = Sample(index=1, prompt="prompt-1", status=Sample.Status.COMPLETED)
return [sample, sibling]
monkeypatch.setattr(sglang_rollout, "load_function", lambda _path: custom_generate_list)
result = asyncio.run(
generate_and_rm(
make_args(custom_generate_function_path="plugin_contracts.fake_generate", group_rm=True),
Sample(index=0, prompt="prompt-0"),
sampling_params={"temperature": 0.3},
evaluation=False,
)
)
assert isinstance(result, list)
assert len(result) == 2
assert all(isinstance(sample, Sample) for sample in result)
if __name__ == "__main__":
run_contract_test_file()
|