feat: FunPack Timing also logs a sampler's step 1 / step 2 / later-step times

Co-Authored-By: Claude Sonnet 5.5 <noreply@anthropic.com>

core/node_timing.py

@@ -11,6 +11,7 @@ from . import log
 
 _runs = {}
 _END = ("execution_success", "execution_error", "execution_interrupted")
+STEPS = 3          # a node that reports at least this many steps gets a per-step line
 TAG = "_funpack_node_timing"
 
 
@@ -22,8 +23,14 @@ def observe(event, data, now=None):
         _runs[pid] = {"start": now, "node": "(before the first node)", "at": now, "spent": {}}
         return
     run = _runs.get(pid)
-    if run is None or (event != "executing" and event not in _END):
+    if run is None:
         return
+    if event == "progress":
+        run.setdefault("marks", []).append(now)
+        return
+    if event != "executing" and event not in _END:
+        return
+    _steps(run, now)
     run["spent"][run["node"]] = run["spent"].get(run["node"], 0.0) + now - run["at"]
     if event == "executing" and data.get("node") is not None:
         run["node"], run["at"] = str(data.get("display_node") or data["node"]), now
@@ -35,6 +42,18 @@ def observe(event, data, now=None):
         log.info("FunPack Timing", f"{now - run['start']:.1f}s in nodes: {shown or 'nothing over 0.3s'}")
 
 
+def _steps(run, now):
+    """The sampler's own step times, once its node is over: how long until step 1 was done (loading and
+    first-call setup included), how long step 2 took, and what the later steps average."""
+    marks, run["marks"] = run.get("marks") or [], []
+    if len(marks) < STEPS:
+        return
+    first, second = marks[0] - run["at"], marks[1] - marks[0]
+    rest = (marks[-1] - marks[1]) / (len(marks) - 2)
+    log.info("FunPack Timing", f"{run['node']}: step 1 done {first:.1f}s after the node started, step 2 took "
+             f"{second:.1f}s, the {len(marks) - 2} after that average {rest:.1f}s")
+
+
 def install(server):
     """Listen to the server's outgoing messages; once, and never in the way of sending them."""
     send = getattr(server, "send_sync", None)

core/tests/test_node_timing.py

@@ -20,3 +20,16 @@ def test_a_run_is_split_by_node_and_logged_once_at_the_end():
     assert line == ["70.0s in nodes: sampler 30.0s · te 20.0s · save 19.0s · (before the first node) 1.0s"]
     s.send_sync("status", {})
     assert seen == ["status"], "messages still go out"
+
+
+def test_a_sampler_reports_its_first_steps_apart_from_the_rest():
+    log._reset()
+    obs = node_timing.observe
+    obs("execution_start", {"prompt_id": "p"}, now=0.0)
+    obs("executing", {"prompt_id": "p", "node": "sampler"}, now=1.0)
+    for i, at in enumerate((31.0, 33.0, 35.0, 37.0)):
+        obs("progress", {"prompt_id": "p", "value": i + 1, "max": 4}, now=at)
+    obs("execution_success", {"prompt_id": "p"}, now=38.0)
+    lines = [r["message"] for r in log.history() if r["source"] == "FunPack Timing"]
+    assert lines[0] == "sampler: step 1 done 30.0s after the node started, step 2 took 2.0s, the 2 after that average 2.0s"
+    assert len(lines) == 2, "the per-node total still follows"