feat(deps-viewer): mark alloc tasks as fire sources

simpler_setup/tools/deps_viewer.py

@@ -750,27 +750,30 @@ 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 +810,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};",

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():