From 1e1e0dd68809315389e1a70503a93b444f5d8922 Mon Sep 17 00:00:00 2001 From: Youhezhen Date: Mon, 13 Jul 2026 02:22:49 -0700 Subject: [PATCH 1/6] feat(deps-viewer): mark alloc tasks as fire sources --- simpler_setup/tools/deps_viewer.py | 23 ++++++++++++++--------- tests/ut/py/test_deps_viewer.py | 15 ++++++++++++++- 2 files changed, 28 insertions(+), 10 deletions(-) diff --git a/simpler_setup/tools/deps_viewer.py b/simpler_setup/tools/deps_viewer.py index e56330613c..9d28bd196a 100644 --- a/simpler_setup/tools/deps_viewer.py +++ b/simpler_setup/tools/deps_viewer.py @@ -750,27 +750,32 @@ def emit_text(edges, nodes, meta, deps_path, annotations=None, tensor_table=None return "\n".join(lines) + "\n" -def _task_markers(nodes, edges, task_table): +def _task_markers(nodes, edges, meta, task_table): """Map task_id -> marker string for the node label. - 🔥 (fire): the task itself is a flagged early-dispatch producer - (deps.json ``early_dispatch`` — the submit had allow_early_resolve). + 🔥 (fire): the task is either a flagged early-dispatch producer + (deps.json ``early_dispatch`` — the submit had allow_early_resolve) + or an alloc task. Alloc tasks are immediate graph sources, so they + are treated as fire-marked producers by default. ⭐ (star): every one of the task's predecessors is 🔥 (and it has at - least one), so the task is fully fed by flagged producers. + least one), so the task is fully fed by flagged producers or alloc + sources. A task can carry both. """ pred_map: dict[int, set] = {} for pred, succ in edges: pred_map.setdefault(succ, set()).add(pred) - def _flagged(tid): - return bool((task_table.get(tid) or {}).get("early_dispatch")) + def _fire_marked(tid): + return _task_kind(tid, meta, task_table) == "alloc" or bool( + (task_table.get(tid) or {}).get("early_dispatch") + ) markers = {} for tid in nodes: - fire = "🔥" if _flagged(tid) else "" + fire = "🔥" if _fire_marked(tid) else "" preds = pred_map.get(tid, set()) - star = "⭐" if preds and all(_flagged(p) for p in preds) else "" + star = "⭐" if preds and all(_fire_marked(p) for p in preds) else "" if fire or star: markers[tid] = fire + star return markers @@ -807,7 +812,7 @@ def emit_dot( task_table = task_table or {} hidden_edges = set(hidden_edges or ()) show_tensor = bool(task_table) if show_tensor_info is None else bool(show_tensor_info and task_table) - markers = _task_markers(nodes, edges, task_table) + markers = _task_markers(nodes, edges, meta, task_table) lines = [ "digraph deps {", f" rankdir={direction};", diff --git a/tests/ut/py/test_deps_viewer.py b/tests/ut/py/test_deps_viewer.py index cfd10eeaab..2f533ee614 100644 --- a/tests/ut/py/test_deps_viewer.py +++ b/tests/ut/py/test_deps_viewer.py @@ -240,7 +240,20 @@ def test_spmd_badges_json_includes_only_multiblock_tasks(): def test_emit_dot_handles_missing_task_table(): dot = deps_viewer.emit_dot(edges=[], nodes=[1], meta={}, task_table=None) - assert 'label="1 · alloc"' in dot + assert 'label="🔥 1 · alloc"' in dot + + +def test_emit_dot_treats_alloc_as_a_fire_marked_source(): + dot = deps_viewer.emit_dot( + edges=[(1, 2)], + nodes=[1, 2], + meta={}, + task_table={2: {"task_id": 2, "kernel_ids": [-1, 7, -1]}}, + show_tensor_info=False, + ) + + assert 'label="🔥 1 · alloc"' in dot + assert 'label="⭐ 2"' in dot def test_emit_dot_hides_selected_edges_with_background_color(): From 8714ac9756d20a3f0dca58317f0c6c5c4f3b2015 Mon Sep 17 00:00:00 2001 From: Youhezhen Date: Mon, 13 Jul 2026 03:47:54 -0700 Subject: [PATCH 2/6] Fix: format deps viewer markers --- simpler_setup/tools/deps_viewer.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/simpler_setup/tools/deps_viewer.py b/simpler_setup/tools/deps_viewer.py index 9d28bd196a..bdf1597fb3 100644 --- a/simpler_setup/tools/deps_viewer.py +++ b/simpler_setup/tools/deps_viewer.py @@ -767,9 +767,7 @@ def _task_markers(nodes, edges, meta, task_table): pred_map.setdefault(succ, set()).add(pred) def _fire_marked(tid): - return _task_kind(tid, meta, task_table) == "alloc" or bool( - (task_table.get(tid) or {}).get("early_dispatch") - ) + return _task_kind(tid, meta, task_table) == "alloc" or bool((task_table.get(tid) or {}).get("early_dispatch")) markers = {} for tid in nodes: From 6aaa7e79b661b55f06c899b9feba3eef22cd79ef Mon Sep 17 00:00:00 2001 From: Youhezhen Date: Tue, 14 Jul 2026 18:42:33 -0700 Subject: [PATCH 3/6] chore: retrigger CI From 4f63aba77dfc058dbe427c28645401bf3eaed011 Mon Sep 17 00:00:00 2001 From: Youhezhen Date: Wed, 15 Jul 2026 01:14:08 -0700 Subject: [PATCH 4/6] chore: retrigger CI From 72205ddd253afc567e04fd5616449d1571c5b484 Mon Sep 17 00:00:00 2001 From: Youhezhen Date: Wed, 15 Jul 2026 01:41:26 -0700 Subject: [PATCH 5/6] chore: retrigger CI From 479adde34bc3effc40c52a29da6392132b20e056 Mon Sep 17 00:00:00 2001 From: Youhezhen Date: Wed, 15 Jul 2026 02:01:31 -0700 Subject: [PATCH 6/6] chore: retrigger CI