7d11dfa688
Implementation of a mock scheduler for jobs/tasks that will evolve into a job manager for running training jobs over vast.ai instances.
265 lines
8.3 KiB
Python
265 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()
|