[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
+3 -3
View File
@@ -191,9 +191,9 @@ def bench_kineto(
kernel_names = (kernel_names,) if isinstance(kernel_names, str) else kernel_names
assert all([isinstance(name, str) for name in kernel_names])
for name in kernel_names:
assert (
sum([name in line for line in prof_lines]) == 1
), f"Errors of the kernel {name} in the profiling table"
assert sum([name in line for line in prof_lines]) == 1, (
f"Errors of the kernel {name} in the profiling table"
)
# Save chrome traces
if trace_path is not None:
+21 -15
View File
@@ -155,7 +155,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="",
)
@@ -198,9 +198,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
@@ -325,11 +325,14 @@ 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,
config_kwargs,
best_time, best_results = (
t,
(
num_sms,
nvl_chunk_size,
rdma_chunk_size,
config_kwargs,
),
)
if local_rank == 0:
print(
@@ -338,7 +341,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)
@@ -399,11 +402,14 @@ def test_main(
flush=True,
)
if t < best_time:
best_time, best_results = t, (
num_sms,
nvl_chunk_size,
rdma_chunk_size,
config_kwargs,
best_time, best_results = (
t,
(
num_sms,
nvl_chunk_size,
rdma_chunk_size,
config_kwargs,
),
)
if local_rank == 0: