minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

Report per-token weight-version spans in generation meta info (#35926)

Browse Source
This commit is contained in:
fzyzcjy
2026-08-24 20:18:52 +08:00
committed by GitHub
parent 981dfa2b83
commit 3b24d8981b
20 changed files with 1607 additions and 25 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/registered/rl/test_weight_version_spans.py
+590
View File
@@ -0,0 +1,590 @@
import json
import time
import unittest
from concurrent.futures import ThreadPoolExecutor
import requests
from sglang.srt.utils import kill_process_tree
from sglang.test.ci.ci_register import register_cuda_ci
from sglang.test.test_utils import (
DEFAULT_MODEL_NAME_FOR_TEST_MLA,
DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
DEFAULT_URL_FOR_TEST,
CustomTestCase,
popen_launch_server,
)
register_cuda_ci(
est_time=180, stage="nightly", runner_config="2-gpu-large", nightly=True
)
_REQUEST_TIMEOUT = 180
def _assert_spans_contiguous(test, meta_info):
spans = meta_info["weight_versions"]
test.assertGreater(len(spans), 0)
test.assertEqual(spans[0]["start"], 0)
for prev, cur in zip(spans, spans[1:]):
test.assertEqual(prev["end"], cur["start"])
test.assertNotEqual(prev["version"], cur["version"])
test.assertEqual(meta_info["weight_version"], spans[-1]["version"])
return spans
class TestWeightVersionSpans(CustomTestCase):
@classmethod
def setUpClass(cls):
cls.model = DEFAULT_MODEL_NAME_FOR_TEST_MLA
cls.base_url = DEFAULT_URL_FOR_TEST
cls.process = popen_launch_server(
cls.model,
base_url=cls.base_url,
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
other_args=[
"--weight-version",
"base-v0",
"--trust-remote-code",
"--tp-size",
"2",
"--dp-size",
"2",
"--enable-dp-attention",
"--speculative-algorithm",
"EAGLE",
"--speculative-draft-model-path",
DEFAULT_MODEL_NAME_FOR_TEST_MLA,
"--speculative-num-steps",
"2",
"--speculative-eagle-topk",
"3",
"--speculative-num-draft-tokens",
"3",
"--cuda-graph-max-bs-decode",
"32",
"--max-running-requests",
"8",
],
)
@classmethod
def tearDownClass(cls):
kill_process_tree(cls.process.pid)
def _generate(self, max_new_tokens: int, prompt: str = "The capital of France is"):
response = requests.post(
f"{self.base_url}/generate",
json={
"text": prompt,
"sampling_params": {
"temperature": 0.8,
"max_new_tokens": max_new_tokens,
"ignore_eos": True,
},
},
timeout=_REQUEST_TIMEOUT,
)
self.assertEqual(response.status_code, 200)
return response.json()
def _current_version(self):
response = requests.get(f"{self.base_url}/get_model_info", timeout=30)
self.assertEqual(response.status_code, 200)
return response.json()["weight_version"]
def _pause(self, mode: str):
requests.post(
f"{self.base_url}/pause_generation",
json={"mode": mode},
timeout=30,
).raise_for_status()
def _continue(self):
requests.post(
f"{self.base_url}/continue_generation",
json={},
timeout=30,
).raise_for_status()
def _set_weight_version(self, new_version: str, abort_all_requests: bool = False):
response = requests.post(
f"{self.base_url}/update_weight_version",
json={
"new_version": new_version,
"abort_all_requests": abort_all_requests,
},
timeout=30,
)
self.assertEqual(response.status_code, 200)
return response.json()
def _update_weights_from_disk(self, **fields) -> None:
response = requests.post(
f"{self.base_url}/update_weights_from_disk",
json={"model_path": self.model, "flush_cache": False, **fields},
timeout=_REQUEST_TIMEOUT,
)
self.assertEqual(response.status_code, 200)
self.assertTrue(response.json()["success"])
def _run_while_paused(
self,
num_requests: int,
while_paused,
mode: str = "retract",
max_new_tokens: int = 1024,
):
with ThreadPoolExecutor(max_workers=num_requests) as executor:
futures = [
executor.submit(
self._generate,
max_new_tokens=max_new_tokens - 16 * i,
prompt=f"Write a long story about the number {i}.",
)
for i in range(num_requests)
]
time.sleep(2)
self._pause(mode)
try:
while_paused()
finally:
self._continue()
return [future.result() for future in futures]
def test_01_single_span_without_update(self):
"""A request untouched by updates reports one span covering all output tokens."""
data = self._generate(max_new_tokens=8)
meta_info = data["meta_info"]
spans = _assert_spans_contiguous(self, meta_info)
self.assertEqual(len(spans), 1)
self.assertEqual(spans[0]["version"], "base-v0")
self.assertEqual(spans[0]["end"], meta_info["completion_tokens"])
self.assertEqual(meta_info["weight_version"], "base-v0")
def test_02_update_weight_version_endpoint_applies_to_new_requests(self):
"""The endpoint returns only once every scheduler stamps new requests with the new version."""
self._set_weight_version("endpoint-v1")
self.assertEqual(self._current_version(), "endpoint-v1")
with ThreadPoolExecutor(max_workers=4) as executor:
futures = [
executor.submit(self._generate, max_new_tokens=8) for _ in range(4)
]
results = [future.result() for future in futures]
for data in results:
meta_info = data["meta_info"]
spans = _assert_spans_contiguous(self, meta_info)
self.assertEqual(len(spans), 1)
self.assertEqual(spans[0]["version"], "endpoint-v1")
def test_03_spans_split_across_pause_update_continue(self):
"""Requests spanning pause -> update_weights_from_disk -> continue report one span per version."""
base_version = self._current_version()
results = self._run_while_paused(
num_requests=4,
while_paused=lambda: self._update_weights_from_disk(
weight_version="disk-v2"
),
)
multi_span_count = 0
for data in results:
meta_info = data["meta_info"]
spans = _assert_spans_contiguous(self, meta_info)
self.assertEqual(spans[-1]["end"], meta_info["completion_tokens"])
versions = [span["version"] for span in spans]
self.assertEqual(versions[0], base_version)
self.assertIn(versions[-1], (base_version, "disk-v2"))
if len(spans) > 1:
multi_span_count += 1
self.assertEqual(versions, [base_version, "disk-v2"])
self.assertGreater(spans[0]["end"], 0)
self.assertGreater(
multi_span_count,
0,
"No request spanned the weight update -- no boundary was recorded.",
)
def test_04_openai_metadata_contains_weight_versions(self):
"""OpenAI-compatible responses surface the spans under response metadata."""
response = requests.post(
f"{self.base_url}/v1/chat/completions",
json={
"model": self.model,
"messages": [{"role": "user", "content": "Hello"}],
"max_tokens": 8,
"temperature": 0.0,
},
timeout=_REQUEST_TIMEOUT,
)
self.assertEqual(response.status_code, 200)
data = response.json()
metadata = data["metadata"]
self.assertIn("weight_versions", metadata)
spans = metadata["weight_versions"]
self.assertEqual(len(spans), 1)
self.assertEqual(spans[0]["version"], metadata["weight_version"])
self.assertEqual(metadata["weight_version"], self._current_version())
self.assertEqual(spans[0]["start"], 0)
self.assertEqual(spans[0]["end"], data["usage"]["completion_tokens"])
def test_04b_openai_metadata_reports_the_first_choice_when_n_is_greater_than_one(
self,
):
"""With n > 1 the single metadata block describes the first choice instead of going missing."""
response = requests.post(
f"{self.base_url}/v1/chat/completions",
json={
"model": self.model,
"messages": [{"role": "user", "content": "Hello"}],
"max_tokens": 8,
"temperature": 0.8,
"n": 2,
},
timeout=_REQUEST_TIMEOUT,
)
self.assertEqual(response.status_code, 200)
data = response.json()
self.assertEqual(len(data["choices"]), 2)
metadata = data["metadata"]
spans = metadata["weight_versions"]
self.assertEqual(len(spans), 1)
self.assertEqual(spans[0]["version"], metadata["weight_version"])
self.assertEqual(metadata["weight_version"], self._current_version())
self.assertEqual(spans[0]["start"], 0)
def test_05_aborted_retracted_requests_report_spans(self):
"""Requests aborted while retracted in the waiting queue still report their spans."""
def abort_all():
requests.post(
f"{self.base_url}/abort_request",
json={"abort_all": True},
timeout=30,
).raise_for_status()
results = self._run_while_paused(num_requests=4, while_paused=abort_all)
aborted_with_spans = 0
for data in results:
meta_info = data["meta_info"]
if meta_info["finish_reason"]["type"] != "abort":
continue
if "weight_versions" not in meta_info:
continue
spans = _assert_spans_contiguous(self, meta_info)
self.assertEqual(spans[-1]["end"], meta_info["completion_tokens"])
aborted_with_spans += 1
self.assertGreater(
aborted_with_spans,
0,
"No aborted request carried weight_versions -- the AbortReq path lost the spans.",
)
def test_06_running_requests_split_without_abort(self):
"""A version bump with abort_all_requests=False splits requests still in the running batch."""
previous_version = self._current_version()
results = self._run_while_paused(
num_requests=4,
while_paused=lambda: self._set_weight_version("inplace-v3"),
mode="in_place",
max_new_tokens=1024,
)
self.assertEqual(self._current_version(), "inplace-v3")
split_count = 0
for data in results:
meta_info = data["meta_info"]
self.assertNotEqual(meta_info["finish_reason"]["type"], "abort")
spans = _assert_spans_contiguous(self, meta_info)
self.assertEqual(spans[-1]["end"], meta_info["completion_tokens"])
self.assertEqual(spans[0]["version"], previous_version)
if len(spans) > 1:
split_count += 1
self.assertEqual(
[span["version"] for span in spans],
[previous_version, "inplace-v3"],
)
self.assertGreater(spans[0]["end"], 0)
self.assertGreater(
split_count,
0,
"No running request was split -- the running batch was not visited.",
)
def test_07_update_without_weight_version_does_not_split(self):
"""A refit that carries no weight_version leaves attribution untouched."""
version = self._current_version()
results = self._run_while_paused(
num_requests=4, while_paused=self._update_weights_from_disk
)
self.assertEqual(self._current_version(), version)
for data in results:
meta_info = data["meta_info"]
spans = _assert_spans_contiguous(self, meta_info)
self.assertEqual(len(spans), 1)
self.assertEqual(spans[0]["version"], version)
self.assertEqual(spans[0]["end"], meta_info["completion_tokens"])
def test_08_reannouncing_current_version_is_a_noop(self):
"""Re-announcing the version the server already has must not split anything."""
version = self._current_version()
results = self._run_while_paused(
num_requests=4,
while_paused=lambda: self._set_weight_version(version),
mode="in_place",
max_new_tokens=128,
)
self.assertEqual(self._current_version(), version)
for data in results:
spans = _assert_spans_contiguous(self, data["meta_info"])
self.assertEqual(len(spans), 1)
self.assertEqual(spans[0]["version"], version)
def test_09_three_spans_across_two_updates(self):
"""Two updates during one request produce three ordered, non-empty spans."""
first_version = self._current_version()
with ThreadPoolExecutor(max_workers=4) as executor:
futures = [
executor.submit(
self._generate,
max_new_tokens=2048,
prompt=f"Write a long story about the number {i}.",
)
for i in range(4)
]
for new_version in ("multi-a", "multi-b"):
time.sleep(2)
self._pause("in_place")
try:
self._set_weight_version(new_version)
finally:
self._continue()
results = [future.result() for future in futures]
expected = [first_version, "multi-a", "multi-b"]
three_span_count = 0
for data in results:
meta_info = data["meta_info"]
spans = _assert_spans_contiguous(self, meta_info)
versions = [span["version"] for span in spans]
self.assertEqual(versions, expected[: len(versions)])
for span in spans:
self.assertGreater(span["end"], span["start"])
self.assertEqual(spans[-1]["end"], meta_info["completion_tokens"])
if len(spans) == 3:
three_span_count += 1
self.assertGreater(
three_span_count,
0,
"No request spanned both updates -- multi-event accumulation was not exercised.",
)
def test_10_abort_before_first_token_reports_empty_span(self):
"""A request aborted before producing a token still reports a well-formed span."""
version = self._current_version()
def abort_all():
requests.post(
f"{self.base_url}/abort_request",
json={"abort_all": True},
timeout=30,
).raise_for_status()
results = self._run_while_paused(
num_requests=16,
while_paused=abort_all,
)
empty_aborts = [
data["meta_info"]
for data in results
if data["meta_info"]["finish_reason"]["type"] == "abort"
and data["meta_info"]["completion_tokens"] == 0
]
self.assertGreater(len(empty_aborts), 0)
for meta_info in empty_aborts:
self.assertEqual(
meta_info["weight_versions"],
[{"version": version, "start": 0, "end": 0}],
)
self.assertEqual(meta_info["weight_version"], version)
def test_11_streaming_reports_spans_only_on_the_final_chunk(self):
"""Intermediate stream chunks carry no spans; the finishing chunk carries them all."""
max_new_tokens = 32
response = requests.post(
f"{self.base_url}/generate",
json={
"text": "The capital of France is",
"sampling_params": {
"temperature": 0.8,
"max_new_tokens": max_new_tokens,
"ignore_eos": True,
},
"stream": True,
},
stream=True,
timeout=_REQUEST_TIMEOUT,
)
self.assertEqual(response.status_code, 200)
chunks = []
for line in response.iter_lines(decode_unicode=True):
if not line or not line.startswith("data:"):
continue
payload = line[len("data:") :].strip()
if payload == "[DONE]":
break
chunks.append(json.loads(payload))
version = self._current_version()
self.assertGreater(len(chunks), 1)
for chunk in chunks[:-1]:
self.assertNotIn("weight_versions", chunk["meta_info"])
self.assertEqual(chunk["meta_info"]["weight_version"], version)
meta_info = chunks[-1]["meta_info"]
spans = _assert_spans_contiguous(self, meta_info)
self.assertEqual(spans[-1]["end"], meta_info["completion_tokens"])
self.assertEqual(spans[-1]["end"], max_new_tokens)
def test_12_completions_endpoint_reports_metadata(self):
"""/v1/completions surfaces the spans the same way /v1/chat/completions does."""
response = requests.post(
f"{self.base_url}/v1/completions",
json={
"model": self.model,
"prompt": "The capital of France is",
"max_tokens": 8,
"temperature": 0.0,
},
timeout=_REQUEST_TIMEOUT,
)
self.assertEqual(response.status_code, 200)
data = response.json()
metadata = data["metadata"]
self.assertEqual(metadata["weight_version"], self._current_version())
spans = metadata["weight_versions"]
self.assertEqual(len(spans), 1)
self.assertEqual(spans[0]["version"], metadata["weight_version"])
self.assertEqual(spans[0]["start"], 0)
self.assertEqual(spans[0]["end"], data["usage"]["completion_tokens"])
def test_13_new_requests_after_the_endpoint_returns_see_the_new_version(self):
"""Once the endpoint returns, every concurrently admitted request stamps the new version."""
self._set_weight_version("dp-v1")
self.assertEqual(self._current_version(), "dp-v1")
with ThreadPoolExecutor(max_workers=8) as executor:
futures = [
executor.submit(self._generate, max_new_tokens=8) for _ in range(8)
]
results = [future.result() for future in futures]
for data in results:
spans = _assert_spans_contiguous(self, data["meta_info"])
self.assertEqual(len(spans), 1)
self.assertEqual(spans[0]["version"], "dp-v1")
def test_14_all_inflight_requests_are_swept(self):
"""A version change must split every in-flight request, not just the one the sweep happens to reach first."""
previous_version = self._current_version()
with ThreadPoolExecutor(max_workers=8) as executor:
futures = [
executor.submit(
self._generate,
max_new_tokens=1024 - 16 * i,
prompt=f"Write a long story about the number {i}.",
)
for i in range(8)
]
time.sleep(2)
self._set_weight_version("dp-v2")
results = [future.result() for future in futures]
split_count = 0
for data in results:
meta_info = data["meta_info"]
spans = _assert_spans_contiguous(self, meta_info)
self.assertEqual(spans[-1]["end"], meta_info["completion_tokens"])
self.assertEqual(spans[0]["version"], previous_version)
if len(spans) > 1:
split_count += 1
self.assertEqual(
[span["version"] for span in spans],
[previous_version, "dp-v2"],
)
self.assertGreater(spans[0]["end"], 0)
self.assertGreater(
split_count,
1,
"At most one in-flight request was split -- an attention-DP rank was "
"likely missed by the weight version sweep.",
)
def test_15_single_span_covers_all_accepted_tokens(self):
"""Draft-token overshoot must not leak past the reported completion length."""
data = self._generate(max_new_tokens=32)
meta_info = data["meta_info"]
spans = _assert_spans_contiguous(self, meta_info)
self.assertEqual(len(spans), 1)
self.assertEqual(spans[0]["version"], self._current_version())
self.assertEqual(spans[0]["end"], meta_info["completion_tokens"])
def test_16_retract_update_continue_keeps_exact_boundaries(self):
"""Spans stay contiguous and clamped when a retract-paused update lands mid-speculation."""
previous_version = self._current_version()
results = self._run_while_paused(
num_requests=4,
while_paused=lambda: self._set_weight_version("spec-v1"),
)
split_count = 0
for data in results:
meta_info = data["meta_info"]
spans = _assert_spans_contiguous(self, meta_info)
self.assertEqual(spans[-1]["end"], meta_info["completion_tokens"])
self.assertEqual(spans[0]["version"], previous_version)
if len(spans) > 1:
split_count += 1
self.assertEqual(
[span["version"] for span in spans],
[previous_version, "spec-v1"],
)
self.assertGreater(spans[0]["end"], 0)
self.assertGreater(
split_count,
0,
"No request spanned the update -- the retract boundary under "
"speculative decoding was not recorded.",
)
if __name__ == "__main__":
unittest.main()
test/registered/unit/managers/test_msgpack_ipc_roundtrip.py
+31
View File
@@ -15,6 +15,7 @@ import msgspec
from sglang.srt.managers import io_struct
from sglang.srt.managers.io_struct import (
AbortReq,
BackupDramReq,
ChecksumInfo,
CheckWeightsReqOutput,
@@ -37,6 +38,7 @@ from sglang.srt.model_executor.cuda_graph_config import CudaGraphConfig
from sglang.srt.utils.msgspec_utils import msgspec_to_builtins
from sglang.srt.utils.weight_checker import ChecksumInfo as PydanticChecksumInfo
from sglang.srt.utils.weight_checker import ParallelismInfo as PydanticParallelismInfo
from sglang.srt.utils.weight_versions import WeightVersionSpan
from sglang.test.ci.ci_register import register_cpu_ci
from sglang.test.test_utils import CustomTestCase
@@ -237,5 +239,34 @@ class TestMsgpackIpcRoundtrip(CustomTestCase):
self.assertEqual(_round_trip(output), output)
class TestWeightVersionSpansRoundTrip(CustomTestCase):
"""The per-request weight-version spans ride the same msgpack IPC path."""
def test_abort_req_carries_spans(self):
"""A scheduler-side abort keeps its spans across the wire."""
obj = AbortReq(
rid="r0",
weight_versions=[
WeightVersionSpan(version="v1", start=0, end=3),
WeightVersionSpan(version="v2", start=3, end=7),
],
)
decoded = _double_hop(obj)
self.assertEqual(
decoded.weight_versions,
[
WeightVersionSpan(version="v1", start=0, end=3),
WeightVersionSpan(version="v2", start=3, end=7),
],
)
self.assertIsInstance(decoded.weight_versions[0], WeightVersionSpan)
def test_abort_req_defaults_to_no_spans(self):
"""The field is optional on the wire, so an abort without spans decodes to None."""
self.assertIsNone(_round_trip(AbortReq(rid="r0")).weight_versions)
if __name__ == "__main__":
unittest.main()
test/registered/unit/managers/test_multi_tokenizer_mixin.py
+33
View File
@@ -1,5 +1,6 @@
import unittest
from sglang.srt.utils.weight_versions import WeightVersionSpan
from sglang.test.ci.ci_register import register_cpu_ci
from sglang.test.test_utils import maybe_stub_sgl_kernel
@@ -76,6 +77,13 @@ def _make_batch_str_output() -> BatchStrOutput:
placeholder_tokens_idx=[None, None],
placeholder_tokens_val=[None, None],
retraction_counts=[0, 0],
weight_versions=[
[
WeightVersionSpan(version="v1", start=0, end=3),
WeightVersionSpan(version="v2", start=3, end=5),
],
[WeightVersionSpan(version="v2", start=0, end=2)],
],
)
@@ -92,6 +100,31 @@ class TestMultiTokenizerMixin(unittest.TestCase):
[{"device": 1, "host": 3}],
)
def test_batch_str_output_keeps_weight_versions_nested_per_request(self):
"""Per-request segment lists stay one level nested after the split."""
output = _make_batch_str_output()
self.assertEqual(
_handle_output_by_index(output, 0).weight_versions,
[
[
WeightVersionSpan(version="v1", start=0, end=3),
WeightVersionSpan(version="v2", start=3, end=5),
]
],
)
self.assertEqual(
_handle_output_by_index(output, 1).weight_versions,
[[WeightVersionSpan(version="v2", start=0, end=2)]],
)
def test_batch_str_output_without_weight_versions_stays_none(self):
"""An output from an older server without the field splits into None."""
output = _make_batch_str_output()
output.weight_versions = None
self.assertIsNone(_handle_output_by_index(output, 0).weight_versions)
def test_get_tokenizer_worker_class_uses_default(self):
self.assertIs(get_tokenizer_worker_class(DefaultServerArgs()), TokenizerWorker)
test/registered/unit/managers/test_output_streamer_customized_info.py
+55 -16
View File
@@ -9,6 +9,10 @@ from sglang.srt.managers.scheduler_components.output_streamer import (
_GenerationStreamAccumulator,
)
from sglang.srt.speculative.spec_info import SpeculativeAlgorithm
from sglang.srt.utils.weight_versions import (
WeightVersionSpan,
record_weight_version_events,
)
from sglang.test.ci.ci_register import register_cpu_ci
register_cpu_ci(est_time=1, suite="base-a-test-cpu")
@@ -58,6 +62,7 @@ class _FakeReq:
self.mm_video_tokens = 0
self.multimodal_inputs = None
self.customized_info = customized_info
self.weight_version_events = []
def finished(self):
return self._finished
@@ -69,11 +74,26 @@ class _FakeReq:
return False
def _accumulator(current_weight_version="default"):
return _GenerationStreamAccumulator(
return_logprob=False,
return_hidden_states=False,
return_routed_experts=False,
return_indexer_topk=False,
spec_algorithm=SpeculativeAlgorithm.NONE,
disaggregation_mode=DisaggregationMode.NULL,
default_stream_interval=1,
default_force_stream_interval=1,
get_cached_tokens_details=lambda req: None,
current_weight_version=current_weight_version,
)
class TestOutputStreamerCustomizedInfo(unittest.TestCase):
def setUp(self):
serving_patch = patch(
"sglang.srt.managers.scheduler_components.output_streamer.get_serving",
return_value=SimpleNamespace(stream_interval=1),
return_value=SimpleNamespace(stream_interval=1, weight_version="default"),
)
observability_patch = patch(
"sglang.srt.managers.scheduler_components.output_streamer.get_observability",
@@ -84,22 +104,8 @@ class TestOutputStreamerCustomizedInfo(unittest.TestCase):
self.addCleanup(serving_patch.stop)
self.addCleanup(observability_patch.stop)
@staticmethod
def _accumulator():
return _GenerationStreamAccumulator(
return_logprob=False,
return_hidden_states=False,
return_routed_experts=False,
return_indexer_topk=False,
spec_algorithm=SpeculativeAlgorithm.NONE,
disaggregation_mode=DisaggregationMode.NULL,
default_stream_interval=1,
default_force_stream_interval=1,
get_cached_tokens_details=lambda req: None,
)
def test_customized_info_is_padded_for_mixed_batches(self):
accumulator = self._accumulator()
accumulator = _accumulator()
accumulator.accept(req=_FakeReq("r0", [10, 11]))
accumulator.accept(
@@ -334,5 +340,38 @@ class TestOutputStreamerCustomizedInfo(unittest.TestCase):
)
class TestOutputStreamerWeightVersions(unittest.TestCase):
def test_payload_carries_spans_for_finished_requests(self):
"""Finished requests report their spans; still-generating ones report nothing."""
streaming_req = _FakeReq("r0", [10, 11])
finished_req = _FakeReq("r1", [20, 21, 22], finished=True)
record_weight_version_events([finished_req], old_version="v1")
finished_req.output_ids.extend([23, 24])
accumulator = _accumulator(current_weight_version="v2")
accumulator.accept(req=streaming_req)
accumulator.accept(req=finished_req)
payload = accumulator.to_payload(dp_rank=0, is_idle_batch=False)
self.assertEqual(
payload.weight_versions,
[
None,
[
WeightVersionSpan(version="v1", start=0, end=3),
WeightVersionSpan(version="v2", start=3, end=5),
],
],
)
def test_payload_omits_spans_while_all_requests_stream(self):
"""A batch of unfinished requests puts nothing on the wire."""
accumulator = _accumulator(current_weight_version="v2")
accumulator.accept(req=_FakeReq("r0", [10]))
payload = accumulator.to_payload(dp_rank=0, is_idle_batch=False)
self.assertIsNone(payload.weight_versions)
if __name__ == "__main__":
unittest.main()
test/registered/unit/managers/test_output_streamer_logprobs.py
+1
View File
@@ -73,6 +73,7 @@ def _make_accumulator() -> _GenerationStreamAccumulator:
default_stream_interval=1,
default_force_stream_interval=1,
get_cached_tokens_details=lambda req: None,
current_weight_version=None,
)
test/registered/unit/managers/test_priority_scheduling_disaggregation.py
+7 -1
View File
@@ -41,6 +41,8 @@ class TestDisaggregationPriorityQueueing(unittest.TestCase):
req = MagicMock()
req.priority = priority
req.rid = "req"
req.output_ids = []
req.weight_version_events = []
req.time_stats = MagicMock()
req.time_stats.trace_ctx = MagicMock()
return req
@@ -73,7 +75,11 @@ class TestDisaggregationPriorityQueueing(unittest.TestCase):
scheduler.abort_on_priority_when_disabled = True
req = self._new_req(priority=10)
scheduler._add_request_to_queue(req)
with patch(
"sglang.srt.managers.scheduler.get_serving",
return_value=SimpleNamespace(weight_version="v0"),
):
scheduler._add_request_to_queue(req)
scheduler.disagg_decode_prealloc_queue.add.assert_not_called()
scheduler.ipc_channels.send_to_tokenizer.send_output.assert_called_once()
test/registered/unit/managers/test_scheduler_timeouts.py
+11 -1
View File
@@ -9,7 +9,7 @@ scheduler/test_scheduler_control.py.
import time
import unittest
from types import SimpleNamespace
from unittest.mock import MagicMock
from unittest.mock import MagicMock, patch
from sglang.srt.environ import envs
from sglang.test.ci.ci_register import register_cpu_ci
@@ -29,6 +29,8 @@ class _FakeReq:
self.rid = rid
self.to_finish = None
self._finished = is_finished
self.output_ids = []
self.weight_version_events = []
self.time_stats = SimpleNamespace(
wait_queue_entry_time=wait_entry,
forward_entry_time=forward_entry,
@@ -54,6 +56,14 @@ def _scheduler(waiting_queue):
class TestWaitingTimeout(CustomTestCase):
def setUp(self):
patcher = patch(
"sglang.srt.managers.scheduler.get_serving",
return_value=SimpleNamespace(weight_version="v0"),
)
patcher.start()
self.addCleanup(patcher.stop)
def test_drops_only_reqs_past_the_deadline(self):
now = time.perf_counter()
stale = _req("stale", wait_entry=now - 10)
test/registered/unit/managers/test_scheduler_weight_version_tracking.py
+37 -3
View File
@@ -37,11 +37,27 @@ class TestSchedulerRecordWeightVersionChange(CustomTestCase):
self.addCleanup(patcher.stop)
return serving
def _scheduler(
self, *, inflight=(), waiting=(), chunked=None, staging=()
) -> SimpleNamespace:
return SimpleNamespace(
collect_inflight_reqs=lambda: set(inflight),
waiting_queue=list(waiting),
chunked_req=chunked,
hisparse_coordinator=(
SimpleNamespace(
ack_staging_queue=[SimpleNamespace(req=req) for req in staging]
)
if staging
else None
),
)
def test_a_new_version_is_adopted(self):
"""The scheduler has to end up on the version it was told about, or nothing downstream can read it."""
serving = self._serving("v1")
Scheduler.record_weight_version_change(SimpleNamespace(), new_version="v2")
Scheduler.record_weight_version_change(self._scheduler(), new_version="v2")
self.assertEqual(serving.weight_version, "v2")
@@ -49,7 +65,7 @@ class TestSchedulerRecordWeightVersionChange(CustomTestCase):
"""Re-announcing the current version must not be treated as a change."""
serving = self._serving("v1")
Scheduler.record_weight_version_change(SimpleNamespace(), new_version="v1")
Scheduler.record_weight_version_change(self._scheduler(), new_version="v1")
self.assertEqual(serving.weight_version, "v1")
@@ -57,10 +73,28 @@ class TestSchedulerRecordWeightVersionChange(CustomTestCase):
"""An update that carries no version must leave the recorded one alone."""
serving = self._serving("v1")
Scheduler.record_weight_version_change(SimpleNamespace(), new_version=None)
Scheduler.record_weight_version_change(self._scheduler(), new_version=None)
self.assertEqual(serving.weight_version, "v1")
def test_every_source_of_live_requests_is_stamped(self):
"""A request missed here keeps attributing its next tokens to the superseded version."""
self._serving("v1")
inflight, queued, chunked, staged = (object() for _ in range(4))
scheduler = self._scheduler(
inflight=[inflight], waiting=[queued], chunked=chunked, staging=[staged]
)
with patch(
"sglang.srt.managers.scheduler.record_weight_version_events",
return_value=0,
) as recorder:
Scheduler.record_weight_version_change(scheduler, new_version="v2")
self.assertEqual(
set(recorder.call_args.args[0]), {inflight, queued, chunked, staged}
)
class TestRecordWeightVersionAfterUpdate(CustomTestCase):
def _updater(
test/registered/unit/utils/test_weight_versions.py
+629
View File
@@ -0,0 +1,629 @@
import random
import unittest
from types import SimpleNamespace
from unittest.mock import patch
from sglang.srt.managers.io_struct import AbortReq
from sglang.srt.managers.scheduler import Scheduler, _make_abort_req
from sglang.srt.utils.weight_versions import (
WeightVersionEvent,
WeightVersionSpan,
add_weight_versions_to_meta_info,
build_endpoint_weight_version_metadata,
compute_weight_version_spans,
record_weight_version_events,
truncate_weight_version_events,
)
from sglang.test.ci.ci_register import register_cpu_ci
from sglang.test.test_utils import CustomTestCase
register_cpu_ci(est_time=4, suite="base-a-test-cpu")
class _ReqStub:
def __init__(self, output_len: int):
self.output_ids = [0] * output_len
self.weight_version_events = []
def record_weight_version_change(self, old_version):
record_weight_version_events([self], old_version=old_version)
def compute_weight_version_spans(self, current_version, num_output_tokens):
return compute_weight_version_spans(
self.weight_version_events,
current_version=current_version,
num_output_tokens=num_output_tokens,
)
def _expected_spans(events, current_version, num_output_tokens):
"""Reference model: name the version that sampled each token, then run-length encode."""
per_token = []
for index in range(num_output_tokens):
owner = next(
(event.old_version for event in events if event.num_output_tokens > index),
current_version,
)
per_token.append(owner)
if not per_token:
first_event_end_at_zero = next(
(event for event in events if event.num_output_tokens >= 0), None
)
version = (
first_event_end_at_zero.old_version
if first_event_end_at_zero is not None
else current_version
)
return [WeightVersionSpan(version=version, start=0, end=0)]
spans = []
for index, version in enumerate(per_token):
if spans and spans[-1].version == version:
spans[-1].end = index + 1
else:
spans.append(WeightVersionSpan(version=version, start=index, end=index + 1))
return spans
class TestComputeWeightVersionSpans(CustomTestCase):
def test_no_events_returns_single_span(self):
"""A request untouched by updates gets one span covering all output tokens."""
self.assertEqual(
_ReqStub(5).compute_weight_version_spans(
current_version="v1", num_output_tokens=5
),
[WeightVersionSpan(version="v1", start=0, end=5)],
)
def test_zero_output_tokens_returns_empty_span_span(self):
"""A request finishing with no output still reports the current version."""
self.assertEqual(
_ReqStub(0).compute_weight_version_spans(
current_version="v1", num_output_tokens=0
),
[WeightVersionSpan(version="v1", start=0, end=0)],
)
def test_one_update_splits_into_two_spans(self):
"""An update at 3 output tokens attributes [0,3) to the old version and the rest to the new."""
req = _ReqStub(3)
req.record_weight_version_change(old_version="v1")
req.output_ids.extend([0] * 4)
self.assertEqual(
req.compute_weight_version_spans(current_version="v2", num_output_tokens=7),
[
WeightVersionSpan(version="v1", start=0, end=3),
WeightVersionSpan(version="v2", start=3, end=7),
],
)
def test_two_updates_split_into_three_spans(self):
"""Each update the request lives through adds one span."""
req = _ReqStub(2)
req.record_weight_version_change(old_version="v1")
req.output_ids.extend([0] * 3)
req.record_weight_version_change(old_version="v2")
req.output_ids.extend([0] * 1)
self.assertEqual(
req.compute_weight_version_spans(current_version="v3", num_output_tokens=6),
[
WeightVersionSpan(version="v1", start=0, end=2),
WeightVersionSpan(version="v2", start=2, end=5),
WeightVersionSpan(version="v3", start=5, end=6),
],
)
def test_update_before_first_token_records_no_event(self):
"""An update while the request has no output leaves the whole output on the new version."""
req = _ReqStub(0)
req.record_weight_version_change(old_version="v1")
self.assertEqual(req.weight_version_events, [])
req.output_ids.extend([0] * 4)
self.assertEqual(
req.compute_weight_version_spans(current_version="v2", num_output_tokens=4),
[WeightVersionSpan(version="v2", start=0, end=4)],
)
def test_back_to_back_updates_skip_empty_span(self):
"""Two updates with no tokens in between produce no empty span."""
req = _ReqStub(2)
req.record_weight_version_change(old_version="v1")
req.record_weight_version_change(old_version="v2")
req.output_ids.extend([0] * 2)
self.assertEqual(
req.compute_weight_version_spans(current_version="v3", num_output_tokens=4),
[
WeightVersionSpan(version="v1", start=0, end=2),
WeightVersionSpan(version="v3", start=2, end=4),
],
)
def test_update_at_final_length_yields_no_trailing_span(self):
"""An event recorded exactly at the final output length adds no empty trailing span."""
req = _ReqStub(4)
req.record_weight_version_change(old_version="v1")
self.assertEqual(
req.compute_weight_version_spans(current_version="v2", num_output_tokens=4),
[WeightVersionSpan(version="v1", start=0, end=4)],
)
def test_event_beyond_reported_length_is_clamped(self):
"""A spec-decode overshoot event is clamped to the reported output length."""
req = _ReqStub(6)
req.record_weight_version_change(old_version="v1")
self.assertEqual(
req.compute_weight_version_spans(current_version="v2", num_output_tokens=4),
[WeightVersionSpan(version="v1", start=0, end=4)],
)
def test_clamp_to_zero_reports_the_pre_update_version(self):
"""With no visible tokens the single span keeps the version that sampled them."""
req = _ReqStub(3)
req.record_weight_version_change(old_version="v1")
self.assertEqual(
req.compute_weight_version_spans(current_version="v2", num_output_tokens=0),
[WeightVersionSpan(version="v1", start=0, end=0)],
)
def test_clamped_output_can_end_before_the_current_version(self):
"""When the visible tokens stop early the newest version never appears."""
req = _ReqStub(3)
req.record_weight_version_change(old_version="v1")
req.output_ids.extend([0] * 2)
req.record_weight_version_change(old_version="v2")
self.assertEqual(
req.compute_weight_version_spans(current_version="v3", num_output_tokens=4),
[
WeightVersionSpan(version="v1", start=0, end=3),
WeightVersionSpan(version="v2", start=3, end=4),
],
)
def test_clamping_collapses_events_and_merges_into_one_span(self):
"""Events beyond the visible range collapse instead of leaving empty spans."""
req = _ReqStub(3)
req.record_weight_version_change(old_version="v1")
req.output_ids.extend([0] * 2)
req.record_weight_version_change(old_version="v2")
self.assertEqual(
req.compute_weight_version_spans(current_version="v1", num_output_tokens=2),
[WeightVersionSpan(version="v1", start=0, end=2)],
)
def test_duplicate_event_at_the_same_length_is_idempotent(self):
"""Recording twice at the same token count matches recording once."""
req = _ReqStub(3)
req.record_weight_version_change(old_version="v1")
req.record_weight_version_change(old_version="v1")
req.output_ids.extend([0] * 2)
self.assertEqual(
req.compute_weight_version_spans(current_version="v2", num_output_tokens=5),
[
WeightVersionSpan(version="v1", start=0, end=3),
WeightVersionSpan(version="v2", start=3, end=5),
],
)
def test_version_returning_after_tokens_keeps_separate_spans(self):
"""A version that comes back after another one sampled tokens gets its own span."""
req = _ReqStub(2)
req.record_weight_version_change(old_version="v1")
req.output_ids.extend([0] * 2)
req.record_weight_version_change(old_version="v2")
req.output_ids.extend([0] * 2)
self.assertEqual(
req.compute_weight_version_spans(current_version="v1", num_output_tokens=6),
[
WeightVersionSpan(version="v1", start=0, end=2),
WeightVersionSpan(version="v2", start=2, end=4),
WeightVersionSpan(version="v1", start=4, end=6),
],
)
def test_spans_are_never_empty_even_when_everything_is_clamped_away(self):
"""add_weight_versions_to_meta_info reads the newest span by index, so an empty list would crash it."""
req = _ReqStub(3)
req.record_weight_version_change(old_version="v1")
req.record_weight_version_change(old_version="v2")
for current_version in ("v1", "v3"):
self.assertTrue(
req.compute_weight_version_spans(
current_version=current_version, num_output_tokens=0
)
)
self.assertTrue(
_ReqStub(0).compute_weight_version_spans(
current_version="v1", num_output_tokens=0
)
)
def test_spans_satisfy_the_contract_for_random_event_sequences(self):
"""Randomized event sequences always yield ordered, contiguous, non-duplicated spans."""
rng = random.Random(0)
versions = ["v0", "v1", "v2"]
for _ in range(300):
counts = sorted(rng.randint(1, 8) for _ in range(rng.randint(0, 5)))
events = [
WeightVersionEvent(
old_version=rng.choice(versions), num_output_tokens=count
)
for count in counts
]
current_version = rng.choice(versions)
num_output_tokens = rng.randint(0, 10)
spans = compute_weight_version_spans(
events,
current_version=current_version,
num_output_tokens=num_output_tokens,
)
self.assertEqual(
spans,
_expected_spans(events, current_version, num_output_tokens),
)
self.assertTrue(spans)
self.assertEqual(spans[0].start, 0)
self.assertEqual(spans[-1].end, num_output_tokens)
for previous, current in zip(spans, spans[1:]):
self.assertEqual(previous.end, current.start)
self.assertNotEqual(previous.version, current.version)
if num_output_tokens == 0:
self.assertEqual(len(spans), 1)
self.assertEqual(spans[0].end, 0)
else:
for span in spans:
self.assertLess(span.start, span.end)
class _ServingStub:
def __init__(self, weight_version):
self.weight_version = weight_version
class _ContextStub:
def __init__(self, serving: _ServingStub):
self.serving = serving
def override(self, source, **fields):
self.serving.weight_version = fields["weight_version"]
class _SchedulerStub:
collect_inflight_reqs = Scheduler.collect_inflight_reqs
def __init__(
self,
version,
running,
waiting,
chunked=None,
last_batch=None,
pp_size=1,
hisparse=None,
):
self.serving = _ServingStub(version)
self.ps = SimpleNamespace(pp_size=pp_size)
self.running_batch = SimpleNamespace(reqs=running)
self.last_batch = last_batch
self.waiting_queue = waiting
self.chunked_req = chunked
self.hisparse_coordinator = hisparse
class TestSchedulerRecordWeightVersionChange(CustomTestCase):
def _scheduler(self, *args, **kwargs):
scheduler = _SchedulerStub(*args, **kwargs)
for name, value in (
("get_serving", scheduler.serving),
("get_context", _ContextStub(scheduler.serving)),
):
patcher = patch(f"sglang.srt.managers.scheduler.{name}", return_value=value)
patcher.start()
self.addCleanup(patcher.stop)
return scheduler
def test_records_on_running_waiting_and_chunked_requests(self):
"""A version change records an event on every live request holding output tokens."""
running_req = _ReqStub(3)
waiting_req = _ReqStub(5)
chunked_req = _ReqStub(2)
scheduler = self._scheduler(
"v1", [running_req], [waiting_req], chunked=chunked_req
)
Scheduler.record_weight_version_change(scheduler, new_version="v2")
self.assertEqual(scheduler.serving.weight_version, "v2")
for req, output_len in (
(running_req, 3),
(waiting_req, 5),
(chunked_req, 2),
):
self.assertEqual(len(req.weight_version_events), 1)
self.assertEqual(req.weight_version_events[0].old_version, "v1")
self.assertEqual(req.weight_version_events[0].num_output_tokens, output_len)
def test_records_on_last_batch_requests_exactly_once(self):
"""A request in both the running and last batch is recorded once, and last-batch-only requests are recorded too."""
shared_req = _ReqStub(3)
prefill_req = _ReqStub(1)
scheduler = self._scheduler(
"v1",
[shared_req],
[],
last_batch=SimpleNamespace(reqs=[shared_req, prefill_req]),
)
Scheduler.record_weight_version_change(scheduler, new_version="v2")
self.assertEqual(len(shared_req.weight_version_events), 1)
self.assertEqual(len(prefill_req.weight_version_events), 1)
def test_records_on_all_pp_microbatches(self):
"""With pipeline parallelism every microbatch is visited, not just the selected one."""
mb0_req = _ReqStub(2)
mb1_req = _ReqStub(4)
pending_req = _ReqStub(6)
scheduler = self._scheduler("v1", [], [], pp_size=2)
scheduler.running_mbs = [
SimpleNamespace(reqs=[mb0_req]),
SimpleNamespace(reqs=[mb1_req]),
]
scheduler.mbs = [None, SimpleNamespace(reqs=[pending_req])]
Scheduler.record_weight_version_change(scheduler, new_version="v2")
for req in (mb0_req, mb1_req, pending_req):
self.assertEqual(len(req.weight_version_events), 1)
def test_records_on_hisparse_staging_requests(self):
"""Requests parked in the HiSparse staging queue are swept like any live request."""
staging_req = _ReqStub(2)
coordinator = SimpleNamespace(
ack_staging_queue=[SimpleNamespace(req=staging_req)]
)
scheduler = self._scheduler("v1", [], [], hisparse=coordinator)
Scheduler.record_weight_version_change(scheduler, new_version="v2")
self.assertEqual(len(staging_req.weight_version_events), 1)
self.assertEqual(staging_req.weight_version_events[0].old_version, "v1")
self.assertEqual(staging_req.weight_version_events[0].num_output_tokens, 2)
def test_same_version_is_a_noop(self):
"""Re-announcing the current version must not record events."""
req = _ReqStub(3)
scheduler = self._scheduler("v1", [req], [])
Scheduler.record_weight_version_change(scheduler, new_version="v1")
self.assertEqual(scheduler.serving.weight_version, "v1")
self.assertEqual(req.weight_version_events, [])
def test_none_version_is_a_noop(self):
"""An update without a weight version must not disturb attribution."""
req = _ReqStub(3)
scheduler = self._scheduler("v1", [req], [])
Scheduler.record_weight_version_change(scheduler, new_version=None)
self.assertEqual(scheduler.serving.weight_version, "v1")
self.assertEqual(req.weight_version_events, [])
class TestRecordWeightVersionEvents(CustomTestCase):
def test_records_only_for_requests_that_have_output(self):
"""Requests without output tokens are skipped without stopping the sweep."""
empty_req = _ReqStub(0)
started_req = _ReqStub(4)
num_recorded = record_weight_version_events(
[empty_req, started_req, _ReqStub(2)], old_version="v1"
)
self.assertEqual(num_recorded, 2)
self.assertEqual(empty_req.weight_version_events, [])
self.assertEqual(len(started_req.weight_version_events), 1)
self.assertEqual(started_req.weight_version_events[0].old_version, "v1")
self.assertEqual(started_req.weight_version_events[0].num_output_tokens, 4)
class TestTruncateWeightVersionEvents(CustomTestCase):
def test_events_within_the_kept_prefix_are_preserved(self):
"""Events entirely inside the streamed prefix survive a retract untouched."""
events = [WeightVersionEvent(old_version="v1", num_output_tokens=3)]
self.assertEqual(
truncate_weight_version_events(events, num_kept_tokens=5), events
)
def test_events_beyond_the_kept_prefix_are_clamped(self):
"""An event past the streamed prefix is clamped so re-generated tokens forget it."""
events = [
WeightVersionEvent(old_version="v1", num_output_tokens=3),
WeightVersionEvent(old_version="v2", num_output_tokens=9),
]
self.assertEqual(
truncate_weight_version_events(events, num_kept_tokens=5),
[
WeightVersionEvent(old_version="v1", num_output_tokens=3),
WeightVersionEvent(old_version="v2", num_output_tokens=5),
],
)
def test_zero_kept_tokens_drops_all_events(self):
"""With nothing streamed yet the whole history is discarded, as before the fix."""
events = [WeightVersionEvent(old_version="v1", num_output_tokens=3)]
self.assertEqual(truncate_weight_version_events(events, num_kept_tokens=0), [])
class TestSpanlessAbortPaths(CustomTestCase):
def test_priority_disabled_rejection_attaches_spans(self):
"""The pre-scheduler priority rejection reports spans like every other abort."""
req = _ReqStub(0)
req.rid = "r0"
req.priority = 5
req.time_stats = SimpleNamespace(
trace_ctx=SimpleNamespace(abort=lambda abort_info: None)
)
sent = []
scheduler = SimpleNamespace(
enable_priority_scheduling=False,
abort_on_priority_when_disabled=True,
ipc_channels=SimpleNamespace(
send_to_tokenizer=SimpleNamespace(
send_output=lambda obj, req_arg: sent.append(obj)
)
),
)
with patch(
"sglang.srt.managers.scheduler.get_serving",
return_value=SimpleNamespace(weight_version="v1"),
):
accepted = Scheduler._set_or_validate_priority(scheduler, req)
self.assertFalse(accepted)
self.assertEqual(
sent[0].weight_versions, [WeightVersionSpan(version="v1", start=0, end=0)]
)
class TestMakeAbortReq(CustomTestCase):
def test_abort_req_carries_spans_clamped_to_sampled_tokens(self):
"""An aborted request reports the versions that sampled the tokens it produced."""
req = _ReqStub(4)
req.rid = "r0"
req.weight_version_events.append(
WeightVersionEvent(old_version="v1", num_output_tokens=3)
)
with patch(
"sglang.srt.managers.scheduler.get_serving",
return_value=SimpleNamespace(weight_version="v2"),
):
abort_req = _make_abort_req(req)
self.assertIsInstance(abort_req, AbortReq)
self.assertEqual(abort_req.rid, "r0")
self.assertEqual(
abort_req.weight_versions,
[
WeightVersionSpan(version="v1", start=0, end=3),
WeightVersionSpan(version="v2", start=3, end=4),
],
)
class TestBuildEndpointWeightVersionMetadata(CustomTestCase):
def test_metadata_projects_only_the_weight_fields(self):
"""Endpoint metadata exposes the version fields and nothing else from meta_info."""
spans = [
{"version": "v1", "start": 0, "end": 3},
{"version": "v2", "start": 3, "end": 7},
]
metadata = build_endpoint_weight_version_metadata(
{
"weight_version": "v2",
"weight_versions": spans,
"id": "r0",
"completion_tokens": 7,
}
)
self.assertEqual(metadata, {"weight_version": "v2", "weight_versions": spans})
def test_metadata_omits_spans_when_absent(self):
"""Responses without spans still report the legacy version alone."""
metadata = build_endpoint_weight_version_metadata({"weight_version": "v2"})
self.assertEqual(metadata, {"weight_version": "v2"})
class TestAddWeightVersionsToMetaInfo(CustomTestCase):
def test_meta_info_gets_dict_spans_and_legacy_field(self):
"""Spans become dicts and the legacy weight_version reports the newest span."""
meta_info = {"weight_version": "stale"}
add_weight_versions_to_meta_info(
meta_info,
[
WeightVersionSpan(version="v1", start=0, end=3),
WeightVersionSpan(version="v2", start=3, end=7),
],
num_output_tokens=7,
)
self.assertEqual(
meta_info["weight_versions"],
[
{"version": "v1", "start": 0, "end": 3},
{"version": "v2", "start": 3, "end": 7},
],
)
self.assertEqual(meta_info["weight_version"], "v2")
def test_spans_are_clamped_to_returned_tokens(self):
"""Aborts returning fewer tokens than sampled clamp spans to the visible range."""
meta_info = {}
add_weight_versions_to_meta_info(
meta_info,
[
WeightVersionSpan(version="v1", start=0, end=3),
WeightVersionSpan(version="v2", start=3, end=7),
],
num_output_tokens=5,
)
self.assertEqual(
meta_info["weight_versions"],
[
{"version": "v1", "start": 0, "end": 3},
{"version": "v2", "start": 3, "end": 5},
],
)
self.assertEqual(meta_info["weight_version"], "v2")
def test_clamp_never_extends_spans(self):
"""A response longer than the sampled range leaves the spans untouched."""
meta_info = {}
add_weight_versions_to_meta_info(
meta_info,
[WeightVersionSpan(version="v1", start=0, end=3)],
num_output_tokens=10,
)
self.assertEqual(
meta_info["weight_versions"], [{"version": "v1", "start": 0, "end": 3}]
)
self.assertEqual(meta_info["weight_version"], "v1")
def test_clamp_drops_trailing_spans_and_rewrites_the_legacy_field(self):
"""Dropping invisible spans also moves the legacy version back."""
meta_info = {}
add_weight_versions_to_meta_info(
meta_info,
[
WeightVersionSpan(version="v1", start=0, end=3),
WeightVersionSpan(version="v2", start=3, end=7),
],
num_output_tokens=3,
)
self.assertEqual(
meta_info["weight_versions"], [{"version": "v1", "start": 0, "end": 3}]
)
self.assertEqual(meta_info["weight_version"], "v1")
def test_clamp_to_zero_tokens_keeps_a_degenerate_span(self):
"""A response with no visible tokens still reports a well-formed empty span."""
meta_info = {}
add_weight_versions_to_meta_info(
meta_info,
[
WeightVersionSpan(version="v1", start=0, end=3),
WeightVersionSpan(version="v2", start=3, end=7),
],
num_output_tokens=0,
)
self.assertEqual(
meta_info["weight_versions"], [{"version": "v1", "start": 0, "end": 0}]
)
self.assertEqual(meta_info["weight_version"], "v1")
if __name__ == "__main__":
unittest.main()