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

[CI][RFC] Replace black-jupyter with ruff-format (#37210)

Browse Source 
Co-authored-by: Alison Shao <a.shao@wustl.edu>
This commit is contained in:
Alex Nails
2026-09-02 19:46:08 -07:00
committed by GitHub
co-authored by Alison Shao
parent 2641e427be
commit 28262c20df
1411 changed files with 7766 additions and 8176 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/registered/unit/mem_cache/test_hicache_file_lru_unit.py
+8 -6
View File
@@ -391,9 +391,10 @@ class TestMinFreeSpaceWatermark(HiCacheFileLRUTestBase):
free[0] += os.path.getsize(p)
return original_remove(p)
with mock.patch.object(
b._evictor, "_fs_stats", side_effect=fake_fs_stats
), mock.patch("os.remove", side_effect=tracked_remove):
with (
mock.patch.object(b._evictor, "_fs_stats", side_effect=fake_fs_stats),
mock.patch("os.remove", side_effect=tracked_remove),
):
self.assertTrue(b.set("newk", _t(60)))
self.assertFalse(b.exists("victim"))
self.assertTrue(b.exists("newk"))
@@ -516,9 +517,10 @@ class TestHiCacheFileMetadataIntegration(HiCacheFileLRUTestBase):
b.set("k2", _t(50))
# Now patch os.scandir and os.path.exists
with mock.patch("os.scandir") as mock_scandir, mock.patch(
"os.path.exists"
) as mock_exists:
with (
mock.patch("os.scandir") as mock_scandir,
mock.patch("os.path.exists") as mock_exists,
):
mock_exists.return_value = True
# batch_exists_v2 for k1 and k2 should hit the metadata cache and NOT call os.scandir or os.path.exists
test/registered/unit/mem_cache/test_hybrid_pool_assembler.py
+1 -3
View File
@@ -49,9 +49,7 @@ class TestDeviceAllocEviction(CustomTestCase):
def test_sufficient_capacity_skips_eviction(self):
cache = MagicMock()
cache.token_to_kv_pool_allocator.swa_available_size.return_value = 10
cache.req_to_token_pool.mamba_allocator.schedulable_available_size.return_value = (
10
)
cache.req_to_token_pool.mamba_allocator.schedulable_available_size.return_value = 10
_evict_swa_for_device_alloc(cache, required_size=10)
_evict_mamba_for_device_alloc(cache, required_size=10)
test/registered/unit/mem_cache/test_kv_index_translator.py
-1
View File
@@ -182,7 +182,6 @@ def _alloc_and_fill(allocator, ps, lens):
class TestReadTableBuild(unittest.TestCase):
def test_read_table_matches_reference_across_multipliers(self):
"""The load-bearing formula pin: full AND swa read tables equal
the independent per-element derivation, across page sizes and both
test/registered/unit/mem_cache/test_mamba_donated_alloc_ratio.py
+3 -1
View File
@@ -117,7 +117,9 @@ class TestMambaRatioEnvGate(unittest.TestCase):
strategy = (
"extra_buffer_lazy"
if lazy
else "extra_buffer" if extra_buffer else "no_buffer"
else "extra_buffer"
if extra_buffer
else "no_buffer"
)
from sglang.srt import runtime_context as rc
test/registered/unit/mem_cache/test_mamba_path_state_cap.py
+6 -3
View File
@@ -105,9 +105,12 @@ class TestMambaPathStateCap(unittest.TestCase):
def test_server_arg_rejects_zero_and_values_below_negative_one(self):
for value in (0, -2):
args = ServerArgs(model_path="dummy", mamba_max_states_per_path=value)
with self.subTest(value=value), self.assertRaisesRegex(
ValueError,
"must be -1 \\(unlimited\\) or a positive integer",
with (
self.subTest(value=value),
self.assertRaisesRegex(
ValueError,
"must be -1 \\(unlimited\\) or a positive integer",
),
):
handle_mamba_backend(args)
test/registered/unit/mem_cache/test_mamba_unittest.py
+9 -9
View File
@@ -300,17 +300,17 @@ class TestMamba(unittest.TestCase):
full_num_tokens = 1
print(f"evicting {full_num_tokens} full token")
result = tree.evict(EvictParams(num_tokens=full_num_tokens))
assert (
result.num_tokens_evicted >= full_num_tokens
), f"evicted {result.num_tokens_evicted} full tokens, expected {full_num_tokens}"
assert result.num_tokens_evicted >= full_num_tokens, (
f"evicted {result.num_tokens_evicted} full tokens, expected {full_num_tokens}"
)
tree.pretty_print()
mamba_num = 1
print(f"evicting {mamba_num} mamba")
result = tree.evict(EvictParams(num_tokens=0, mamba_num=mamba_num))
assert (
result.mamba_num_evicted >= mamba_num
), f"evicted {result.mamba_num_evicted} mamba states, expected {mamba_num}"
assert result.mamba_num_evicted >= mamba_num, (
f"evicted {result.mamba_num_evicted} mamba states, expected {mamba_num}"
)
tree.pretty_print()
req5_token_ids = [1, 2, 3, 4, 5]
@@ -348,9 +348,9 @@ class TestMamba(unittest.TestCase):
mamba_num = 1
print(f"evicting {mamba_num} mamba")
result = tree.evict(EvictParams(num_tokens=0, mamba_num=mamba_num))
assert (
result.mamba_num_evicted >= mamba_num
), f"evicted {result.mamba_num_evicted} mamba states, expected {mamba_num}"
assert result.mamba_num_evicted >= mamba_num, (
f"evicted {result.mamba_num_evicted} mamba states, expected {mamba_num}"
)
tree.pretty_print()
req8_token_ids = [1, 2, 3, 4, 5, 60, 70]
test/registered/unit/mem_cache/test_mem_pool_host.py
+14 -7
View File
@@ -242,10 +242,11 @@ class TestHostMemoryBudget(CustomTestCase):
# Deliberate single-accessor stub: isolates the budget math from the
# topology derivation, which the ranks_per_host case below covers.
fake_mem = unittest.mock.Mock(available=self._AVAILABLE)
with unittest.mock.patch.object(
base, "ranks_per_host", return_value=ranks
), unittest.mock.patch.object(
base.psutil, "virtual_memory", return_value=fake_mem
with (
unittest.mock.patch.object(base, "ranks_per_host", return_value=ranks),
unittest.mock.patch.object(
base.psutil, "virtual_memory", return_value=fake_mem
),
):
return base.host_memory_budget_bytes()
@@ -264,9 +265,15 @@ class TestHostMemoryBudget(CustomTestCase):
# The launcher slices ranks uniformly across nodes, so the co-located
# rank count is world_size // nnodes — no hostname collective.
fake_group = unittest.mock.Mock(world_size=16)
with get_context().override_server_args(nnodes=2), unittest.mock.patch.object(
torch.distributed, "is_initialized", return_value=True
), unittest.mock.patch.object(base, "get_world_group", return_value=fake_group):
with (
get_context().override_server_args(nnodes=2),
unittest.mock.patch.object(
torch.distributed, "is_initialized", return_value=True
),
unittest.mock.patch.object(
base, "get_world_group", return_value=fake_group
),
):
self.assertEqual(base.ranks_per_host(), 8)
test/registered/unit/mem_cache/test_mmap_allocator.py
+6 -3
View File
@@ -83,9 +83,12 @@ class TestMmapAllocator(unittest.TestCase):
# MAP_POPULATE is unreachable on a 5.14+ kernel, so CI never runs it;
# force the branch or it ships untested.
with self.subTest(path="map_populate"), unittest.mock.patch(
"sglang.srt.mem_cache.storage.mmap.mmap_allocator._has_madv_populate_write",
return_value=False,
with (
self.subTest(path="map_populate"),
unittest.mock.patch(
"sglang.srt.mem_cache.storage.mmap.mmap_allocator._has_madv_populate_write",
return_value=False,
),
):
mm = _mmap_prefaulted(-1, alloc_bytes, flags)
try:
test/registered/unit/mem_cache/test_multi_ended_allocator.py
+10 -6
View File
@@ -3141,12 +3141,16 @@ class TestFloatMultiEndedAllocator(unittest.TestCase):
_, _, fla, _, kv = self._build_tri()
v = fla.alloc(4)
self._stamp(fla, kv, v)
with mock.patch.object(
torch.Tensor, "tolist", side_effect=AssertionError("tolist = D2H")
), mock.patch.object(
torch.Tensor, "item", side_effect=AssertionError("item = D2H")
), mock.patch.object(
torch, "unique", side_effect=AssertionError("unique = host sync")
with (
mock.patch.object(
torch.Tensor, "tolist", side_effect=AssertionError("tolist = D2H")
),
mock.patch.object(
torch.Tensor, "item", side_effect=AssertionError("item = D2H")
),
mock.patch.object(
torch, "unique", side_effect=AssertionError("unique = host sync")
),
):
fla.free(v[:2], _pages=v[:2])
test/registered/unit/mem_cache/test_radix_cache_slru_accuracy.py
-1
View File
@@ -18,7 +18,6 @@ register_cpu_ci(est_time=8, suite="base-a-test-cpu")
class TestSLRUAccuracy(unittest.TestCase):
def setUp(self):
"""Setup minimal memory pools for testing"""
torch.set_default_device(None)
test/registered/unit/mem_cache/test_rust_tree_core_integration.py
+3 -2
View File
@@ -882,8 +882,9 @@ def test_prefetch_node_accessors_round_trip():
assert not core.is_backuped(leaf)
assert not core.is_root(leaf)
assert core.get_last_hash_value(leaf) == (
mem_cache.get_hash_str(array("q", [1, 2]), None, 2)[-1]
assert (
core.get_last_hash_value(leaf)
== (mem_cache.get_hash_str(array("q", [1, 2]), None, 2)[-1])
)
assert core.get_prefix_hash_values(leaf) == []
test/registered/unit/mem_cache/test_session_token_share_unit.py
-1
View File
@@ -53,7 +53,6 @@ def _recv(rid, input_ids, max_new_tokens=8):
class TestSessionTokenShare(CustomTestCase):
def setUp(self):
self.session = Session(capacity_of_str_len=0, session_id="s", streaming=True)
test/registered/unit/mem_cache/test_swa_eviction_boundary.py
-1
View File
@@ -134,7 +134,6 @@ def _make_batch(tree, allocator, pool):
class TestSWAEvictionBoundary(unittest.TestCase):
# -- Eviction formula: page_size > window --
def test_formula_page_gt_window_sweep(self):
test/registered/unit/mem_cache/test_swa_unittest.py
+9 -9
View File
@@ -727,9 +727,9 @@ class TestSWA(unittest.TestCase):
EvictParams(num_tokens=full_num_tokens, swa_num_tokens=swa_num_tokens)
)
assert isinstance(evict_result, EvictResult)
assert (
evict_result.swa_num_tokens_evicted >= swa_num_tokens
), f"evicted {evict_result.swa_num_tokens_evicted} swa tokens, expected {swa_num_tokens}"
assert evict_result.swa_num_tokens_evicted >= swa_num_tokens, (
f"evicted {evict_result.swa_num_tokens_evicted} swa tokens, expected {swa_num_tokens}"
)
tree.pretty_print()
full_num_tokens, swa_num_tokens = 1, 2
@@ -738,12 +738,12 @@ class TestSWA(unittest.TestCase):
EvictParams(num_tokens=full_num_tokens, swa_num_tokens=swa_num_tokens)
)
assert isinstance(evict_result, EvictResult)
assert (
evict_result.num_tokens_evicted >= full_num_tokens
), f"evicted {evict_result.num_tokens_evicted} full tokens, expected {full_num_tokens}"
assert (
evict_result.swa_num_tokens_evicted >= swa_num_tokens
), f"evicted {evict_result.swa_num_tokens_evicted} swa tokens, expected {swa_num_tokens}"
assert evict_result.num_tokens_evicted >= full_num_tokens, (
f"evicted {evict_result.num_tokens_evicted} full tokens, expected {full_num_tokens}"
)
assert evict_result.swa_num_tokens_evicted >= swa_num_tokens, (
f"evicted {evict_result.swa_num_tokens_evicted} swa tokens, expected {swa_num_tokens}"
)
tree.pretty_print()
req5_token_ids = [1, 2, 3, 4, 5]
test/registered/unit/mem_cache/test_unified_free_no_host_sync.py
+7 -4
View File
@@ -457,10 +457,13 @@ class TestFreeSwaWindowRatchetNoHostSync(unittest.TestCase):
alloc = self._swa_composite(lazy=True)
v = alloc.alloc(8 * self.PS)
self.assertIsNotNone(v)
with mock.patch.object(
torch, "unique", side_effect=AssertionError("unique = host sync")
), mock.patch.object(
torch.Tensor, "item", side_effect=AssertionError("item = host sync")
with (
mock.patch.object(
torch, "unique", side_effect=AssertionError("unique = host sync")
),
mock.patch.object(
torch.Tensor, "item", side_effect=AssertionError("item = host sync")
),
):
alloc.free_swa(v[: 4 * self.PS], start_pos=0)
alloc.free_swa(v[4 * self.PS :], start_pos=4 * self.PS)
test/registered/unit/mem_cache/test_unified_radix_cache_bench.py
+3 -3
View File
@@ -424,9 +424,9 @@ def bench_api(
(excluded from latency measurement).
"""
items = setup_fn()
assert (
len(items) >= num_ops + warmup
), f"need {num_ops + warmup} items, got {len(items)}"
assert len(items) >= num_ops + warmup, (
f"need {num_ops + warmup} items, got {len(items)}"
)
for i in range(warmup):
op_fn(items[i])
test/registered/unit/mem_cache/test_unified_radix_cache_unittest.py
+51 -30
View File
@@ -299,11 +299,11 @@ class TestUnifiedTreeCoreLoadBackPending(CustomTestCase):
core.node_by_id.side_effect = nodes.__getitem__
core.components_by_type = {ComponentType.FULL: mock.Mock()}
core.full_host_duplicates = {}
core._is_settled_full_host_duplicate.side_effect = (
lambda node: UnifiedTreeCore._is_settled_full_host_duplicate(core, node)
core._is_settled_full_host_duplicate.side_effect = lambda node: (
UnifiedTreeCore._is_settled_full_host_duplicate(core, node)
)
core._update_duplicate_tracking.side_effect = (
lambda node: UnifiedTreeCore._update_duplicate_tracking(core, node)
core._update_duplicate_tracking.side_effect = lambda node: (
UnifiedTreeCore._update_duplicate_tracking(core, node)
)
return core, shared, anchor_a, anchor_b
@@ -1265,7 +1265,6 @@ class TestUnifiedRadixCacheKVEvents(CustomTestCase):
class UnifiedRadixCacheSuite:
cfg: CacheConfig
_rid: int = 0
@@ -3572,8 +3571,9 @@ class UnifiedRadixCacheSuite:
self.assertIn(n, pipeline.inflight_backup_node_ids)
self._pump_hicache_until(
cache,
lambda: not pipeline.inflight_backup_node_ids
and not pipeline.ongoing_backup,
lambda: (
not pipeline.inflight_backup_node_ids and not pipeline.ongoing_backup
),
"buffer backup pipeline did not drain",
)
@@ -3689,8 +3689,10 @@ class UnifiedRadixCacheSuite:
self.assertEqual((stats["attempts"], stats["issued"]), (1, 1))
self._pump_hicache_until(
cons,
lambda: cons.check_prefetch_progress(req_id)
and cons.buffer_pipeline.has_staged(req_id),
lambda: (
cons.check_prefetch_progress(req_id)
and cons.buffer_pipeline.has_staged(req_id)
),
"prefetch did not stage",
)
# Staged: bounce occupies host staging; nothing device-side; span
@@ -3753,8 +3755,10 @@ class UnifiedRadixCacheSuite:
# loaded KV bytes equal the producer's.
self._pump_hicache_until(
cons,
lambda: not cons.buffer_pipeline.ongoing_buffer_load_back
and self._host_avail_sizes(cons) == avail0,
lambda: (
not cons.buffer_pipeline.ongoing_buffer_load_back
and self._host_avail_sizes(cons) == avail0
),
"load-back ack did not free the bounce",
)
mc = cons.match_prefix(MatchPrefixParams(key=RadixKey(array("q", seq))))
@@ -3818,8 +3822,10 @@ class UnifiedRadixCacheSuite:
)
self._pump_hicache_until(
cons,
lambda: cons.check_prefetch_progress(root_req)
and cons.buffer_pipeline.has_staged(root_req),
lambda: (
cons.check_prefetch_progress(root_req)
and cons.buffer_pipeline.has_staged(root_req)
),
"salted root prefetch did not stage",
)
held = cons.buffer_pipeline.staged_prefetches[root_req]
@@ -3888,8 +3894,10 @@ class UnifiedRadixCacheSuite:
)
self._pump_hicache_until(
cons2,
lambda: cons2.check_prefetch_progress(anchored_req)
and cons2.buffer_pipeline.has_staged(anchored_req),
lambda: (
cons2.check_prefetch_progress(anchored_req)
and cons2.buffer_pipeline.has_staged(anchored_req)
),
"salted mid-tree prefetch did not stage",
)
self._consume_staged_prefetch(
@@ -3950,8 +3958,10 @@ class UnifiedRadixCacheSuite:
)
self._pump_hicache_until(
cons,
lambda: cons.check_prefetch_progress(req_id)
and cons.buffer_pipeline.has_staged(req_id),
lambda: (
cons.check_prefetch_progress(req_id)
and cons.buffer_pipeline.has_staged(req_id)
),
"retried prefetch did not stage",
)
self.assertFalse(cons.pop_storage_prefetch_miss(req_id))
@@ -4047,8 +4057,10 @@ class UnifiedRadixCacheSuite:
)
self._pump_hicache_until(
cons,
lambda: cons.check_prefetch_progress(req_id)
and cons.buffer_pipeline.has_staged(req_id),
lambda: (
cons.check_prefetch_progress(req_id)
and cons.buffer_pipeline.has_staged(req_id)
),
"prefetch did not stage",
)
@@ -4151,8 +4163,10 @@ class UnifiedRadixCacheSuite:
)
self._pump_hicache_until(
cons,
lambda: cons.check_prefetch_progress(req_id)
and cons.buffer_pipeline.has_staged(req_id),
lambda: (
cons.check_prefetch_progress(req_id)
and cons.buffer_pipeline.has_staged(req_id)
),
"prefetch did not stage",
)
cons.pop_prefetch_loaded_tokens(req_id)
@@ -4222,8 +4236,10 @@ class UnifiedRadixCacheSuite:
)
self._pump_hicache_until(
cons,
lambda: cons.check_prefetch_progress(req_id)
and cons.buffer_pipeline.has_staged(req_id),
lambda: (
cons.check_prefetch_progress(req_id)
and cons.buffer_pipeline.has_staged(req_id)
),
"prefetch did not stage",
)
cons.pop_prefetch_loaded_tokens(req_id)
@@ -4331,8 +4347,10 @@ class UnifiedRadixCacheSuite:
)
self._pump_hicache_until(
cons,
lambda: cons.check_prefetch_progress(req_id)
and cons.buffer_pipeline.has_staged(req_id),
lambda: (
cons.check_prefetch_progress(req_id)
and cons.buffer_pipeline.has_staged(req_id)
),
"prefetch did not stage",
)
cons.pop_prefetch_loaded_tokens(req_id)
@@ -4398,8 +4416,10 @@ class UnifiedRadixCacheSuite:
)
self._pump_hicache_until(
cons,
lambda: cons.check_prefetch_progress(req_id)
and cons.buffer_pipeline.has_staged(req_id),
lambda: (
cons.check_prefetch_progress(req_id)
and cons.buffer_pipeline.has_staged(req_id)
),
"prefetch did not stage",
)
cons.pop_prefetch_loaded_tokens(req_id)
@@ -4519,8 +4539,10 @@ class UnifiedRadixCacheSuite:
)
self._pump_hicache_until(
cons2,
lambda: cons2.check_prefetch_progress("subwin-req")
and cons2.buffer_pipeline.has_staged("subwin-req"),
lambda: (
cons2.check_prefetch_progress("subwin-req")
and cons2.buffer_pipeline.has_staged("subwin-req")
),
"sub-window prefetch did not stage",
)
self.assertTrue(
@@ -7081,7 +7103,6 @@ class UnifiedRadixCacheSuite:
class UnifiedLRUListBoundedRefreshTest(CustomTestCase):
components = (ComponentType.FULL, ComponentType.SWA)
def _make_node(self, key_len: int) -> UnifiedTreeNode:
test/registered/unit/mem_cache/test_unified_tri_pool.py
+7 -4
View File
@@ -465,10 +465,13 @@ class TestTriFreeSwaNoHostSync(unittest.TestCase):
self.assertIsNotNone(v)
from unittest import mock
with mock.patch.object(
torch, "unique", side_effect=AssertionError("unique = host sync")
), mock.patch.object(
torch.Tensor, "item", side_effect=AssertionError("item = host sync")
with (
mock.patch.object(
torch, "unique", side_effect=AssertionError("unique = host sync")
),
mock.patch.object(
torch.Tensor, "item", side_effect=AssertionError("item = host sync")
),
):
alloc.free_swa(v[: 4 * self.PS], start_pos=0)
self.assertEqual(alloc.verify_byte_accounting(), [])