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

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
+19 -13
View File
@@ -149,7 +149,7 @@ def test_main(
for with_topk in (False, True):
if local_rank == 0:
print(
f'[testing] Running with {"FP8" if isinstance(current_x, tuple) else "BF16"}, {"with" if with_topk else "without"} top-k (async={async_mode}, previous={previous_mode}) ...',
f"[testing] Running with {'FP8' if isinstance(current_x, tuple) else 'BF16'}, {'with' if with_topk else 'without'} top-k (async={async_mode}, previous={previous_mode}) ...",
flush=True,
end="",
)
@@ -192,9 +192,9 @@ def test_main(
# Checks
recv_gbl_rank_prefix_sum = handle[-4]
assert gbl_num_tokens_per_rank[rank].item() == recv_x.size(
0
), f"{gbl_num_tokens_per_rank[rank].item()} != {recv_x.size(0)}"
assert gbl_num_tokens_per_rank[rank].item() == recv_x.size(0), (
f"{gbl_num_tokens_per_rank[rank].item()} != {recv_x.size(0)}"
)
assert (
gbl_num_tokens_per_expert.view(num_ranks, -1)[rank].tolist()
== recv_num_tokens_per_expert_list
@@ -316,10 +316,13 @@ def test_main(
tune_args = {"x": current_x, "handle": handle, "config": config}
t = bench(lambda: buffer.dispatch(**tune_args))[0]
if t < best_time:
best_time, best_results = t, (
num_sms,
nvl_chunk_size,
rdma_chunk_size,
best_time, best_results = (
t,
(
num_sms,
nvl_chunk_size,
rdma_chunk_size,
),
)
if local_rank == 0:
print(
@@ -328,7 +331,7 @@ def test_main(
)
if local_rank == 0:
print(
f'[tuning] Best dispatch ({"FP8" if isinstance(current_x, tuple) else "BF16"}): SMs {best_results[0]}, NVL chunk {best_results[1]}, RDMA chunk {best_results[2]}: {rdma_send_bytes / 1e9 / best_time:.2f} GB/s (RDMA), {nvl_recv_bytes / 1e9 / best_time:.2f} GB/s (NVL)',
f"[tuning] Best dispatch ({'FP8' if isinstance(current_x, tuple) else 'BF16'}): SMs {best_results[0]}, NVL chunk {best_results[1]}, RDMA chunk {best_results[2]}: {rdma_send_bytes / 1e9 / best_time:.2f} GB/s (RDMA), {nvl_recv_bytes / 1e9 / best_time:.2f} GB/s (NVL)",
flush=True,
)
print("", flush=True)
@@ -385,10 +388,13 @@ def test_main(
flush=True,
)
if t < best_time:
best_time, best_results = t, (
num_sms,
nvl_chunk_size,
rdma_chunk_size,
best_time, best_results = (
t,
(
num_sms,
nvl_chunk_size,
rdma_chunk_size,
),
)
if local_rank == 0:
+5 -5
View File
@@ -120,7 +120,7 @@ def test_main(
for with_topk in (False, True):
if local_rank == 0:
print(
f'[testing] Running with {"FP8" if isinstance(current_x, tuple) else "BF16"}, {"with" if with_topk else "without"} top-k (async={async_mode}, previous={previous_mode}) ...',
f"[testing] Running with {'FP8' if isinstance(current_x, tuple) else 'BF16'}, {'with' if with_topk else 'without'} top-k (async={async_mode}, previous={previous_mode}) ...",
flush=True,
end="",
)
@@ -162,9 +162,9 @@ def test_main(
# Checks
rank_prefix_matrix = handle[0]
assert gbl_num_tokens_per_rank[rank].item() == recv_x.size(
0
), f"{gbl_num_tokens_per_rank[rank].item()} != {recv_x.size(0)}"
assert gbl_num_tokens_per_rank[rank].item() == recv_x.size(0), (
f"{gbl_num_tokens_per_rank[rank].item()} != {recv_x.size(0)}"
)
assert (
gbl_num_tokens_per_expert.view(num_ranks, -1)[rank].tolist()
== recv_num_tokens_per_expert_list
@@ -280,7 +280,7 @@ def test_main(
)
if local_rank == 0:
print(
f'[tuning] Best dispatch ({"FP8" if isinstance(current_x, tuple) else "BF16"}): SMs {best_results[0]}, NVL chunk {best_results[1]}, {nvl_recv_bytes / 1e9 / best_time:.2f} GB/s (NVL)',
f"[tuning] Best dispatch ({'FP8' if isinstance(current_x, tuple) else 'BF16'}): SMs {best_results[0]}, NVL chunk {best_results[1]}, {nvl_recv_bytes / 1e9 / best_time:.2f} GB/s (NVL)",
flush=True,
)
print("", flush=True)
+19 -18
View File
@@ -36,9 +36,9 @@ def test_main(
# NOTES: the integers greater than 256 exceeds the BF16 precision limit
rank_offset = 128
assert (
num_ranks - rank_offset < 257
), "Too many ranks (exceeding test precision limit)"
assert num_ranks - rank_offset < 257, (
"Too many ranks (exceeding test precision limit)"
)
x = torch.ones((num_tokens, hidden), dtype=torch.bfloat16, device="cuda") * (
rank - rank_offset
@@ -55,9 +55,9 @@ def test_main(
# Randomly mask some positions
for i in range(10):
topk_idx[random.randint(0, num_tokens - 1), random.randint(0, num_topk - 1)] = (
-1
)
topk_idx[
random.randint(0, num_tokens - 1), random.randint(0, num_topk - 1)
] = -1
# Check dispatch correctness
do_check = True
@@ -114,9 +114,9 @@ def test_main(
assert (
num_valid_tokens == (recv_layout_range & int_mask).sum().item()
), f"{num_valid_tokens} != {recv_layout_range & int_mask}.sum().item()"
assert (
num_valid_tokens == (all_topk_idx == expert_id).sum().item()
), f"{num_valid_tokens} != {(all_topk_idx == expert_id).sum().item()}"
assert num_valid_tokens == (all_topk_idx == expert_id).sum().item(), (
f"{num_valid_tokens} != {(all_topk_idx == expert_id).sum().item()}"
)
# Check received data
recv_x = recv_x[:num_valid_tokens]
@@ -127,9 +127,10 @@ def test_main(
recv_x[:, -128:] - recv_src_info.view(-1, 1) % num_tokens
).sum().item() == 0
for j in range(num_ranks):
begin_idx, count = (recv_layout_range[j] >> 32).item(), (
recv_layout_range[j] & int_mask
).item()
begin_idx, count = (
(recv_layout_range[j] >> 32).item(),
(recv_layout_range[j] & int_mask).item(),
)
assert (recv_x_amin == j - rank_offset).sum().item() == (
all_topk_idx[j] == expert_id
).sum().item()
@@ -145,9 +146,9 @@ def test_main(
# Check combine correctness
for zero_copy in (False, True):
if zero_copy:
buffer.get_next_low_latency_combine_buffer(handle)[
:, :, :
] = simulated_gemm_x
buffer.get_next_low_latency_combine_buffer(handle)[:, :, :] = (
simulated_gemm_x
)
out = torch.empty(
(num_tokens, hidden), dtype=torch.bfloat16, device="cuda"
)
@@ -203,9 +204,9 @@ def test_main(
)
large_gemm_with_hook(hook) if return_recv_hook else None
if zero_copy:
buffer.get_next_low_latency_combine_buffer(handle)[
:, :, :
] = simulated_gemm_x
buffer.get_next_low_latency_combine_buffer(handle)[:, :, :] = (
simulated_gemm_x
)
combined_x, event, hook = buffer.low_latency_combine(
simulated_gemm_x,
topk_idx,
+1 -2
View File
@@ -238,8 +238,7 @@ class _ElasticScaleUpEndToEndBase(CustomTestCase):
join_end = rank_offset + join_tp
if join_end > len(visible_devices):
raise RuntimeError(
f"Scale-up requires {join_end} visible GPUs, got "
f"{len(visible_devices)}"
f"Scale-up requires {join_end} visible GPUs, got {len(visible_devices)}"
)
env["CUDA_VISIBLE_DEVICES"] = ",".join(visible_devices[rank_offset:join_end])
base_joining_log = os.environ.get(
+2 -5
View File
@@ -16,7 +16,6 @@ from sglang.test.test_utils import CustomTestCase
class TestFlashinferDispatcher(CustomTestCase):
@classmethod
def setUpClass(cls):
server_args = ServerArgs(model_path="dummy")
@@ -125,8 +124,7 @@ class TestFlashinferDispatcher(CustomTestCase):
self.assertTrue(
torch.all(
received_hidden_states[
expected_source_rank
* num_tokens : (expected_source_rank + 1)
expected_source_rank * num_tokens : (expected_source_rank + 1)
* num_tokens
]
== 100.0 + expected_source_rank
@@ -229,8 +227,7 @@ class TestFlashinferDispatcher(CustomTestCase):
self.assertTrue(
torch.all(
received_hidden_states[
expected_source_rank
* num_tokens : (expected_source_rank + 1)
expected_source_rank * num_tokens : (expected_source_rank + 1)
* num_tokens
]
== 100.0 + expected_source_rank