[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
@@ -473,8 +473,7 @@ class TestFormatComparisonRichMinimal:
assert_rich_tags_balanced(result)
assert result == (
"[red]❌[/] [bold red]hidden_states [/] "
"rel_diff=5.00e-01"
"[red]❌[/] [bold red]hidden_states [/] rel_diff=5.00e-01"
)
def test_shape_mismatch(self) -> None:
@@ -1046,7 +1045,7 @@ class TestFormatAbsDiffPercentilesRich:
result: str = _format_abs_diff_percentiles_rich(diff)
assert result == (
"p1=1.00e-04 p5=1.00e-04 p50=2.00e-04 " "p95=4.00e-04 p99=5.00e-04"
"p1=1.00e-04 p5=1.00e-04 p50=2.00e-04 p95=4.00e-04 p99=5.00e-04"
)
def test_high_p99_coloring(self) -> None:
@@ -1123,7 +1122,7 @@ class TestFormatReplicatedChecks:
result: str = format_replicated_checks(checks)
assert result == (
"Replicated checks:\n" " ✅ axis=tp group=0 idx=1 vs 0: n/a diff"
"Replicated checks:\n ✅ axis=tp group=0 idx=1 vs 0: n/a diff"
)
@@ -3299,9 +3299,9 @@ def _create_thd_cp_zigzag_dumps(
# Dump each rank
for cp_rank in range(cp_size):
rank_tensor: torch.Tensor = torch.cat(rank_segments[cp_rank], dim=0)
assert (
rank_tensor.shape[0] == total_per_rank
), f"rank {cp_rank}: expected {total_per_rank} tokens, got {rank_tensor.shape[0]}"
assert rank_tensor.shape[0] == total_per_rank, (
f"rank {cp_rank}: expected {total_per_rank} tokens, got {rank_tensor.shape[0]}"
)
_create_rank_dump(
directory,
@@ -4008,11 +4008,13 @@ class TestEntrypointMetaOverride:
baseline_path, target_path = self._create_single_rank_pair(tmp_path)
yaml_path: Path = tmp_path / "override.yaml"
yaml_path.write_text(textwrap.dedent("""\
yaml_path.write_text(
textwrap.dedent("""\
overrides:
- match: "hidden"
dims: "t h"
"""))
""")
)
argv = _make_argv(
baseline_path,
@@ -4148,13 +4150,15 @@ class TestEntrypointMetaOverride:
baseline_path, target_path = self._create_single_rank_pair(tmp_path)
yaml_path: Path = tmp_path / "override.yaml"
yaml_path.write_text(textwrap.dedent("""\
yaml_path.write_text(
textwrap.dedent("""\
overrides:
- match: "hidden"
dims: "t h"
- match: "hidden"
dims: "a b"
"""))
""")
)
argv = _make_argv(
baseline_path,
@@ -4169,11 +4173,13 @@ class TestEntrypointMetaOverride:
baseline_path, target_path = self._create_single_rank_pair(tmp_path)
yaml_path: Path = tmp_path / "override.yaml"
yaml_path.write_text(textwrap.dedent("""\
yaml_path.write_text(
textwrap.dedent("""\
overrides:
- match: "hidden"
dims: "a b"
"""))
""")
)
argv = _make_argv(
baseline_path,
@@ -194,11 +194,13 @@ class TestFromArgsAndConfig:
def test_cli_before_yaml(self, tmp_path: Path) -> None:
"""CLI rules are ordered before YAML rules (CLI wins on conflict)."""
yaml_path = tmp_path / "override.yaml"
yaml_path.write_text(textwrap.dedent("""\
yaml_path.write_text(
textwrap.dedent("""\
overrides:
- match: "hidden"
dims: "FROM_YAML"
"""))
""")
)
overrider = MetaOverrider.from_args_and_config(
override_dims=["hidden:FROM_CLI"],
@@ -256,14 +258,16 @@ class TestLoadYamlRules:
def test_valid_yaml(self, tmp_path: Path) -> None:
"""Valid YAML with override rules loads correctly."""
yaml_path = tmp_path / "override.yaml"
yaml_path.write_text(textwrap.dedent("""\
yaml_path.write_text(
textwrap.dedent("""\
overrides:
- match: "hidden"
dims: "b s h d"
- match: "logits"
dims: "b s v[tp]"
side: baseline
"""))
""")
)
rules = _load_yaml_rules(yaml_path)
assert len(rules) == 2
assert rules[0].dims == "b s h d"
@@ -591,9 +591,7 @@ class TestFormatAlignerPlan:
)
result: str = _format_aligner_plan(_wrap_plan(plan))
assert result == (
"Aligner Plan:\n" " baseline: (no steps)\n" " target: (no steps)"
)
assert result == ("Aligner Plan:\n baseline: (no steps)\n target: (no steps)")
def test_unsharder(self) -> None:
unsharder: UnsharderPlan = UnsharderPlan(
@@ -614,9 +612,7 @@ class TestFormatAlignerPlan:
result: str = _format_aligner_plan(_wrap_plan(plan))
assert result == (
"Aligner Plan:\n"
" baseline: (no steps)\n"
" target: [step=0: unsharder(tp)]"
"Aligner Plan:\n baseline: (no steps)\n target: [step=0: unsharder(tp)]"
)
def test_reorderer(self) -> None:
@@ -12,7 +12,7 @@ class TestApplyEdits:
"""Tests for the apply_edits() source text transformation function."""
def test_single_line_match_to_multiline_replacement(self) -> None:
source = "def foo():\n" " x = compute()\n" " return x\n"
source = "def foo():\n x = compute()\n return x\n"
edits = [
EditSpec(
match="x = compute()",
@@ -20,12 +20,10 @@ class TestApplyEdits:
)
]
result = apply_edits(source=source, edits=edits)
assert result == (
"def foo():\n" " x = compute()\n" " print(x)\n" " return x\n"
)
assert result == ("def foo():\n x = compute()\n print(x)\n return x\n")
def test_pure_insertion(self) -> None:
source = "def foo():\n" " a = 1\n" " b = 2\n"
source = "def foo():\n a = 1\n b = 2\n"
edits = [
EditSpec(
match="a = 1",
@@ -33,10 +31,10 @@ class TestApplyEdits:
)
]
result = apply_edits(source=source, edits=edits)
assert result == ("def foo():\n" " a = 1\n" " print(a)\n" " b = 2\n")
assert result == ("def foo():\n a = 1\n print(a)\n b = 2\n")
def test_pure_deletion_via_empty_replacement(self) -> None:
source = "def foo():\n" " debug_log()\n" " return 42\n"
source = "def foo():\n debug_log()\n return 42\n"
edits = [
EditSpec(
match="debug_log()",
@@ -44,10 +42,10 @@ class TestApplyEdits:
)
]
result = apply_edits(source=source, edits=edits)
assert result == ("def foo():\n" " return 42\n")
assert result == ("def foo():\n return 42\n")
def test_deletion_fewer_lines(self) -> None:
source = "def foo():\n" " a = 1\n" " b = 2\n" " c = 3\n"
source = "def foo():\n a = 1\n b = 2\n c = 3\n"
edits = [
EditSpec(
match="a = 1\nb = 2",
@@ -55,7 +53,7 @@ class TestApplyEdits:
)
]
result = apply_edits(source=source, edits=edits)
assert result == ("def foo():\n" " ab = 3\n" " c = 3\n")
assert result == ("def foo():\n ab = 3\n c = 3\n")
def test_multiline_match_to_multiline_replacement(self) -> None:
source = (
@@ -133,15 +131,7 @@ class TestApplyEdits:
def test_not_found_diagnostic_single_window_with_marker(self) -> None:
"""first line is present once but full match doesn't fit: one window with '>' on the match-region line."""
source = (
"line0\n"
"line1\n"
"line2\n"
"anchor()\n"
"wrong_next()\n"
"line5\n"
"line6\n"
)
source = "line0\nline1\nline2\nanchor()\nwrong_next()\nline5\nline6\n"
edits = [EditSpec(match="anchor()\nright_next()", replacement="x")]
with pytest.raises(PatchApplicationError) as exc_info:
apply_edits(source=source, edits=edits)
@@ -232,24 +222,22 @@ class TestApplyEdits:
assert "filler9" not in msg
def test_match_found_multiple_times_raises(self) -> None:
source = "def foo():\n" " print(1)\n" " print(1)\n"
source = "def foo():\n print(1)\n print(1)\n"
edits = [EditSpec(match="print(1)", replacement="print(2)")]
with pytest.raises(PatchApplicationError, match="multiple"):
apply_edits(source=source, edits=edits)
def test_multiple_edits_applied_sequentially(self) -> None:
source = "def foo():\n" " a = 1\n" " b = 2\n" " return a + b\n"
source = "def foo():\n a = 1\n b = 2\n return a + b\n"
edits = [
EditSpec(match="a = 1", replacement="a = 10"),
EditSpec(match="b = 2", replacement="b = 20"),
]
result = apply_edits(source=source, edits=edits)
assert result == (
"def foo():\n" " a = 10\n" " b = 20\n" " return a + b\n"
)
assert result == ("def foo():\n a = 10\n b = 20\n return a + b\n")
def test_strip_matching_ignores_leading_trailing_whitespace(self) -> None:
source = "def foo():\n" " x = compute()\n" " return x\n"
source = "def foo():\n x = compute()\n return x\n"
edits = [
EditSpec(
match=" x = compute() ",
@@ -257,11 +245,11 @@ class TestApplyEdits:
)
]
result = apply_edits(source=source, edits=edits)
assert result == ("def foo():\n" " x = replaced()\n" " return x\n")
assert result == ("def foo():\n x = replaced()\n return x\n")
def test_replacement_indented_text_realigned(self) -> None:
"""replacement text with its own indentation gets realigned to match source."""
source = "def foo():\n" " x = compute()\n" " return x\n"
source = "def foo():\n x = compute()\n return x\n"
edits = [
EditSpec(
match="x = compute()",
@@ -270,15 +258,12 @@ class TestApplyEdits:
]
result = apply_edits(source=source, edits=edits)
assert result == (
"def foo():\n"
" x = compute()\n"
" print(x)\n"
" return x\n"
"def foo():\n x = compute()\n print(x)\n return x\n"
)
def test_replacement_with_existing_indent_realigned(self) -> None:
"""replacement text already has indentation that should be rebased."""
source = "def foo():\n" " if True:\n" " x = 1\n" " return x\n"
source = "def foo():\n if True:\n x = 1\n return x\n"
edits = [
EditSpec(
match="x = 1",
@@ -296,12 +281,10 @@ class TestApplyEdits:
)
def test_append_keeps_match_and_adds_after(self) -> None:
source = "def foo():\n" " x = compute()\n" " return x\n"
source = "def foo():\n x = compute()\n return x\n"
edits = [EditSpec(match="x = compute()", append="print(x)")]
result = apply_edits(source=source, edits=edits)
assert result == (
"def foo():\n" " x = compute()\n" " print(x)\n" " return x\n"
)
assert result == ("def foo():\n x = compute()\n print(x)\n return x\n")
def test_append_multiline_match(self) -> None:
source = (
@@ -330,21 +313,18 @@ class TestApplyEdits:
)
def test_prepend_adds_before_match(self) -> None:
source = "def foo():\n" " x = compute()\n" " return x\n"
source = "def foo():\n x = compute()\n return x\n"
edits = [EditSpec(match="x = compute()", prepend="print('before')")]
result = apply_edits(source=source, edits=edits)
assert result == (
"def foo():\n"
" print('before')\n"
" x = compute()\n"
" return x\n"
"def foo():\n print('before')\n x = compute()\n return x\n"
)
def test_prepend_multiline(self) -> None:
source = "def foo():\n" " return x\n"
source = "def foo():\n return x\n"
edits = [EditSpec(match="return x", prepend="a = 1\nb = 2")]
result = apply_edits(source=source, edits=edits)
assert result == ("def foo():\n" " a = 1\n" " b = 2\n" " return x\n")
assert result == ("def foo():\n a = 1\n b = 2\n return x\n")
def test_prepend_deep_indent(self) -> None:
source = (
@@ -365,11 +345,7 @@ class TestApplyEdits:
def test_prepend_multiline_match(self) -> None:
source = (
"def foo():\n"
" result = call(\n"
" a=1,\n"
" )\n"
" return result\n"
"def foo():\n result = call(\n a=1,\n )\n return result\n"
)
edits = [
EditSpec(
@@ -401,13 +377,13 @@ class TestApplyEdits:
def test_second_edit_sees_result_of_first(self) -> None:
"""Edits are applied sequentially; second edit matches modified source."""
source = "def foo():\n" " x = 1\n" " return x\n"
source = "def foo():\n x = 1\n return x\n"
edits = [
EditSpec(match="x = 1", replacement="x = 1\ny = 2"),
EditSpec(match="y = 2", replacement="y = 20"),
]
result = apply_edits(source=source, edits=edits)
assert result == ("def foo():\n" " x = 1\n" " y = 20\n" " return x\n")
assert result == ("def foo():\n x = 1\n y = 20\n return x\n")
if __name__ == "__main__":
+52 -49
View File
@@ -734,9 +734,9 @@ def _assert_files(filenames, *, exist=(), not_exist=()):
for p in exist:
assert any(p in f for f in filenames), f"{p} not found in {filenames}"
for p in not_exist:
assert not any(
p in f for f in filenames
), f"{p} should not exist in {filenames}"
assert not any(p in f for f in filenames), (
f"{p} should not exist in {filenames}"
)
def _load_dump(path: Path) -> dict:
@@ -750,9 +750,9 @@ def _find_dump_file(tmpdir, *, rank: int = 0, name: str) -> Path:
for f in Path(tmpdir).glob("*/*.pt")
if f"rank={rank}" in f.name and name in f.name
]
assert (
len(matches) == 1
), f"Expected 1 file matching rank={rank} name={name}, got {matches}"
assert len(matches) == 1, (
f"Expected 1 file matching rank={rank} name={name}, got {matches}"
)
return matches[0]
@@ -1657,9 +1657,9 @@ class TestZmqPortIsolation:
)
resp.raise_for_status()
states = resp.json()
assert (
len(states) == 2
), f"Instance {i} (port {port}): expected 2 ranks, got {len(states)}"
assert len(states) == 2, (
f"Instance {i} (port {port}): expected 2 ranks, got {len(states)}"
)
finally:
for event in stop_events:
event.set()
@@ -1719,9 +1719,9 @@ class TestDumperHttp:
val = state
for k in keys:
val = val[k]
assert (
val == expected
), f"rank {rank}: {path}={val!r}, expected {expected!r}"
assert val == expected, (
f"rank {rank}: {path}={val!r}, expected {expected!r}"
)
def test_configure_enable_toggle(self, dumper_http_url: str):
for enable in [True, False]:
@@ -1915,9 +1915,9 @@ class TestNonIntrusiveDumper(_NonIntrusiveTestBase):
)
dumped_output = captured[f"{P}model.mutator.output"]["value"]
assert (
dumped_output == 999.0
).all(), "post-hook should capture outputs after forward"
assert (dumped_output == 999.0).all(), (
"post-hook should capture outputs after forward"
)
def test_hooks_all_module_levels(self, tmp_path):
class Attention(torch.nn.Module):
@@ -2374,9 +2374,9 @@ class TestDumperE2E:
states = requests.post(f"{base_url}/dumper/get_state", json={}).json()
assert len(states) == 2
for rank, state in enumerate(states):
assert (
state["config"]["enable"] is True
), f"rank {rank}: enable should be True after configure"
assert state["config"]["enable"] is True, (
f"rank {rank}: enable should be True after configure"
)
assert state["config"]["dir"] == dump_dir
resp = requests.post(
@@ -2403,16 +2403,16 @@ class TestDumperE2E:
)
for rank in range(2):
assert any(
f"rank={rank}" in f for f in filenames
), f"No dump files for rank {rank}"
assert any(f"rank={rank}" in f for f in filenames), (
f"No dump files for rank {rank}"
)
sample_file = dump_files[0]
loaded = torch.load(sample_file, map_location="cpu", weights_only=False)
assert isinstance(loaded, dict), f"Expected dict, got {type(loaded)}"
assert (
"value" in loaded and "meta" in loaded
), f"Missing value/meta keys: {loaded.keys()}"
assert "value" in loaded and "meta" in loaded, (
f"Missing value/meta keys: {loaded.keys()}"
)
assert "name" in loaded["meta"]
assert "rank" in loaded["meta"]
assert "step" in loaded["meta"]
@@ -2438,22 +2438,22 @@ class TestDumperE2E:
"attn_cp_size",
]
for key in expected_keys:
assert (
key in par
), f"Missing {key} in sglang_parallel_info, got: {sorted(par)}"
assert key in par, (
f"Missing {key} in sglang_parallel_info, got: {sorted(par)}"
)
rids_files = [f for f in dump_files if "name=rids" in f.name]
rids_loaded = torch.load(
rids_files[0], map_location="cpu", weights_only=False
)
rids_value = rids_loaded["value"]
assert isinstance(
rids_value, list
), f"rids should be a list, got {type(rids_value)}"
assert isinstance(rids_value, list), (
f"rids should be a list, got {type(rids_value)}"
)
assert len(rids_value) > 0, "rids should be non-empty"
assert all(
isinstance(r, str) for r in rids_value
), f"each rid should be a str, got {[type(r) for r in rids_value]}"
assert all(isinstance(r, str) for r in rids_value), (
f"each rid should be a str, got {[type(r) for r in rids_value]}"
)
finally:
kill_process_tree(proc.pid)
@@ -2912,9 +2912,9 @@ class TestRecomputeStatus:
model(torch.randn(2, 4))
for key, data in captured.items():
assert (
"recompute_status" in data["meta"]
), f"missing recompute_status in {key}"
assert "recompute_status" in data["meta"], (
f"missing recompute_status in {key}"
)
assert data["meta"]["recompute_status"] == "disabled"
def test_detect_recompute_status_default(self) -> None:
@@ -3553,8 +3553,7 @@ class TestGrafterDistributed:
# worker prepends tmp_path to sys.path so import_module sees it.
module_name = "_xform_user_basic"
(tmp_path / f"{module_name}.py").write_text(
"def transform(graft_input):\n"
" return graft_input.received_list[0] * 2\n"
"def transform(graft_input):\n return graft_input.received_list[0] * 2\n"
)
graft_port = find_available_port(29610)
_run_graft_test(
@@ -3646,7 +3645,9 @@ class TestGrafterDistributed:
7.0,
7.0,
7.0,
], f"target should be unchanged after shape-mismatch graft, got {target.tolist()}"
], (
f"target should be unchanged after shape-mismatch graft, got {target.tolist()}"
)
finally:
if grafter._pg is not None:
dist.destroy_process_group(grafter._pg)
@@ -3694,7 +3695,9 @@ class TestGrafterDistributed:
9.0,
9.0,
9.0,
], f"target must be unchanged when transform throws, got {target.tolist()}"
], (
f"target must be unchanged when transform throws, got {target.tolist()}"
)
output = captured.getvalue()
assert "transform/copy_ raised RuntimeError" in output, output
assert "intentional test error" in output, output
@@ -3783,9 +3786,9 @@ class TestGrafterDistributed:
grafter.maybe_intercept(value=target, tags={"name": "x"})
output = captured.getvalue()
if rank == 0:
assert (
"WARNING" in output
), f"expected WARNING in rank 0 output: {output}"
assert "WARNING" in output, (
f"expected WARNING in rank 0 output: {output}"
)
assert "has not completed after 2s" in output, output
finally:
if grafter._pg is not None:
@@ -3852,9 +3855,9 @@ class TestGrafterDistributed:
pg_after_first = grafter._pg
assert pg_after_first is not None
grafter.maybe_intercept(value=t2, tags={"name": "x"})
assert (
grafter._pg is pg_after_first
), "_pg must be cached across calls, not re-initialized"
assert grafter._pg is pg_after_first, (
"_pg must be cached across calls, not re-initialized"
)
else:
target1 = torch.zeros(3, device="cuda:1")
target2 = torch.zeros(3, device="cuda:1")
@@ -4190,9 +4193,9 @@ def _e2e_transform(graft_input):
the transform is just identity. Real workflows would compute a
non-trivial override (scale, reshape, decode, ...) using the extras.
"""
assert (
graft_input.received_extras_list[0]["my_extra_key"] == "my_extra_value"
), graft_input.received_extras_list
assert graft_input.received_extras_list[0]["my_extra_key"] == "my_extra_value", (
graft_input.received_extras_list
)
return graft_input.received_list[0]
@@ -290,7 +290,7 @@ def _run_e2e_scenario(
print(f"Comparator debug output: {debug_file}")
assert result.returncode == 0, (
f"Comparator failed (rc={result.returncode}). " f"Debug output: {debug_file}"
f"Comparator failed (rc={result.returncode}). Debug output: {debug_file}"
)
@@ -31,7 +31,6 @@ TEST_HIDDEN_SIZE = 32
class SimpleModel(nn.Module):
def __init__(self) -> None:
super().__init__()
self.hidden_size = TEST_HIDDEN_SIZE