51 lines
2.0 KiB
Python
51 lines
2.0 KiB
Python
"""install_oracle_sampler registration into sglang's sampler-backend registry.
|
|
|
|
Instantiating the registered _OracleSampler factory requires a live distributed (TP) group
|
|
plus a populated global ServerArgs, so the forward-path behavior of _OracleSampler is covered
|
|
by the e2e harness rather than this unit file. Here we only assert the registration-side
|
|
contract: the backend name shows up in the registry / choice set, and second install replaces
|
|
the factory with one bound to the new oracle.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import os
|
|
import unittest
|
|
|
|
os.environ["SGLANG_KV_CANARY_ENABLE_TOKEN_ORACLE"] = "1"
|
|
|
|
from sglang.srt.kv_canary.token_oracle.oracle import HashOracle
|
|
from sglang.srt.kv_canary.token_oracle.sampler import install_oracle_sampler
|
|
from sglang.srt.layers.sampler import _CUSTOM_SAMPLER_FACTORIES
|
|
from sglang.srt.server_args import SAMPLING_BACKEND_CHOICES
|
|
from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
|
|
from sglang.test.test_utils import CustomTestCase
|
|
|
|
register_cuda_ci(est_time=60, stage="extra-a", runner_config="1-gpu-small")
|
|
register_amd_ci(est_time=60, suite="extra-a-test-1-gpu-small-amd")
|
|
|
|
|
|
class TestInstallOracleSampler(CustomTestCase):
|
|
def test_install_oracle_sampler_twice_returns_distinct_hooks_with_replaced_oracle(
|
|
self,
|
|
) -> None:
|
|
"""Verify reinstalling the oracle sampler replaces the registered factory."""
|
|
oracle_a = HashOracle(vocab_size=100)
|
|
oracle_b = HashOracle(vocab_size=100)
|
|
|
|
hook_a = install_oracle_sampler(oracle=oracle_a)
|
|
self.assertIn("token_oracle", _CUSTOM_SAMPLER_FACTORIES)
|
|
self.assertIn("token_oracle", SAMPLING_BACKEND_CHOICES)
|
|
factory_a = _CUSTOM_SAMPLER_FACTORIES["token_oracle"]
|
|
self.assertIs(hook_a.oracle, oracle_a)
|
|
|
|
hook_b = install_oracle_sampler(oracle=oracle_b)
|
|
factory_b = _CUSTOM_SAMPLER_FACTORIES["token_oracle"]
|
|
self.assertIs(hook_b.oracle, oracle_b)
|
|
self.assertIsNot(hook_a, hook_b)
|
|
self.assertIsNot(factory_a, factory_b)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|