266 lines
8.3 KiB
Python
266 lines
8.3 KiB
Python
|
"""Tests for the tracing + visualization pipeline.
|
|||
|
|
|
|||
|
|
Runs a real training loop with MockTransport + Tracer, then verifies
|
|||
|
|
all three output modes produce correct data.
|
|||
|
|
"""
|
|||
|
|
|
|||
|
|
from __future__ import annotations
|
|||
|
|
|
|||
|
|
import io
|
|||
|
|
import os
|
|||
|
|
|
|||
|
|
import pytest
|
|||
|
|
import torch
|
|||
|
|
|
|||
|
|
from sched.job import Job
|
|||
|
|
from sched.report import generate_html
|
|||
|
|
from sched.scheduler import Node, Scheduler
|
|||
|
|
from sched.trace import Tracer
|
|||
|
|
from tests.mock_transport import MockTransport
|
|||
|
|
|
|||
|
|
JOBS_DIR = os.path.join(os.path.dirname(__file__), os.pardir, "jobs")
|
|||
|
|
MNIST_SCRIPT = os.path.abspath(os.path.join(JOBS_DIR, "mnist.py"))
|
|||
|
|
|
|||
|
|
|
|||
|
|
async def _run_traced(
|
|||
|
|
rounds=3, local_steps=10, num_nodes=2, kill_node: int | None = None, kill_after_round: int | None = None,
|
|||
|
|
) -> tuple[Scheduler, Tracer, MockTransport]:
|
|||
|
|
"""Helper: run a traced job and return the scheduler + tracer."""
|
|||
|
|
buf = io.StringIO()
|
|||
|
|
tracer = Tracer(file=buf)
|
|||
|
|
job = Job(
|
|||
|
|
script=MNIST_SCRIPT,
|
|||
|
|
num_nodes=num_nodes,
|
|||
|
|
rounds=rounds,
|
|||
|
|
local_steps=local_steps,
|
|||
|
|
lr=0.01,
|
|||
|
|
batch_size=32,
|
|||
|
|
)
|
|||
|
|
transport = MockTransport()
|
|||
|
|
sched = Scheduler(job, transport, tracer=tracer)
|
|||
|
|
|
|||
|
|
for rank in range(num_nodes):
|
|||
|
|
conn = await transport.connect(f"mock{rank}", 22)
|
|||
|
|
sched.nodes.append(
|
|||
|
|
Node(rank=rank, host=f"mock{rank}", port=22, conn=conn, status="ready")
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
await sched.deploy()
|
|||
|
|
|
|||
|
|
job_mod = sched._load_job_module()
|
|||
|
|
state = job_mod.make_model().state_dict()
|
|||
|
|
|
|||
|
|
for r in range(rounds):
|
|||
|
|
if kill_node is not None and kill_after_round is not None and r == kill_after_round + 1:
|
|||
|
|
sched.nodes[kill_node].conn.kill()
|
|||
|
|
state = await sched.run_round(state, r)
|
|||
|
|
|
|||
|
|
return sched, tracer, transport
|
|||
|
|
|
|||
|
|
|
|||
|
|
# ------------------------------------------------------------------
|
|||
|
|
# C) text log tests
|
|||
|
|
# ------------------------------------------------------------------
|
|||
|
|
|
|||
|
|
|
|||
|
|
@pytest.mark.asyncio
|
|||
|
|
async def test_text_log_events_emitted():
|
|||
|
|
"""Tracer collects the right number and kinds of events."""
|
|||
|
|
sched, tracer, transport = await _run_traced(rounds=2, local_steps=5)
|
|||
|
|
try:
|
|||
|
|
kinds = [ev.kind for ev in tracer.events]
|
|||
|
|
assert kinds.count("round_start") == 2
|
|||
|
|
assert kinds.count("round_end") == 2
|
|||
|
|
assert kinds.count("aggregate") == 2
|
|||
|
|
# 2 workers × (push + exec + pull) × 2 rounds = 12
|
|||
|
|
# plus deploy pushes (4: worker.py + job_module.py per node)
|
|||
|
|
# plus deploy exec (2: "test -f" check per node)
|
|||
|
|
assert kinds.count("push") >= 4 # at least deploy pushes
|
|||
|
|
assert kinds.count("exec") == 6 # 2 deploy checks + 2 workers × 2 rounds
|
|||
|
|
assert kinds.count("pull") == 4
|
|||
|
|
finally:
|
|||
|
|
transport.cleanup()
|
|||
|
|
|
|||
|
|
|
|||
|
|
@pytest.mark.asyncio
|
|||
|
|
async def test_text_log_captures_errors():
|
|||
|
|
"""Dead worker produces events with error field."""
|
|||
|
|
sched, tracer, transport = await _run_traced(
|
|||
|
|
rounds=2, local_steps=5, kill_node=1, kill_after_round=0
|
|||
|
|
)
|
|||
|
|
try:
|
|||
|
|
error_events = [ev for ev in tracer.events if ev.data.get("error")]
|
|||
|
|
assert len(error_events) > 0
|
|||
|
|
assert any(ev.rank == 1 for ev in error_events)
|
|||
|
|
finally:
|
|||
|
|
transport.cleanup()
|
|||
|
|
|
|||
|
|
|
|||
|
|
@pytest.mark.asyncio
|
|||
|
|
async def test_text_log_output_to_buffer():
|
|||
|
|
"""Live text log is written to the provided file object."""
|
|||
|
|
buf = io.StringIO()
|
|||
|
|
tracer = Tracer(file=buf)
|
|||
|
|
tracer.emit("round_start", round_num=0, active_nodes=2)
|
|||
|
|
|
|||
|
|
output = buf.getvalue()
|
|||
|
|
assert "round 0 start" in output
|
|||
|
|
assert "2 workers" in output
|
|||
|
|
|
|||
|
|
|
|||
|
|
# ------------------------------------------------------------------
|
|||
|
|
# A) gantt tests
|
|||
|
|
# ------------------------------------------------------------------
|
|||
|
|
|
|||
|
|
|
|||
|
|
@pytest.mark.asyncio
|
|||
|
|
async def test_gantt_basic_structure():
|
|||
|
|
"""Gantt output contains expected sections."""
|
|||
|
|
sched, tracer, transport = await _run_traced(rounds=2, local_steps=5)
|
|||
|
|
try:
|
|||
|
|
text = tracer.gantt(width=40)
|
|||
|
|
assert "round 0" in text
|
|||
|
|
assert "round 1" in text
|
|||
|
|
assert "scheduler" in text
|
|||
|
|
assert "worker 0" in text
|
|||
|
|
assert "worker 1" in text
|
|||
|
|
# legend
|
|||
|
|
assert "push" in text
|
|||
|
|
assert "train" in text
|
|||
|
|
finally:
|
|||
|
|
transport.cleanup()
|
|||
|
|
|
|||
|
|
|
|||
|
|
@pytest.mark.asyncio
|
|||
|
|
async def test_gantt_shows_training_blocks():
|
|||
|
|
"""Worker rows contain training blocks (█)."""
|
|||
|
|
sched, tracer, transport = await _run_traced(rounds=1, local_steps=10)
|
|||
|
|
try:
|
|||
|
|
text = tracer.gantt(width=50)
|
|||
|
|
worker_lines = [l for l in text.splitlines() if "worker" in l and "█" in l]
|
|||
|
|
# both workers should have training blocks
|
|||
|
|
assert len(worker_lines) >= 2, f"Expected training blocks in:\n{text}"
|
|||
|
|
finally:
|
|||
|
|
transport.cleanup()
|
|||
|
|
|
|||
|
|
|
|||
|
|
@pytest.mark.asyncio
|
|||
|
|
async def test_gantt_shows_dead_worker():
|
|||
|
|
"""Dead worker shows ✗ markers."""
|
|||
|
|
sched, tracer, transport = await _run_traced(
|
|||
|
|
rounds=2, local_steps=5, kill_node=1, kill_after_round=0
|
|||
|
|
)
|
|||
|
|
try:
|
|||
|
|
text = tracer.gantt(width=40)
|
|||
|
|
# In round 1, worker 1 should have dead markers
|
|||
|
|
lines = text.splitlines()
|
|||
|
|
# find lines for round 1
|
|||
|
|
in_round_1 = False
|
|||
|
|
for line in lines:
|
|||
|
|
if "round 1" in line:
|
|||
|
|
in_round_1 = True
|
|||
|
|
if in_round_1 and "worker 1" in line:
|
|||
|
|
assert "✗" in line, f"Expected dead markers in: {line}"
|
|||
|
|
break
|
|||
|
|
finally:
|
|||
|
|
transport.cleanup()
|
|||
|
|
|
|||
|
|
|
|||
|
|
@pytest.mark.asyncio
|
|||
|
|
async def test_gantt_three_workers():
|
|||
|
|
"""Gantt handles 3 workers."""
|
|||
|
|
sched, tracer, transport = await _run_traced(
|
|||
|
|
rounds=1, local_steps=5, num_nodes=3
|
|||
|
|
)
|
|||
|
|
try:
|
|||
|
|
text = tracer.gantt(width=40)
|
|||
|
|
assert "worker 0" in text
|
|||
|
|
assert "worker 1" in text
|
|||
|
|
assert "worker 2" in text
|
|||
|
|
finally:
|
|||
|
|
transport.cleanup()
|
|||
|
|
|
|||
|
|
|
|||
|
|
# ------------------------------------------------------------------
|
|||
|
|
# B) HTML report tests
|
|||
|
|
# ------------------------------------------------------------------
|
|||
|
|
|
|||
|
|
|
|||
|
|
@pytest.mark.asyncio
|
|||
|
|
async def test_html_report_generated(tmp_path):
|
|||
|
|
"""HTML report is written and contains key elements."""
|
|||
|
|
sched, tracer, transport = await _run_traced(rounds=2, local_steps=5)
|
|||
|
|
try:
|
|||
|
|
path = str(tmp_path / "report.html")
|
|||
|
|
result = generate_html(tracer, path)
|
|||
|
|
assert os.path.exists(result)
|
|||
|
|
|
|||
|
|
content = open(result).read()
|
|||
|
|
assert "<!DOCTYPE html>" in content
|
|||
|
|
assert "<svg" in content
|
|||
|
|
assert "Training Run Report" in content
|
|||
|
|
assert "scheduler" in content
|
|||
|
|
assert "worker 0" in content
|
|||
|
|
# stats
|
|||
|
|
assert "Rounds" in content
|
|||
|
|
assert "Workers" in content
|
|||
|
|
assert "Data Transferred" in content
|
|||
|
|
finally:
|
|||
|
|
transport.cleanup()
|
|||
|
|
|
|||
|
|
|
|||
|
|
@pytest.mark.asyncio
|
|||
|
|
async def test_html_report_shows_dead_worker(tmp_path):
|
|||
|
|
"""HTML report marks dead workers in red."""
|
|||
|
|
sched, tracer, transport = await _run_traced(
|
|||
|
|
rounds=2, local_steps=5, kill_node=1, kill_after_round=0
|
|||
|
|
)
|
|||
|
|
try:
|
|||
|
|
path = str(tmp_path / "report.html")
|
|||
|
|
generate_html(tracer, path)
|
|||
|
|
content = open(path).read()
|
|||
|
|
# should contain error color and dead marker
|
|||
|
|
assert "#ef4444" in content # dead color
|
|||
|
|
assert "FAILED" in content or "err" in content
|
|||
|
|
finally:
|
|||
|
|
transport.cleanup()
|
|||
|
|
|
|||
|
|
|
|||
|
|
@pytest.mark.asyncio
|
|||
|
|
async def test_html_report_event_log(tmp_path):
|
|||
|
|
"""HTML report includes the text event log."""
|
|||
|
|
sched, tracer, transport = await _run_traced(rounds=1, local_steps=5)
|
|||
|
|
try:
|
|||
|
|
path = str(tmp_path / "report.html")
|
|||
|
|
generate_html(tracer, path)
|
|||
|
|
content = open(path).read()
|
|||
|
|
assert "event-log" in content
|
|||
|
|
assert "round 0" in content
|
|||
|
|
finally:
|
|||
|
|
transport.cleanup()
|
|||
|
|
|
|||
|
|
|
|||
|
|
# ------------------------------------------------------------------
|
|||
|
|
# summary table tests
|
|||
|
|
# ------------------------------------------------------------------
|
|||
|
|
|
|||
|
|
|
|||
|
|
@pytest.mark.asyncio
|
|||
|
|
async def test_summary_table():
|
|||
|
|
"""Summary table has correct structure."""
|
|||
|
|
sched, tracer, transport = await _run_traced(rounds=2, local_steps=5)
|
|||
|
|
try:
|
|||
|
|
text = tracer.summary()
|
|||
|
|
assert "Round" in text
|
|||
|
|
assert "Workers" in text
|
|||
|
|
assert "Δ norm" in text
|
|||
|
|
assert "total" in text
|
|||
|
|
# should have 2 data rows
|
|||
|
|
data_lines = [
|
|||
|
|
l for l in text.splitlines()
|
|||
|
|
if l.strip() and l.strip()[0].isdigit()
|
|||
|
|
]
|
|||
|
|
assert len(data_lines) == 2
|
|||
|
|
finally:
|
|||
|
|
transport.cleanup()
|