7d11dfa688
Implementation of a mock scheduler for jobs/tasks that will evolve into a job manager for running training jobs over vast.ai instances.
424 lines
15 KiB
Python
424 lines
15 KiB
Python
"""Tracing infrastructure for observing data flow through the system.
|
|
|
|
Three output modes, all fed from the same event list:
|
|
C) Live text log — Tracer._print() streams events as they happen
|
|
A) Terminal Gantt — Tracer.gantt() renders a Unicode timeline per round
|
|
B) HTML report — report.generate_html(tracer, path) writes a self-contained SVG timeline
|
|
|
|
Usage:
|
|
tracer = Tracer()
|
|
sched = Scheduler(job, transport, tracer=tracer)
|
|
await sched.run()
|
|
tracer.gantt()
|
|
tracer.summary()
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import os
|
|
import sys
|
|
import time
|
|
from dataclasses import dataclass, field
|
|
|
|
|
|
# ------------------------------------------------------------------
|
|
# Events
|
|
# ------------------------------------------------------------------
|
|
|
|
@dataclass
|
|
class Event:
|
|
time: float # seconds since tracer start
|
|
kind: str # push | pull | exec | round_start | round_end | aggregate
|
|
rank: int | None # None for scheduler-level events
|
|
data: dict = field(default_factory=dict)
|
|
|
|
|
|
def _fmt_bytes(n: int) -> str:
|
|
if n < 1024:
|
|
return f"{n}B"
|
|
if n < 1024 * 1024:
|
|
return f"{n / 1024:.1f}KB"
|
|
return f"{n / (1024 * 1024):.1f}MB"
|
|
|
|
|
|
# ------------------------------------------------------------------
|
|
# Tracer
|
|
# ------------------------------------------------------------------
|
|
|
|
class Tracer:
|
|
def __init__(self, file=None):
|
|
self._t0 = time.monotonic()
|
|
self.events: list[Event] = []
|
|
self._file = file or sys.stderr
|
|
|
|
def emit(self, kind: str, rank: int | None = None, **data) -> Event:
|
|
ev = Event(
|
|
time=time.monotonic() - self._t0,
|
|
kind=kind,
|
|
rank=rank,
|
|
data=data,
|
|
)
|
|
self.events.append(ev)
|
|
self._print(ev)
|
|
return ev
|
|
|
|
# ------------------------------------------------------------------
|
|
# C) Live text log
|
|
# ------------------------------------------------------------------
|
|
|
|
def _print(self, ev: Event):
|
|
t = f"[{ev.time:7.2f}s]"
|
|
node = f"node={ev.rank}" if ev.rank is not None else " "
|
|
err = ev.data.get("error")
|
|
|
|
match ev.kind:
|
|
case "push":
|
|
sz = _fmt_bytes(ev.data.get("size_bytes", 0))
|
|
dur = ev.data.get("duration_s", 0)
|
|
path = os.path.basename(ev.data.get("path", ""))
|
|
if err:
|
|
line = f"{t} push {node} {path:<20s} !! {err}"
|
|
else:
|
|
line = f"{t} push {node} {path:<20s} {sz:>8s} {dur:.3f}s"
|
|
case "pull":
|
|
sz = _fmt_bytes(ev.data.get("size_bytes", 0))
|
|
dur = ev.data.get("duration_s", 0)
|
|
path = os.path.basename(ev.data.get("path", ""))
|
|
if err:
|
|
line = f"{t} pull {node} {path:<20s} !! {err}"
|
|
else:
|
|
line = f"{t} pull {node} {path:<20s} {sz:>8s} {dur:.3f}s"
|
|
case "exec":
|
|
dur = ev.data.get("duration_s", 0)
|
|
exit_code = ev.data.get("exit_code", "?")
|
|
if err:
|
|
line = f"{t} exec {node} !! {err}"
|
|
else:
|
|
line = f"{t} exec {node} exit={exit_code:<3} {dur:.3f}s"
|
|
output = ev.data.get("output_tail", "")
|
|
if output:
|
|
line += f" | {output}"
|
|
case "round_start":
|
|
r = ev.data.get("round_num", "?")
|
|
n = ev.data.get("active_nodes", "?")
|
|
line = f"{t} {'─'*4} round {r} start ({n} workers) {'─'*20}"
|
|
case "round_end":
|
|
r = ev.data.get("round_num", "?")
|
|
surv = ev.data.get("survivors", "?")
|
|
total = ev.data.get("total_nodes", "?")
|
|
dur = ev.data.get("duration_s", 0)
|
|
line = f"{t} {'─'*4} round {r} end ({surv}/{total} survived, {dur:.2f}s) {'─'*12}"
|
|
case "aggregate":
|
|
n = ev.data.get("num_workers", "?")
|
|
wnorm = ev.data.get("weight_norm", 0)
|
|
dnorm = ev.data.get("delta_norm", 0)
|
|
dur = ev.data.get("duration_s", 0)
|
|
line = (
|
|
f"{t} agg {n} workers"
|
|
f" |W|={wnorm:.2f}"
|
|
f" Δ={dnorm:.4f}"
|
|
f" {dur:.3f}s"
|
|
)
|
|
case _:
|
|
line = f"{t} {ev.kind:6s} {node} {ev.data}"
|
|
|
|
print(line, file=self._file, flush=True)
|
|
|
|
# ------------------------------------------------------------------
|
|
# helpers
|
|
# ------------------------------------------------------------------
|
|
|
|
def _collect_rounds(self) -> dict[int, dict]:
|
|
"""Group events into per-round buckets with timing metadata."""
|
|
rounds: dict[int, dict] = {}
|
|
current_round: int | None = None
|
|
|
|
for ev in self.events:
|
|
if ev.kind == "round_start":
|
|
current_round = ev.data["round_num"]
|
|
rounds[current_round] = {
|
|
"start": ev.time,
|
|
"end": ev.time,
|
|
"duration": 0.0,
|
|
"num_workers": ev.data.get("active_nodes", 0),
|
|
"survivors": 0,
|
|
"total": 0,
|
|
"events": [],
|
|
"weight_norm": 0.0,
|
|
"delta_norm": 0.0,
|
|
}
|
|
elif ev.kind == "round_end":
|
|
r = ev.data["round_num"]
|
|
if r in rounds:
|
|
rounds[r]["end"] = ev.time
|
|
rounds[r]["duration"] = ev.data.get("duration_s", 0)
|
|
rounds[r]["survivors"] = ev.data.get("survivors", 0)
|
|
rounds[r]["total"] = ev.data.get("total_nodes", 0)
|
|
elif ev.kind == "aggregate":
|
|
if current_round is not None and current_round in rounds:
|
|
rounds[current_round]["weight_norm"] = ev.data.get("weight_norm", 0)
|
|
rounds[current_round]["delta_norm"] = ev.data.get("delta_norm", 0)
|
|
rounds[current_round]["events"].append(ev)
|
|
else:
|
|
if current_round is not None and current_round in rounds:
|
|
rounds[current_round]["events"].append(ev)
|
|
|
|
return rounds
|
|
|
|
def _max_rank(self) -> int:
|
|
ranks = [ev.rank for ev in self.events if ev.rank is not None]
|
|
return max(ranks) if ranks else 0
|
|
|
|
# ------------------------------------------------------------------
|
|
# A) Terminal Gantt chart
|
|
# ------------------------------------------------------------------
|
|
|
|
def gantt(self, width: int = 60) -> str:
|
|
"""Render a Unicode Gantt chart showing data flow per round."""
|
|
rounds = self._collect_rounds()
|
|
if not rounds:
|
|
return "(no rounds recorded)"
|
|
|
|
max_rank = self._max_rank()
|
|
|
|
lines: list[str] = []
|
|
lines.append("")
|
|
lines.append(" ▓ push █ train ▒ pull ● agg ✗ dead ░ wait")
|
|
lines.append("")
|
|
|
|
for rnum in sorted(rounds.keys()):
|
|
rd = rounds[rnum]
|
|
t0 = rd["start"]
|
|
dur = rd["duration"] or 0.001
|
|
surv = rd["survivors"]
|
|
total = rd["total"]
|
|
|
|
lines.append(f" round {rnum} [{surv}/{total}] {dur:.2f}s")
|
|
|
|
def to_col(t: float) -> int:
|
|
frac = (t - t0) / dur
|
|
return max(0, min(width - 1, int(frac * width)))
|
|
|
|
def fill(row: list[str], t_start: float, t_end: float, char: str):
|
|
s = to_col(t_start)
|
|
e = to_col(t_end)
|
|
if s == e:
|
|
e = min(s + 1, width - 1)
|
|
for i in range(s, e + 1):
|
|
if 0 <= i < width:
|
|
row[i] = char
|
|
|
|
# --- scheduler row ---
|
|
sched = list("░" * width)
|
|
for ev in rd["events"]:
|
|
d = ev.data.get("duration_s", 0)
|
|
if ev.kind == "push":
|
|
fill(sched, ev.time - d, ev.time, "▓")
|
|
elif ev.kind == "pull":
|
|
fill(sched, ev.time - d, ev.time, "▒")
|
|
elif ev.kind == "aggregate":
|
|
fill(sched, ev.time - d, ev.time, "●")
|
|
lines.append(f" scheduler {''.join(sched)}")
|
|
|
|
# --- worker rows ---
|
|
for rank in range(max_rank + 1):
|
|
row = list("─" * width)
|
|
rank_evs = [e for e in rd["events"] if e.rank == rank]
|
|
dead_col = None
|
|
|
|
for ev in rank_evs:
|
|
d = ev.data.get("duration_s", 0)
|
|
if ev.data.get("error"):
|
|
dead_col = to_col(ev.time)
|
|
fill(row, ev.time - d, ev.time, "✗")
|
|
elif ev.kind == "exec" and ev.data.get("exit_code", 0) != 0:
|
|
dead_col = to_col(ev.time)
|
|
fill(row, ev.time - d, ev.time, "✗")
|
|
elif ev.kind == "exec":
|
|
fill(row, ev.time - d, ev.time, "█")
|
|
elif ev.kind == "push":
|
|
fill(row, ev.time - d, ev.time, "▓")
|
|
elif ev.kind == "pull":
|
|
fill(row, ev.time - d, ev.time, "▒")
|
|
|
|
# fill dead from failure point onward
|
|
if dead_col is not None:
|
|
for i in range(dead_col, width):
|
|
if row[i] == "─":
|
|
row[i] = "✗"
|
|
|
|
# worker was expected but produced nothing
|
|
if not rank_evs and rank < rd.get("num_workers", 0):
|
|
row = list("✗" * width)
|
|
|
|
lines.append(f" worker {rank:<3} {''.join(row)}")
|
|
|
|
lines.append("")
|
|
|
|
text = "\n".join(lines)
|
|
print(text, file=self._file, flush=True)
|
|
return text
|
|
|
|
# ------------------------------------------------------------------
|
|
# Summary table
|
|
# ------------------------------------------------------------------
|
|
|
|
def summary(self) -> str:
|
|
"""Render a round-by-round summary table."""
|
|
rounds = self._collect_rounds()
|
|
if not rounds:
|
|
return "(no rounds recorded)"
|
|
|
|
lines: list[str] = []
|
|
lines.append("")
|
|
|
|
header = (
|
|
f"{'Round':<6} {'Workers':<9} {'Params↑':<10} {'Weights↓':<11}"
|
|
f" {'Train(max)':<11} {'Agg':<8} {'Total':<8} {'Δ norm'}"
|
|
)
|
|
lines.append(header)
|
|
lines.append("─" * len(header))
|
|
|
|
total_push = 0
|
|
total_pull = 0
|
|
total_time = 0.0
|
|
|
|
for rnum in sorted(rounds.keys()):
|
|
rd = rounds[rnum]
|
|
surv = rd["survivors"]
|
|
total = rd["total"]
|
|
|
|
push_bytes = sum(
|
|
e.data.get("size_bytes", 0)
|
|
for e in rd["events"]
|
|
if e.kind == "push" and not e.data.get("error")
|
|
)
|
|
pull_bytes = sum(
|
|
e.data.get("size_bytes", 0)
|
|
for e in rd["events"]
|
|
if e.kind == "pull" and not e.data.get("error")
|
|
)
|
|
exec_durs = [
|
|
e.data.get("duration_s", 0)
|
|
for e in rd["events"]
|
|
if e.kind == "exec" and not e.data.get("error")
|
|
]
|
|
agg_dur = sum(
|
|
e.data.get("duration_s", 0)
|
|
for e in rd["events"]
|
|
if e.kind == "aggregate"
|
|
)
|
|
train_max = max(exec_durs) if exec_durs else 0
|
|
dur = rd["duration"]
|
|
delta = rd["delta_norm"]
|
|
|
|
total_push += push_bytes
|
|
total_pull += pull_bytes
|
|
total_time += dur
|
|
|
|
lines.append(
|
|
f"{rnum:<6} {surv}/{total:<6} {_fmt_bytes(push_bytes):<10}"
|
|
f" {_fmt_bytes(pull_bytes):<11} {train_max:<11.2f}"
|
|
f" {agg_dur:<8.3f} {dur:<8.2f} {delta:.4f}"
|
|
)
|
|
|
|
lines.append("─" * len(header))
|
|
lines.append(
|
|
f"{'total':<6} {'':9} {_fmt_bytes(total_push):<10}"
|
|
f" {_fmt_bytes(total_pull):<11} {'':11}"
|
|
f" {'':8} {total_time:<8.2f}"
|
|
)
|
|
|
|
text = "\n".join(lines)
|
|
print(text, file=self._file, flush=True)
|
|
return text
|
|
|
|
|
|
# ------------------------------------------------------------------
|
|
# TracingConnection wrapper
|
|
# ------------------------------------------------------------------
|
|
|
|
class TracingConnection:
|
|
"""Wraps any Connection to auto-trace push/pull/exec with timing and sizes.
|
|
|
|
Emits events on both success and failure so the Gantt chart can show
|
|
where workers died.
|
|
"""
|
|
|
|
def __init__(self, inner, rank: int, tracer: Tracer):
|
|
self._inner = inner
|
|
self._rank = rank
|
|
self._tracer = tracer
|
|
|
|
# proxy unknown attributes to inner (e.g. MockConnection.kill, .root_dir)
|
|
def __getattr__(self, name):
|
|
return getattr(self._inner, name)
|
|
|
|
async def exec(self, cmd: str, timeout: float = 30.0) -> tuple[int, str]:
|
|
t0 = time.monotonic()
|
|
try:
|
|
ret, output = await self._inner.exec(cmd, timeout=timeout)
|
|
except Exception as exc:
|
|
dur = time.monotonic() - t0
|
|
self._tracer.emit(
|
|
"exec", rank=self._rank,
|
|
duration_s=dur, exit_code=-1, output_tail="",
|
|
error=str(exc),
|
|
)
|
|
raise
|
|
|
|
dur = time.monotonic() - t0
|
|
tail = ""
|
|
for line in reversed(output.strip().splitlines()):
|
|
if line.strip():
|
|
tail = line.strip()[:80]
|
|
break
|
|
|
|
self._tracer.emit(
|
|
"exec", rank=self._rank,
|
|
duration_s=dur, exit_code=ret, output_tail=tail,
|
|
)
|
|
return ret, output
|
|
|
|
async def push(self, local_path: str, remote_path: str) -> None:
|
|
size = os.path.getsize(local_path) if os.path.exists(local_path) else 0
|
|
t0 = time.monotonic()
|
|
try:
|
|
await self._inner.push(local_path, remote_path)
|
|
except Exception as exc:
|
|
dur = time.monotonic() - t0
|
|
self._tracer.emit(
|
|
"push", rank=self._rank,
|
|
path=remote_path, size_bytes=0, duration_s=dur,
|
|
error=str(exc),
|
|
)
|
|
raise
|
|
|
|
dur = time.monotonic() - t0
|
|
self._tracer.emit(
|
|
"push", rank=self._rank,
|
|
path=remote_path, size_bytes=size, duration_s=dur,
|
|
)
|
|
|
|
async def pull(self, remote_path: str, local_path: str) -> None:
|
|
t0 = time.monotonic()
|
|
try:
|
|
await self._inner.pull(remote_path, local_path)
|
|
except Exception as exc:
|
|
dur = time.monotonic() - t0
|
|
self._tracer.emit(
|
|
"pull", rank=self._rank,
|
|
path=remote_path, size_bytes=0, duration_s=dur,
|
|
error=str(exc),
|
|
)
|
|
raise
|
|
|
|
dur = time.monotonic() - t0
|
|
size = os.path.getsize(local_path) if os.path.exists(local_path) else 0
|
|
self._tracer.emit(
|
|
"pull", rank=self._rank,
|
|
path=remote_path, size_bytes=size, duration_s=dur,
|
|
)
|
|
|
|
async def close(self) -> None:
|
|
await self._inner.close()
|