File size: 5,630 Bytes
656477e | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 | # SPDX-FileCopyrightText: © 2026 Tenstorrent AI ULC
# SPDX-License-Identifier: Apache-2.0
import ast
from pathlib import Path
from types import SimpleNamespace
import pytest
from models.common.tests.demos import cleanup_utils
def test_cleanup_dp_model_case_orders_owners_before_parent_and_children(monkeypatch):
calls = []
parent = SimpleNamespace(quiesce_devices=lambda: calls.append("parent-quiesce"))
submeshes = [object(), object()]
group = SimpleNamespace(cleanup=lambda: calls.append("group-cleanup"))
models = [("model-0", submeshes[0]), ("model-1", submeshes[1])]
monkeypatch.setattr(
cleanup_utils,
"cleanup_model_case",
lambda model, submesh: calls.append(("model-cleanup", model, submesh)),
)
monkeypatch.setattr(
cleanup_utils.ttnn,
"close_mesh_device",
lambda submesh: calls.append(("child-close", submesh)),
)
cleanup_utils.cleanup_dp_model_case(group, [], models, parent, submeshes)
assert calls == [
"group-cleanup",
("model-cleanup", "model-0", submeshes[0]),
("model-cleanup", "model-1", submeshes[1]),
"parent-quiesce",
("child-close", submeshes[0]),
("child-close", submeshes[1]),
]
def test_cleanup_dp_model_case_closes_every_carved_child_after_partial_build(monkeypatch):
calls = []
submeshes = [object(), object(), object(), object()]
lanes = [SimpleNamespace(cleanup=lambda: calls.append("lane-cleanup"))]
parent = SimpleNamespace(quiesce_devices=lambda: calls.append("parent-quiesce"))
monkeypatch.setattr(
cleanup_utils,
"cleanup_model_case",
lambda model, submesh: calls.append(("model-cleanup", model, submesh)),
)
monkeypatch.setattr(
cleanup_utils.ttnn,
"close_mesh_device",
lambda submesh: calls.append(("child-close", submesh)),
)
cleanup_utils.cleanup_dp_model_case(None, lanes, [("model-0", submeshes[0])], parent, submeshes)
assert calls[:3] == ["lane-cleanup", ("model-cleanup", "model-0", submeshes[0]), "parent-quiesce"]
assert calls[3:] == [("child-close", submesh) for submesh in submeshes]
def test_cleanup_dp_model_case_attempts_all_children_after_cleanup_failure(monkeypatch, expect_error):
closed = []
submeshes = [object(), object(), object()]
def fail_group_cleanup():
raise RuntimeError("group cleanup failed")
monkeypatch.setattr(cleanup_utils, "cleanup_model_case", lambda *_: None)
monkeypatch.setattr(cleanup_utils.ttnn, "close_mesh_device", closed.append)
with expect_error(RuntimeError, "group cleanup failed"):
cleanup_utils.cleanup_dp_model_case(
SimpleNamespace(cleanup=fail_group_cleanup),
[],
[],
SimpleNamespace(quiesce_devices=lambda: None),
submeshes,
)
assert closed == submeshes
def test_two_sequential_dp_profiles_release_children_between_carves(monkeypatch):
class Child:
def __init__(self, generation):
self.generation = generation
self.closed = False
class Parent:
def __init__(self):
self.children = []
self.quiesce_count = 0
def quiesce_devices(self):
self.quiesce_count += 1
def carve(self, generation):
assert all(child.closed for child in self.children)
children = [Child(generation), Child(generation)]
self.children.extend(children)
return children
parent = Parent()
monkeypatch.setattr(cleanup_utils.ttnn, "close_mesh_device", lambda child: setattr(child, "closed", True))
for generation in ("performance", "accuracy"):
parent.quiesce_devices()
children = parent.carve(generation)
cleanup_utils.cleanup_dp_model_case(None, [], [], parent, children)
assert parent.quiesce_count == 4
assert all(child.closed for child in parent.children)
@pytest.mark.parametrize(
("demo_path", "create_name"),
[
("models/common/tests/demos/qwen2_7b/demo.py", "_create_dp_submeshes"),
("models/common/tests/demos/qwen25_7b/demo.py", "_create_dp_submeshes"),
("models/common/tests/demos/deepseek_r1_distill_qwen_14b/demo.py", "_create_dp_submeshes"),
("models/common/tests/demos/llama32_1b/demo.py", "create_dp_submeshes"),
("models/common/tests/demos/llama32_3b/demo.py", "create_dp_submeshes"),
],
)
def test_dp_demos_quiesce_before_carving_and_use_shared_teardown(demo_path, create_name):
tree = ast.parse(Path(demo_path).read_text(encoding="utf-8"), filename=demo_path)
function = next(node for node in tree.body if isinstance(node, ast.FunctionDef) and node.name == "_run_dp_smoke")
calls = [node for node in ast.walk(function) if isinstance(node, ast.Call)]
create_call = next(node for node in calls if isinstance(node.func, ast.Name) and node.func.id == create_name)
parent_quiesce = next(
node
for node in calls
if isinstance(node.func, ast.Attribute)
and isinstance(node.func.value, ast.Name)
and node.func.value.id == "mesh_device"
and node.func.attr == "quiesce_devices"
)
teardown_call = next(
node for node in calls if isinstance(node.func, ast.Name) and node.func.id == "cleanup_dp_model_case"
)
assert parent_quiesce.lineno < create_call.lineno < teardown_call.lineno
assert [ast.unparse(argument) for argument in teardown_call.args] == [
"group",
"lanes",
"models",
"mesh_device",
"submeshes",
]
|