"""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 "" in content assert "