[CI][RFC] Replace black-jupyter with ruff-format (#37210)
Co-authored-by: Alison Shao <a.shao@wustl.edu>
This commit is contained in:
co-authored by
Alison Shao
parent
2641e427be
commit
28262c20df
@@ -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__":
|
||||
|
||||
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user