diff --git a/.dockerignore b/.dockerignore new file mode 100644 index 0000000..fbe08b2 --- /dev/null +++ b/.dockerignore @@ -0,0 +1,12 @@ +.git +.venv +.pytest_cache +__pycache__ +*.pyc +.python-version +uv.lock +report.html +dtrain +DESIGN.md +README.md +.claude diff --git a/.python-version b/.python-version new file mode 100644 index 0000000..24ee5b1 --- /dev/null +++ b/.python-version @@ -0,0 +1 @@ +3.13 diff --git a/docker/docker-compose.cpu.yml b/docker/docker-compose.cpu.yml new file mode 100644 index 0000000..d169f78 --- /dev/null +++ b/docker/docker-compose.cpu.yml @@ -0,0 +1,27 @@ +services: + scheduler: + build: + context: .. + dockerfile: docker/scheduler.Dockerfile + depends_on: + - worker-0 + - worker-1 + networks: + - fedavg + + worker-0: + build: + context: .. + dockerfile: docker/worker-cpu.Dockerfile + networks: + - fedavg + + worker-1: + build: + context: .. + dockerfile: docker/worker-cpu.Dockerfile + networks: + - fedavg + +networks: + fedavg: diff --git a/docker/docker-compose.yml b/docker/docker-compose.yml new file mode 100644 index 0000000..74978a3 --- /dev/null +++ b/docker/docker-compose.yml @@ -0,0 +1,27 @@ +services: + scheduler: + build: + context: .. + dockerfile: docker/scheduler.Dockerfile + depends_on: + - worker-0 + - worker-1 + networks: + - fedavg + + worker-0: + build: + context: .. + dockerfile: docker/worker.Dockerfile + networks: + - fedavg + + worker-1: + build: + context: .. + dockerfile: docker/worker.Dockerfile + networks: + - fedavg + +networks: + fedavg: diff --git a/docker/scheduler.Dockerfile b/docker/scheduler.Dockerfile new file mode 100644 index 0000000..6db536e --- /dev/null +++ b/docker/scheduler.Dockerfile @@ -0,0 +1,15 @@ +FROM python:3.13-slim + +RUN apt-get update && apt-get install -y openssh-client && rm -rf /var/lib/apt/lists/* + +WORKDIR /app + +COPY pyproject.toml ./ +COPY sched/ ./sched/ +COPY worker/ ./worker/ +COPY jobs/ ./jobs/ +COPY tests/ ./tests/ + +RUN pip install --no-cache-dir torch --index-url https://download.pytorch.org/whl/cpu + +CMD ["python", "-m", "sched", "run", "--script", "jobs/mnist.py", "--compose", "--trace"] diff --git a/docker/smoke-test.sh b/docker/smoke-test.sh new file mode 100755 index 0000000..5f4e242 --- /dev/null +++ b/docker/smoke-test.sh @@ -0,0 +1,18 @@ +#!/bin/bash +set -e + +ROUNDS="${1:-1}" +NODES="${2:-2}" + +cd "$(dirname "$0")" + +mkdir -p ../reports + +docker compose -f docker-compose.cpu.yml run \ + -v "$(pwd)/../reports:/app/reports" \ + scheduler \ + bash -c "python -m sched run --script jobs/mnist.py --compose --trace --rounds $ROUNDS --nodes $NODES --report reports/report.html 2>&1 | tee reports/log.txt" + +docker compose -f docker-compose.cpu.yml down + +echo "Done. See reports/report.html and reports/log.txt" diff --git a/docker/worker-cpu.Dockerfile b/docker/worker-cpu.Dockerfile new file mode 100644 index 0000000..f504d58 --- /dev/null +++ b/docker/worker-cpu.Dockerfile @@ -0,0 +1,22 @@ +FROM python:3.13-slim + +RUN apt-get update && apt-get install -y openssh-server && rm -rf /var/lib/apt/lists/* +RUN mkdir -p /run/sshd + +# Allow root login with no password +RUN sed -i 's/#PermitRootLogin.*/PermitRootLogin yes/' /etc/ssh/sshd_config \ + && sed -i 's/#PermitEmptyPasswords.*/PermitEmptyPasswords yes/' /etc/ssh/sshd_config \ + && passwd -d root + +RUN pip install --no-cache-dir torch --index-url https://download.pytorch.org/whl/cpu + +WORKDIR /workspace + +COPY worker/worker.py /workspace/worker.py +COPY jobs/ /workspace/jobs/ + +COPY docker/worker-entrypoint.sh /entrypoint.sh +RUN chmod +x /entrypoint.sh + +EXPOSE 22 +CMD ["/entrypoint.sh"] diff --git a/docker/worker-entrypoint.sh b/docker/worker-entrypoint.sh new file mode 100644 index 0000000..98c1182 --- /dev/null +++ b/docker/worker-entrypoint.sh @@ -0,0 +1,8 @@ +#!/bin/bash +set -e + +# Start SSH daemon +/usr/sbin/sshd + +# Keep container alive +exec tail -f /dev/null diff --git a/docker/worker.Dockerfile b/docker/worker.Dockerfile new file mode 100644 index 0000000..08966c4 --- /dev/null +++ b/docker/worker.Dockerfile @@ -0,0 +1,20 @@ +FROM pytorch/pytorch:2.5.0-cuda12.4-cudnn9-runtime + +RUN apt-get update && apt-get install -y openssh-server && rm -rf /var/lib/apt/lists/* +RUN mkdir -p /run/sshd + +# Allow root login with no password (vast.ai injects keys at runtime) +RUN sed -i 's/#PermitRootLogin.*/PermitRootLogin yes/' /etc/ssh/sshd_config \ + && sed -i 's/#PermitEmptyPasswords.*/PermitEmptyPasswords yes/' /etc/ssh/sshd_config \ + && passwd -d root + +WORKDIR /workspace + +COPY worker/worker.py /workspace/worker.py +COPY jobs/ /workspace/jobs/ + +COPY docker/worker-entrypoint.sh /entrypoint.sh +RUN chmod +x /entrypoint.sh + +EXPOSE 22 +CMD ["/entrypoint.sh"] diff --git a/jobs/__init__.py b/jobs/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/jobs/__pycache__/mnist.cpython-313.pyc b/jobs/__pycache__/mnist.cpython-313.pyc new file mode 100644 index 0000000..fc41ac8 Binary files /dev/null and b/jobs/__pycache__/mnist.cpython-313.pyc differ diff --git a/jobs/mnist.py b/jobs/mnist.py new file mode 100644 index 0000000..7842e83 --- /dev/null +++ b/jobs/mnist.py @@ -0,0 +1,54 @@ +"""MNIST job module. + +Exports make_model() and make_dataloader() for the scheduler + worker. +Falls back to synthetic data if torchvision is not installed. +""" + +from __future__ import annotations + +import torch +import torch.nn as nn +import torch.nn.functional as F +from torch.utils.data import DataLoader, Subset, TensorDataset + + +class MNISTNet(nn.Module): + def __init__(self): + super().__init__() + self.fc1 = nn.Linear(784, 128) + self.fc2 = nn.Linear(128, 10) + + def forward(self, x): + x = x.view(x.size(0), -1) + x = F.relu(self.fc1(x)) + return self.fc2(x) + + +def make_model() -> nn.Module: + return MNISTNet() + + +def make_dataloader( + rank: int, world_size: int, batch_size: int = 64 +) -> DataLoader: + try: + from torchvision import datasets, transforms + + dataset = datasets.MNIST( + "/tmp/mnist_data", + train=True, + download=True, + transform=transforms.ToTensor(), + ) + except (ImportError, Exception): + # synthetic fallback — same shape as MNIST + n = 4096 + x = torch.randn(n, 1, 28, 28) + y = torch.randint(0, 10, (n,)) + dataset = TensorDataset(x, y) + + # shard by rank: interleaved assignment + indices = list(range(rank, len(dataset), world_size)) + subset = Subset(dataset, indices) + + return DataLoader(subset, batch_size=batch_size, shuffle=True, drop_last=True) diff --git a/main.py b/main.py new file mode 100644 index 0000000..61bab30 --- /dev/null +++ b/main.py @@ -0,0 +1,6 @@ +def main(): + print("Hello from vastai-utils!") + + +if __name__ == "__main__": + main() diff --git a/pyproject.toml b/pyproject.toml new file mode 100644 index 0000000..de029ae --- /dev/null +++ b/pyproject.toml @@ -0,0 +1,14 @@ +[project] +name = "vastai-utils" +version = "0.1.0" +description = "Add your description here" +requires-python = ">=3.13" +dependencies = [ + "torch>=2.10.0", +] + +[dependency-groups] +dev = [ + "pytest>=9.0.2", + "pytest-asyncio>=1.3.0", +] diff --git a/report.html b/report.html new file mode 100644 index 0000000..e3f5ca8 --- /dev/null +++ b/report.html @@ -0,0 +1,222 @@ + + + + +Training Run Report + + + +

Training Run Report

+ +
+
+
3
+
Rounds
+
+
+
2
+
Workers
+
+
+
15.6s
+
Total Time
+
+
+
4.7MB
+
Data Transferred
+
+
+
2.3MB
+
Params Pushed
+
+
+
2.3MB
+
Weights Pulled
+
+
+
8.7
+
Final |W|
+
+
+ +
+
Push params
+
Train
+
Pull weights
+
Aggregate
+
Dead / Error
+
+ +
+ + +round 0 [2/2] 4.59s + +0.0s + +0.9s + +1.8s + +2.8s + +3.7s + +4.6s +scheduler + +push 399.7KB → node 0 (0.000s) +push 399.7KB → node 1 (0.000s) +pull 399.7KB ← node 0 (0.000s) +pull 399.7KB ← node 1 (0.000s) +aggregate: 2 workers, |W|=8.75, Δ=0.0891 +worker 0 + +push 399.7KB → node 0 (0.000s) +train node 0: 4.59s exit=0 +cpu = _conversion_method_template(device=torch.device("cpu")) +pull 399.7KB ← node 0 (0.000s) +worker 1 + +push 399.7KB → node 1 (0.000s) +train node 1: 4.31s exit=0 +cpu = _conversion_method_template(device=torch.device("cpu")) +pull 399.7KB ← node 1 (0.000s) +round 1 [2/2] 5.53s + +0.0s + +1.1s + +2.2s + +3.3s + +4.4s + +5.5s +scheduler + +push 399.6KB → node 0 (0.000s) +push 399.6KB → node 1 (0.002s) +pull 399.7KB ← node 0 (0.000s) +pull 399.7KB ← node 1 (0.000s) +aggregate: 2 workers, |W|=8.74, Δ=0.0876 +worker 0 + +push 399.6KB → node 0 (0.000s) +train node 0: 5.35s exit=0 +cpu = _conversion_method_template(device=torch.device("cpu")) +pull 399.7KB ← node 0 (0.000s) +worker 1 + +push 399.6KB → node 1 (0.002s) +train node 1: 5.53s exit=0 +cpu = _conversion_method_template(device=torch.device("cpu")) +pull 399.7KB ← node 1 (0.000s) +round 2 [2/2] 5.50s + +0.0s + +1.1s + +2.2s + +3.3s + +4.4s + +5.5s +scheduler + +push 399.6KB → node 0 (0.000s) +push 399.6KB → node 1 (0.000s) +pull 399.7KB ← node 0 (0.000s) +pull 399.7KB ← node 1 (0.000s) +aggregate: 2 workers, |W|=8.74, Δ=0.0883 +worker 0 + +push 399.6KB → node 0 (0.000s) +train node 0: 5.50s exit=0 +cpu = _conversion_method_template(device=torch.device("cpu")) +pull 399.7KB ← node 0 (0.000s) +worker 1 + +push 399.6KB → node 1 (0.000s) +train node 1: 5.38s exit=0 +cpu = _conversion_method_template(device=torch.device("cpu")) +pull 399.7KB ← node 1 (0.000s) + +
+ +
+
[   0.00s] push  node=0  worker.py               1.8KB  0.000s
+[   0.00s] push  node=0  job_module.py           1.4KB  0.000s
+[   0.00s] push  node=1  worker.py               1.8KB  0.000s
+[   0.00s] push  node=1  job_module.py           1.4KB  0.000s
+[   0.00s] ──── round 0 start (2 workers) ────────────────
+[   0.00s] push  node=0  params.pt             399.7KB  0.000s
+[   0.00s] push  node=1  params.pt             399.7KB  0.000s
+[   4.32s] exec  node=1  exit=0    4.313s  | cpu = _conversion_method_template(device=torch.device("cpu"))
+[   4.59s] exec  node=0  exit=0    4.587s  | cpu = _conversion_method_template(device=torch.device("cpu"))
+[   4.59s] pull  node=0  weights.pt            399.7KB  0.000s
+[   4.59s] pull  node=1  weights.pt            399.7KB  0.000s
+[   4.59s] agg   2 workers  |W|=8.75  Δ=0.0891  0.001s
+[   4.59s] ──── round 0 end (2/2 survived, 4.59s) ────────
+[   4.59s] ──── round 1 start (2 workers) ────────────────
+[   4.59s] push  node=0  params.pt             399.6KB  0.000s
+[   4.60s] push  node=1  params.pt             399.6KB  0.002s
+[   9.95s] exec  node=0  exit=0    5.350s  | cpu = _conversion_method_template(device=torch.device("cpu"))
+[  10.12s] exec  node=1  exit=0    5.525s  | cpu = _conversion_method_template(device=torch.device("cpu"))
+[  10.12s] pull  node=0  weights.pt            399.7KB  0.000s
+[  10.12s] pull  node=1  weights.pt            399.7KB  0.000s
+[  10.13s] agg   2 workers  |W|=8.74  Δ=0.0876  0.001s
+[  10.13s] ──── round 1 end (2/2 survived, 5.53s) ────────
+[  10.13s] ──── round 2 start (2 workers) ────────────────
+[  10.13s] push  node=0  params.pt             399.6KB  0.000s
+[  10.13s] push  node=1  params.pt             399.6KB  0.000s
+[  15.51s] exec  node=1  exit=0    5.378s  | cpu = _conversion_method_template(device=torch.device("cpu"))
+[  15.63s] exec  node=0  exit=0    5.498s  | cpu = _conversion_method_template(device=torch.device("cpu"))
+[  15.63s] pull  node=0  weights.pt            399.7KB  0.000s
+[  15.63s] pull  node=1  weights.pt            399.7KB  0.000s
+[  15.63s] agg   2 workers  |W|=8.74  Δ=0.0883  0.001s
+[  15.63s] ──── round 2 end (2/2 survived, 5.50s) ────────
+
+ + + \ No newline at end of file diff --git a/reports/log.txt b/reports/log.txt new file mode 100644 index 0000000..d0403cc --- /dev/null +++ b/reports/log.txt @@ -0,0 +1,31 @@ +/usr/local/lib/python3.13/site-packages/torch/_subclasses/functional_tensor.py:283: UserWarning: Failed to initialize NumPy: No module named 'numpy' (Triggered internally at /pytorch/torch/csrc/utils/tensor_numpy.cpp:84.) + cpu = _conversion_method_template(device=torch.device("cpu")) +[ 0.15s] exec node=1 exit=0 0.142s | Warning: Permanently added 'worker-1' (ED25519) to the list of known hosts. +[ 0.15s] exec node=0 exit=0 0.144s | Warning: Permanently added 'worker-0' (ED25519) to the list of known hosts. +[ 0.30s] push node=1 job_module.py 1.4KB 0.146s +[ 0.30s] push node=0 job_module.py 1.4KB 0.146s +[ 0.30s] ──── round 0 start (2 workers) ──────────────────── +[ 0.45s] push node=0 params.pt 399.7KB 0.145s +[ 0.45s] push node=1 params.pt 399.7KB 0.146s +[ 9.86s] exec node=1 exit=0 9.416s | cpu = _conversion_method_template(device=torch.device("cpu")) +[ 9.87s] exec node=0 exit=0 9.419s | cpu = _conversion_method_template(device=torch.device("cpu")) +[ 10.04s] pull node=0 weights.pt 399.7KB 0.174s +[ 10.21s] pull node=1 weights.pt 399.7KB 0.165s +[ 10.21s] agg 2 workers |W|=8.72 Δ=0.3177 0.002s +[ 10.21s] ──── round 0 end (2/2 survived, 9.91s) ──────────── +2026-02-05 15:59:23,615 INFO sched.scheduler: Round 0: 2/2 workers + + ▓ push █ train ▒ pull ● agg ✗ dead ░ wait + + round 0 [2/2] 9.91s + scheduler ▓▓░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░▒▒● + worker 0 █████████████████████████████████████████████████████████▒▒─ + worker 1 ██████████████████████████████████████████████████████████▒▒ + + +Round Workers Params↑ Weights↓ Train(max) Agg Total Δ norm +──────────────────────────────────────────────────────────────────────────── +0 2/2 799.4KB 799.4KB 9.42 0.002 9.91 0.3177 +──────────────────────────────────────────────────────────────────────────── +total 799.4KB 799.4KB 9.91 +Report: /app/reports/report.html diff --git a/reports/report.html b/reports/report.html new file mode 100644 index 0000000..e871630 --- /dev/null +++ b/reports/report.html @@ -0,0 +1,140 @@ + + + + +Training Run Report + + + +

Training Run Report

+ +
+
+
1
+
Rounds
+
+
+
2
+
Workers
+
+
+
9.9s
+
Total Time
+
+
+
1.6MB
+
Data Transferred
+
+
+
799.4KB
+
Params Pushed
+
+
+
799.4KB
+
Weights Pulled
+
+
+
8.7
+
Final |W|
+
+
+ +
+
Push params
+
Train
+
Pull weights
+
Aggregate
+
Dead / Error
+
+ +
+ + +round 0 [2/2] 9.91s + +0.0s + +2.0s + +4.0s + +5.9s + +7.9s + +9.9s +scheduler + +push 399.7KB → node 0 (0.145s) +push 399.7KB → node 1 (0.146s) +pull 399.7KB ← node 0 (0.174s) +pull 399.7KB ← node 1 (0.165s) +aggregate: 2 workers, |W|=8.72, Δ=0.3177 +worker 0 + +push 399.7KB → node 0 (0.145s) +train node 0: 9.42s exit=0 +cpu = _conversion_method_template(device=torch.device("cpu")) +pull 399.7KB ← node 0 (0.174s) +worker 1 + +push 399.7KB → node 1 (0.146s) +train node 1: 9.42s exit=0 +cpu = _conversion_method_template(device=torch.device("cpu")) +pull 399.7KB ← node 1 (0.165s) + +
+ +
+
[   0.15s] exec  node=1  exit=0    0.142s  | Warning: Permanently added 'worker-1' (ED25519) to the list of known hosts.
+[   0.15s] exec  node=0  exit=0    0.144s  | Warning: Permanently added 'worker-0' (ED25519) to the list of known hosts.
+[   0.30s] push  node=1  job_module.py           1.4KB  0.146s
+[   0.30s] push  node=0  job_module.py           1.4KB  0.146s
+[   0.30s] ──── round 0 start (2 workers) ────────────────
+[   0.45s] push  node=0  params.pt             399.7KB  0.145s
+[   0.45s] push  node=1  params.pt             399.7KB  0.146s
+[   9.86s] exec  node=1  exit=0    9.416s  | cpu = _conversion_method_template(device=torch.device("cpu"))
+[   9.87s] exec  node=0  exit=0    9.419s  | cpu = _conversion_method_template(device=torch.device("cpu"))
+[  10.04s] pull  node=0  weights.pt            399.7KB  0.174s
+[  10.21s] pull  node=1  weights.pt            399.7KB  0.165s
+[  10.21s] agg   2 workers  |W|=8.72  Δ=0.3177  0.002s
+[  10.21s] ──── round 0 end (2/2 survived, 9.91s) ────────
+
+ + + \ No newline at end of file diff --git a/sched/__init__.py b/sched/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/sched/__main__.py b/sched/__main__.py new file mode 100644 index 0000000..a20fd8d --- /dev/null +++ b/sched/__main__.py @@ -0,0 +1,122 @@ +"""Entry point: python -m sched run --script jobs/mnist.py --trace --report report.html""" + +from __future__ import annotations + +import argparse +import asyncio +import logging +import sys + +from .job import Job +from .scheduler import Scheduler +from .transport import SSHTransport + + +async def _run_mock(sched: Scheduler, job: Job): + """Run with mock transport: create local nodes, skip provisioning.""" + from .scheduler import Node + + for rank in range(job.num_nodes): + conn = await sched.transport.connect(f"mock{rank}", 22) + sched.nodes.append( + Node(rank=rank, host=f"mock{rank}", port=22, conn=conn, status="ready") + ) + sched._wrap_connections() + await sched.deploy() + + job_mod = sched._load_job_module() + state = job_mod.make_model().state_dict() + for r in range(job.rounds): + state = await sched.run_round(state, r) + + +async def _run_compose(sched: Scheduler, job: Job): + """Run against docker-compose workers: worker-0:22, worker-1:22, etc.""" + from .scheduler import Node + + for rank in range(job.num_nodes): + host = f"worker-{rank}" + conn = await sched.transport.connect(host, 22) + sched.nodes.append( + Node(rank=rank, host=host, port=22, conn=conn, status="ready") + ) + sched._wrap_connections() + await sched.deploy() + + job_mod = sched._load_job_module() + state = job_mod.make_model().state_dict() + for r in range(job.rounds): + state = await sched.run_round(state, r) + + +def main(): + logging.basicConfig( + level=logging.INFO, + format="%(asctime)s %(levelname)s %(name)s: %(message)s", + ) + + parser = argparse.ArgumentParser(prog="sched") + sub = parser.add_subparsers(dest="cmd") + + p = sub.add_parser("run", help="Run a training job") + p.add_argument("--script", required=True, help="Path to job module") + p.add_argument("--nodes", type=int, default=2) + p.add_argument("--vram", type=int, default=8) + p.add_argument("--rounds", type=int, default=10) + p.add_argument("--local-steps", type=int, default=100) + p.add_argument("--lr", type=float, default=0.01) + p.add_argument("--batch-size", type=int, default=64) + p.add_argument("--trace", action="store_true", help="Enable live tracing to stderr") + p.add_argument("--report", metavar="PATH", help="Write HTML report to PATH") + p.add_argument("--mock", action="store_true", help="Use local mock transport (no vast.ai)") + p.add_argument("--compose", action="store_true", help="Connect to docker-compose workers (worker-0:22, worker-1:22, ...)") + p.add_argument("--docker-image", metavar="IMAGE", help="Override docker_image in job config") + + args = parser.parse_args() + if not args.cmd: + parser.print_help() + return 1 + + tracer = None + if args.trace or args.report: + from .trace import Tracer + tracer = Tracer(file=sys.stderr) + + job = Job( + script=args.script, + num_nodes=args.nodes, + min_vram=args.vram, + rounds=args.rounds, + local_steps=args.local_steps, + lr=args.lr, + batch_size=args.batch_size, + ) + + if args.docker_image: + job.docker_image = args.docker_image + + if args.mock: + from tests.mock_transport import MockTransport + transport = MockTransport() + sched = Scheduler(job, transport, tracer=tracer) + asyncio.run(_run_mock(sched, job)) + elif args.compose: + sched = Scheduler(job, SSHTransport(), tracer=tracer) + asyncio.run(_run_compose(sched, job)) + else: + sched = Scheduler(job, SSHTransport(), tracer=tracer) + asyncio.run(sched.run()) + + if tracer: + tracer.gantt() + tracer.summary() + if args.report: + from .report import generate_html + path = generate_html(tracer, args.report) + print(f"Report: {path}", file=sys.stderr) + + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/sched/__pycache__/__init__.cpython-313.pyc b/sched/__pycache__/__init__.cpython-313.pyc new file mode 100644 index 0000000..5b17d9e Binary files /dev/null and b/sched/__pycache__/__init__.cpython-313.pyc differ diff --git a/sched/__pycache__/__main__.cpython-313.pyc b/sched/__pycache__/__main__.cpython-313.pyc new file mode 100644 index 0000000..a9a6d1b Binary files /dev/null and b/sched/__pycache__/__main__.cpython-313.pyc differ diff --git a/sched/__pycache__/aggregator.cpython-313.pyc b/sched/__pycache__/aggregator.cpython-313.pyc new file mode 100644 index 0000000..7dfc448 Binary files /dev/null and b/sched/__pycache__/aggregator.cpython-313.pyc differ diff --git a/sched/__pycache__/job.cpython-313.pyc b/sched/__pycache__/job.cpython-313.pyc new file mode 100644 index 0000000..68d59e8 Binary files /dev/null and b/sched/__pycache__/job.cpython-313.pyc differ diff --git a/sched/__pycache__/report.cpython-313.pyc b/sched/__pycache__/report.cpython-313.pyc new file mode 100644 index 0000000..1ff36ad Binary files /dev/null and b/sched/__pycache__/report.cpython-313.pyc differ diff --git a/sched/__pycache__/scheduler.cpython-313.pyc b/sched/__pycache__/scheduler.cpython-313.pyc new file mode 100644 index 0000000..3e8aca5 Binary files /dev/null and b/sched/__pycache__/scheduler.cpython-313.pyc differ diff --git a/sched/__pycache__/trace.cpython-313.pyc b/sched/__pycache__/trace.cpython-313.pyc new file mode 100644 index 0000000..dbaa914 Binary files /dev/null and b/sched/__pycache__/trace.cpython-313.pyc differ diff --git a/sched/__pycache__/transport.cpython-313.pyc b/sched/__pycache__/transport.cpython-313.pyc new file mode 100644 index 0000000..4a5113f Binary files /dev/null and b/sched/__pycache__/transport.cpython-313.pyc differ diff --git a/sched/__pycache__/vastai.cpython-313.pyc b/sched/__pycache__/vastai.cpython-313.pyc new file mode 100644 index 0000000..7335ca2 Binary files /dev/null and b/sched/__pycache__/vastai.cpython-313.pyc differ diff --git a/sched/aggregator.py b/sched/aggregator.py new file mode 100644 index 0000000..e560325 --- /dev/null +++ b/sched/aggregator.py @@ -0,0 +1,18 @@ +from __future__ import annotations + +from collections import OrderedDict + +import torch + + +def average_weights(weight_dicts: list[OrderedDict]) -> OrderedDict: + """Average model state dicts from multiple workers.""" + if not weight_dicts: + raise ValueError("No weights to average") + if len(weight_dicts) == 1: + return weight_dicts[0] + + avg = OrderedDict() + for key in weight_dicts[0]: + avg[key] = torch.stack([w[key].float() for w in weight_dicts]).mean(dim=0) + return avg diff --git a/sched/job.py b/sched/job.py new file mode 100644 index 0000000..e26a248 --- /dev/null +++ b/sched/job.py @@ -0,0 +1,17 @@ +from __future__ import annotations + +from dataclasses import dataclass, field + + +@dataclass +class Job: + script: str + num_nodes: int = 2 + min_vram: int = 8 + docker_image: str = "pytorch/pytorch:2.1.0-cuda12.1-cudnn8-runtime" + rounds: int = 10 + local_steps: int = 100 + lr: float = 0.01 + batch_size: int = 64 + checkpoint_dir: str | None = None + env: dict[str, str] = field(default_factory=dict) diff --git a/sched/report.py b/sched/report.py new file mode 100644 index 0000000..29e7e4e --- /dev/null +++ b/sched/report.py @@ -0,0 +1,403 @@ +"""Generate a self-contained HTML report with an SVG timeline from tracer events.""" + +from __future__ import annotations + +import html +import os +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from .trace import Tracer + +# ------------------------------------------------------------------ +# colours +# ------------------------------------------------------------------ + +COLORS = { + "push": "#3b82f6", + "exec": "#22c55e", + "pull": "#f97316", + "aggregate": "#8b5cf6", + "dead": "#ef4444", +} + +BG = "#0f172a" +CARD = "#1e293b" +TEXT = "#e2e8f0" +MUTED = "#94a3b8" +GRID = "#334155" + + +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" + + +# ------------------------------------------------------------------ +# SVG builder +# ------------------------------------------------------------------ + +def _build_svg(tracer: Tracer) -> str: + rounds = tracer._collect_rounds() + max_rank = tracer._max_rank() + if not rounds: + return '' + + # layout constants + label_w = 100 + right_pad = 20 + row_h = 28 + row_gap = 4 + round_gap = 24 + round_header_h = 22 + chart_w = 700 + total_w = label_w + chart_w + right_pad + + # compute total height + rows_per_round = 1 + max_rank + 1 # scheduler + workers + n_rounds = len(rounds) + total_h = ( + n_rounds * (round_header_h + rows_per_round * (row_h + row_gap) + round_gap) + + 40 # bottom time axis + ) + + parts: list[str] = [] + parts.append( + f'' + ) + parts.append( + f'' + ) + + y_cursor = 12 + + for rnum in sorted(rounds.keys()): + rd = rounds[rnum] + t0 = rd["start"] + dur = rd["duration"] or 0.001 + surv = rd["survivors"] + total_nodes = rd["total"] + + def x_pos(t: float) -> float: + frac = max(0.0, min(1.0, (t - t0) / dur)) + return label_w + frac * chart_w + + def bar_w(d: float) -> float: + return max(3, d / dur * chart_w) # min 3px so it's visible + + # round header + parts.append( + f'' + f"round {rnum} [{surv}/{total_nodes}] {dur:.2f}s" + ) + y_cursor += round_header_h + + # time gridlines + n_ticks = 5 + for i in range(n_ticks + 1): + frac = i / n_ticks + gx = label_w + frac * chart_w + parts.append( + f'' + ) + t_label = frac * dur + parts.append( + f'' + f"{t_label:.1f}s" + ) + + # --- scheduler row --- + row_y = y_cursor + parts.append( + f'scheduler' + ) + # background bar + parts.append( + f'' + ) + for ev in rd["events"]: + d = ev.data.get("duration_s", 0) + if ev.kind in ("push", "pull", "aggregate"): + color = COLORS.get(ev.kind, COLORS["push"]) + bx = x_pos(ev.time - d) + bw = bar_w(d) + tooltip = _tooltip(ev) + parts.append( + f'' + f"{html.escape(tooltip)}" + ) + y_cursor += row_h + row_gap + + # --- worker rows --- + for rank in range(max_rank + 1): + row_y = y_cursor + parts.append( + f'worker {rank}' + ) + parts.append( + f'' + ) + + rank_evs = [e for e in rd["events"] if e.rank == rank] + dead_x: float | None = None + + for ev in rank_evs: + d = ev.data.get("duration_s", 0) + is_err = bool(ev.data.get("error")) + is_bad_exit = ev.kind == "exec" and ev.data.get("exit_code", 0) != 0 + + if is_err or is_bad_exit: + color = COLORS["dead"] + dead_x = x_pos(ev.time) + else: + color = COLORS.get(ev.kind, COLORS["push"]) + + bx = x_pos(ev.time - d) + bw = bar_w(d) + tooltip = _tooltip(ev) + parts.append( + f'' + f"{html.escape(tooltip)}" + ) + + # dead zone: hatched red from failure to end + if dead_x is not None: + dw = label_w + chart_w - dead_x + if dw > 0: + parts.append( + f'' + ) + + # no events at all — full dead bar + if not rank_evs and rank < rd.get("num_workers", 0): + parts.append( + f'' + f"worker {rank}: no response" + ) + + y_cursor += row_h + row_gap + + y_cursor += round_gap + + parts.append("") + return "\n".join(parts) + + +def _tooltip(ev) -> str: + kind = ev.kind + d = ev.data + dur = d.get("duration_s", 0) + err = d.get("error") + if err: + return f"{kind} node={ev.rank}: FAILED — {err}" + if kind == "push": + return f"push {_fmt_bytes(d.get('size_bytes', 0))} → node {ev.rank} ({dur:.3f}s)" + if kind == "pull": + return f"pull {_fmt_bytes(d.get('size_bytes', 0))} ← node {ev.rank} ({dur:.3f}s)" + if kind == "exec": + tail = d.get("output_tail", "") + return f"train node {ev.rank}: {dur:.2f}s exit={d.get('exit_code', '?')}\n{tail}" + if kind == "aggregate": + return ( + f"aggregate: {d.get('num_workers', '?')} workers, " + f"|W|={d.get('weight_norm', 0):.2f}, Δ={d.get('delta_norm', 0):.4f}" + ) + return f"{kind}: {d}" + + +# ------------------------------------------------------------------ +# Full HTML page +# ------------------------------------------------------------------ + +def generate_html(tracer: Tracer, path: str) -> str: + """Write a self-contained HTML report to `path`. Returns the path.""" + rounds = tracer._collect_rounds() + + # compute summary stats + n_rounds = len(rounds) + total_time = sum(r["duration"] for r in rounds.values()) + total_push = 0 + total_pull = 0 + for rd in rounds.values(): + total_push += sum( + e.data.get("size_bytes", 0) for e in rd["events"] + if e.kind == "push" and not e.data.get("error") + ) + total_pull += sum( + e.data.get("size_bytes", 0) for e in rd["events"] + if e.kind == "pull" and not e.data.get("error") + ) + max_workers = max((r["total"] for r in rounds.values()), default=0) + final_wnorm = list(rounds.values())[-1]["weight_norm"] if rounds else 0 + + svg = _build_svg(tracer) + + page = f""" + + + +Training Run Report + + + +

Training Run Report

+ +
+
+
{n_rounds}
+
Rounds
+
+
+
{max_workers}
+
Workers
+
+
+
{total_time:.1f}s
+
Total Time
+
+
+
{_fmt_bytes(total_push + total_pull)}
+
Data Transferred
+
+
+
{_fmt_bytes(total_push)}
+
Params Pushed
+
+
+
{_fmt_bytes(total_pull)}
+
Weights Pulled
+
+
+
{final_wnorm:.1f}
+
Final |W|
+
+
+ +
+
Push params
+
Train
+
Pull weights
+
Aggregate
+
Dead / Error
+
+ +
+{svg} +
+ +
+
{_render_event_log(tracer)}
+
+ + +""" + + path = os.path.abspath(path) + os.makedirs(os.path.dirname(path) or ".", exist_ok=True) + with open(path, "w") as f: + f.write(page) + return path + + +def _render_event_log(tracer: Tracer) -> str: + """Render the text event log as HTML-escaped pre-formatted text.""" + lines = [] + for ev in tracer.events: + err = ev.data.get("error") + t = f"[{ev.time:7.2f}s]" + node = f"node={ev.rank}" if ev.rank is not None else " " + + match ev.kind: + case "push" | "pull": + sz = _fmt_bytes(ev.data.get("size_bytes", 0)) + dur = ev.data.get("duration_s", 0) + p = os.path.basename(ev.data.get("path", "")) + if err: + line = f'{t} {ev.kind:5s} {node} {p:<20s} !! {html.escape(err)}' + else: + line = f"{t} {ev.kind:5s} {node} {p:<20s} {sz:>8s} {dur:.3f}s" + case "exec": + dur = ev.data.get("duration_s", 0) + ec = ev.data.get("exit_code", "?") + if err: + line = f'{t} exec {node} !! {html.escape(err)}' + else: + tail = html.escape(ev.data.get("output_tail", "")) + line = f"{t} exec {node} exit={ec:<3} {dur:.3f}s" + if tail: + line += f" | {tail}" + case "round_start": + r = ev.data.get("round_num", "?") + n = ev.data.get("active_nodes", "?") + line = f"{t} {'─'*4} round {r} start ({n} workers) {'─'*16}" + case "round_end": + r = ev.data.get("round_num", "?") + s = ev.data.get("survivors", "?") + tot = ev.data.get("total_nodes", "?") + dur = ev.data.get("duration_s", 0) + line = f"{t} {'─'*4} round {r} end ({s}/{tot} survived, {dur:.2f}s) {'─'*8}" + case "aggregate": + n = ev.data.get("num_workers", "?") + wn = ev.data.get("weight_norm", 0) + dn = ev.data.get("delta_norm", 0) + dur = ev.data.get("duration_s", 0) + line = f"{t} agg {n} workers |W|={wn:.2f} Δ={dn:.4f} {dur:.3f}s" + case _: + line = f"{t} {ev.kind}" + + lines.append(line) + return "\n".join(lines) diff --git a/sched/scheduler.py b/sched/scheduler.py new file mode 100644 index 0000000..38e0e8f --- /dev/null +++ b/sched/scheduler.py @@ -0,0 +1,300 @@ +from __future__ import annotations + +import asyncio +import importlib.util +import logging +import os +import tempfile +import time +from collections import OrderedDict +from dataclasses import dataclass + +import torch + +from .aggregator import average_weights +from .job import Job +from .trace import Tracer, TracingConnection +from .transport import Connection, Transport + +log = logging.getLogger(__name__) + + +@dataclass +class Node: + rank: int + vast_id: int | None = None + host: str | None = None + port: int | None = None + conn: Connection | None = None + status: str = "pending" # pending | ready | running | dead | done + + +class Scheduler: + def __init__(self, job: Job, transport: Transport, tracer: Tracer | None = None): + self.job = job + self.transport = transport + self.tracer = tracer + self.nodes: list[Node] = [] + self.work_dir = tempfile.mkdtemp(prefix="sched_") + + # ------------------------------------------------------------------ + # helpers + # ------------------------------------------------------------------ + + def _active_nodes(self) -> list[Node]: + return [n for n in self.nodes if n.conn and n.status not in ("dead", "done")] + + def _load_job_module(self): + spec = importlib.util.spec_from_file_location( + "job_module", os.path.abspath(self.job.script) + ) + mod = importlib.util.module_from_spec(spec) + spec.loader.exec_module(mod) + return mod + + def _wrap_connections(self): + """Wrap node connections with TracingConnection if tracer is set.""" + if not self.tracer: + return + for node in self.nodes: + if node.conn and not isinstance(node.conn, TracingConnection): + node.conn = TracingConnection(node.conn, node.rank, self.tracer) + + # ------------------------------------------------------------------ + # lifecycle phases + # ------------------------------------------------------------------ + + async def provision(self): + """Rent instances from vast.ai and wait for SSH. Populates self.nodes.""" + from . import vastai + + offers = await vastai.search_offers(self.job.min_vram, limit=self.job.num_nodes) + if len(offers) < self.job.num_nodes: + raise RuntimeError( + f"Need {self.job.num_nodes} instances, only {len(offers)} offers" + ) + + for rank, offer in enumerate(offers[: self.job.num_nodes]): + await vastai.rent_instance(offer["id"], self.job.docker_image) + self.nodes.append(Node(rank=rank, vast_id=offer["id"], status="provisioning")) + + # poll until SSH ready (up to 5 min) + for _ in range(60): + instances = await vastai.get_instances() + by_id = {i["id"]: i for i in instances} + all_ready = True + for node in self.nodes: + inst = by_id.get(node.vast_id) + if inst and inst.get("ssh_host"): + node.host = inst["ssh_host"] + node.port = inst["ssh_port"] + node.status = "ready" + else: + all_ready = False + if all_ready: + break + await asyncio.sleep(5) + + # open connections + for node in self.nodes: + if node.host: + node.conn = await self.transport.connect(node.host, node.port) + + self._wrap_connections() + + ready = self._active_nodes() + if not ready: + raise RuntimeError("No nodes became ready") + log.info(f"Provisioned {len(ready)}/{self.job.num_nodes} nodes") + + async def deploy(self): + """Push worker script and job module to all active nodes. + + Skips pushing worker.py if it already exists on the remote + (e.g. baked into a Docker image). Always pushes the job module. + """ + self._wrap_connections() + + worker_src = os.path.join( + os.path.dirname(__file__), os.pardir, "worker", "worker.py" + ) + worker_src = os.path.abspath(worker_src) + job_src = os.path.abspath(self.job.script) + + async def _deploy_one(node: Node): + rc, _ = await node.conn.exec("test -f /workspace/worker.py") + if rc != 0: + await node.conn.push(worker_src, "/workspace/worker.py") + else: + log.debug(f"Node {node.rank}: worker.py already present, skipping push") + await node.conn.push(job_src, "/workspace/job_module.py") + + results = await asyncio.gather( + *[_deploy_one(n) for n in self._active_nodes()], + return_exceptions=True, + ) + for node, res in zip(self._active_nodes(), results): + if isinstance(res, Exception): + log.warning(f"Deploy failed on node {node.rank}: {res}") + node.status = "dead" + + if not self._active_nodes(): + raise RuntimeError("Deploy failed on all nodes") + + async def run_round( + self, state_dict: OrderedDict, round_num: int + ) -> OrderedDict: + """Single FedAvg round: push params -> workers train -> pull weights -> average.""" + round_t0 = time.monotonic() + + params_path = os.path.join(self.work_dir, "params.pt") + torch.save(state_dict, params_path) + + active = self._active_nodes() + if not active: + raise RuntimeError("No active workers") + + if self.tracer: + self.tracer.emit( + "round_start", + round_num=round_num, + active_nodes=len(active), + ) + + # --- push params --- + push_res = await asyncio.gather( + *[n.conn.push(params_path, "/workspace/params.pt") for n in active], + return_exceptions=True, + ) + for node, res in zip(active, push_res): + if isinstance(res, Exception): + log.warning(f"Param push failed for node {node.rank}: {res}") + node.status = "dead" + + active = self._active_nodes() + if not active: + raise RuntimeError("No workers alive after param push") + + # --- run training --- + cmd_template = ( + "cd /workspace && python worker.py" + " --params /workspace/params.pt" + " --output /workspace/weights.pt" + " --rank {rank}" + f" --world-size {self.job.num_nodes}" + f" --local-steps {self.job.local_steps}" + f" --lr {self.job.lr}" + f" --batch-size {self.job.batch_size}" + ) + exec_res = await asyncio.gather( + *[ + n.conn.exec(cmd_template.format(rank=n.rank), timeout=300.0) + for n in active + ], + return_exceptions=True, + ) + + survivors = [] + for node, res in zip(active, exec_res): + if isinstance(res, Exception): + log.warning(f"Node {node.rank} raised: {res}") + node.status = "dead" + elif res[0] != 0: + log.warning(f"Node {node.rank} exit {res[0]}: {res[1][:200]}") + node.status = "dead" + else: + survivors.append(node) + + if not survivors: + raise RuntimeError("All workers failed") + + # --- pull weights --- + weight_dicts: list[OrderedDict] = [] + for node in survivors: + local_path = os.path.join(self.work_dir, f"weights_{node.rank}.pt") + try: + await node.conn.pull("/workspace/weights.pt", local_path) + w = torch.load(local_path, map_location="cpu", weights_only=True) + weight_dicts.append(w) + except Exception as e: + log.warning(f"Weight pull failed for node {node.rank}: {e}") + node.status = "dead" + + if not weight_dicts: + raise RuntimeError("No weights collected") + + # --- aggregate --- + agg_t0 = time.monotonic() + new_state = average_weights(weight_dicts) + agg_dur = time.monotonic() - agg_t0 + + if self.tracer: + weight_norm = sum(v.float().norm().item() for v in new_state.values()) + delta_norm = sum( + (new_state[k].float() - state_dict[k].float()).norm().item() + for k in new_state + ) + self.tracer.emit( + "aggregate", + num_workers=len(weight_dicts), + weight_norm=weight_norm, + delta_norm=delta_norm, + duration_s=agg_dur, + ) + + round_dur = time.monotonic() - round_t0 + + if self.tracer: + self.tracer.emit( + "round_end", + round_num=round_num, + survivors=len(weight_dicts), + total_nodes=len(self.nodes), + duration_s=round_dur, + ) + + log.info( + f"Round {round_num}: {len(weight_dicts)}/{len(self.nodes)} workers" + ) + return new_state + + async def run(self) -> OrderedDict: + """Full pipeline: provision -> deploy -> train rounds -> cleanup. + + Returns the final averaged state_dict. + """ + job_mod = self._load_job_module() + model = job_mod.make_model() + state = model.state_dict() + + log.info( + f"Job start: {self.job.rounds} rounds, " + f"{self.job.num_nodes} nodes, " + f"{self.job.local_steps} local steps/round" + ) + + await self.provision() + await self.deploy() + + for r in range(self.job.rounds): + state = await self.run_round(state, r) + + if self.tracer: + self.tracer.summary() + + log.info("Job complete") + await self.cleanup() + return state + + async def cleanup(self): + """Close connections, destroy instances.""" + from . import vastai + + for node in self.nodes: + if node.conn: + await node.conn.close() + if node.vast_id: + try: + await vastai.destroy_instance(node.vast_id) + except Exception: + pass diff --git a/sched/trace.py b/sched/trace.py new file mode 100644 index 0000000..51ddfdc --- /dev/null +++ b/sched/trace.py @@ -0,0 +1,424 @@ +"""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() diff --git a/sched/transport.py b/sched/transport.py new file mode 100644 index 0000000..2e8a3eb --- /dev/null +++ b/sched/transport.py @@ -0,0 +1,84 @@ +from __future__ import annotations + +import asyncio +from typing import Protocol, runtime_checkable + + +@runtime_checkable +class Connection(Protocol): + async def exec(self, cmd: str, timeout: float = 30.0) -> tuple[int, str]: ... + async def push(self, local_path: str, remote_path: str) -> None: ... + async def pull(self, remote_path: str, local_path: str) -> None: ... + async def close(self) -> None: ... + + +@runtime_checkable +class Transport(Protocol): + async def connect(self, host: str, port: int) -> Connection: ... + + +class SSHConnection: + def __init__(self, host: str, port: int): + self.host = host + self.port = port + + async def exec(self, cmd: str, timeout: float = 30.0) -> tuple[int, str]: + proc = await asyncio.create_subprocess_exec( + "ssh", + "-p", str(self.port), + "-o", "StrictHostKeyChecking=no", + "-o", "ConnectTimeout=5", + "-o", "BatchMode=yes", + f"root@{self.host}", + cmd, + stdout=asyncio.subprocess.PIPE, + stderr=asyncio.subprocess.PIPE, + ) + try: + stdout, stderr = await asyncio.wait_for( + proc.communicate(), timeout=timeout + ) + return proc.returncode, (stdout + stderr).decode() + except asyncio.TimeoutError: + proc.kill() + return 1, "timeout" + + async def push(self, local_path: str, remote_path: str) -> None: + proc = await asyncio.create_subprocess_exec( + "scp", + "-P", str(self.port), + "-o", "StrictHostKeyChecking=no", + local_path, + f"root@{self.host}:{remote_path}", + stdout=asyncio.subprocess.PIPE, + stderr=asyncio.subprocess.PIPE, + ) + _, stderr = await proc.communicate() + if proc.returncode != 0: + raise ConnectionError( + f"scp push failed to {self.host}:{self.port}: {stderr.decode()}" + ) + + async def pull(self, remote_path: str, local_path: str) -> None: + proc = await asyncio.create_subprocess_exec( + "scp", + "-P", str(self.port), + "-o", "StrictHostKeyChecking=no", + f"root@{self.host}:{remote_path}", + local_path, + stdout=asyncio.subprocess.PIPE, + stderr=asyncio.subprocess.PIPE, + ) + _, stderr = await proc.communicate() + if proc.returncode != 0: + raise ConnectionError( + f"scp pull failed from {self.host}:{self.port}: {stderr.decode()}" + ) + + async def close(self) -> None: + pass + + +class SSHTransport: + async def connect(self, host: str, port: int) -> SSHConnection: + return SSHConnection(host, port) diff --git a/sched/vastai.py b/sched/vastai.py new file mode 100644 index 0000000..bfccb42 --- /dev/null +++ b/sched/vastai.py @@ -0,0 +1,47 @@ +"""Async wrappers around the vast.ai CLI. Cannibalized from dtrain.""" + +from __future__ import annotations + +import asyncio +import json + + +async def _run(*args: str) -> str | None: + proc = await asyncio.create_subprocess_exec( + "vastai", *args, + stdout=asyncio.subprocess.PIPE, + stderr=asyncio.subprocess.PIPE, + ) + stdout, stderr = await proc.communicate() + if proc.returncode != 0: + return None + return stdout.decode() + + +async def search_offers(min_vram: int, limit: int = 10) -> list[dict]: + query = f"gpu_ram>={min_vram} inet_down>=100" + output = await _run( + "search", "offers", query, "-o", "dph_total", + "--limit", str(limit), "--raw", + ) + if not output: + return [] + return json.loads(output) + + +async def rent_instance(offer_id: int, image: str, disk: int = 20) -> str | None: + return await _run( + "create", "instance", str(offer_id), + "--image", image, "--disk", str(disk), + ) + + +async def get_instances() -> list[dict]: + output = await _run("show", "instances", "--raw") + if not output: + return [] + return json.loads(output) + + +async def destroy_instance(instance_id: int) -> None: + await _run("destroy", "instance", str(instance_id)) diff --git a/tests/__init__.py b/tests/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tests/__pycache__/__init__.cpython-313.pyc b/tests/__pycache__/__init__.cpython-313.pyc new file mode 100644 index 0000000..4657d92 Binary files /dev/null and b/tests/__pycache__/__init__.cpython-313.pyc differ diff --git a/tests/__pycache__/mock_transport.cpython-313.pyc b/tests/__pycache__/mock_transport.cpython-313.pyc new file mode 100644 index 0000000..0112d13 Binary files /dev/null and b/tests/__pycache__/mock_transport.cpython-313.pyc differ diff --git a/tests/__pycache__/test_aggregator.cpython-313-pytest-9.0.2.pyc b/tests/__pycache__/test_aggregator.cpython-313-pytest-9.0.2.pyc new file mode 100644 index 0000000..38f5468 Binary files /dev/null and b/tests/__pycache__/test_aggregator.cpython-313-pytest-9.0.2.pyc differ diff --git a/tests/__pycache__/test_scheduler.cpython-313-pytest-9.0.2.pyc b/tests/__pycache__/test_scheduler.cpython-313-pytest-9.0.2.pyc new file mode 100644 index 0000000..ec061c2 Binary files /dev/null and b/tests/__pycache__/test_scheduler.cpython-313-pytest-9.0.2.pyc differ diff --git a/tests/__pycache__/test_trace.cpython-313-pytest-9.0.2.pyc b/tests/__pycache__/test_trace.cpython-313-pytest-9.0.2.pyc new file mode 100644 index 0000000..23a80b6 Binary files /dev/null and b/tests/__pycache__/test_trace.cpython-313-pytest-9.0.2.pyc differ diff --git a/tests/mock_transport.py b/tests/mock_transport.py new file mode 100644 index 0000000..dca8d55 --- /dev/null +++ b/tests/mock_transport.py @@ -0,0 +1,85 @@ +"""Mock transport that runs workers locally in isolated temp directories. + +Each MockConnection gets its own temp dir standing in for /workspace/. +push/pull copy files into/out of that dir. exec runs commands as local +subprocesses with /workspace paths rewritten to the temp dir. + +Supports injecting failures: call conn.kill() to simulate a node dying. +""" + +from __future__ import annotations + +import asyncio +import os +import shutil +import tempfile + + +class MockConnection: + def __init__(self, root_dir: str): + self.root_dir = root_dir + self._alive = True + + async def exec(self, cmd: str, timeout: float = 30.0) -> tuple[int, str]: + if not self._alive: + raise ConnectionError("node is dead") + + cmd = cmd.replace("/workspace", self.root_dir) + + proc = await asyncio.create_subprocess_shell( + cmd, + cwd=self.root_dir, + stdout=asyncio.subprocess.PIPE, + stderr=asyncio.subprocess.PIPE, + ) + try: + stdout, stderr = await asyncio.wait_for( + proc.communicate(), timeout=timeout + ) + return proc.returncode, (stdout + stderr).decode() + except asyncio.TimeoutError: + proc.kill() + return 1, "timeout" + + async def push(self, local_path: str, remote_path: str) -> None: + if not self._alive: + raise ConnectionError("node is dead") + dest = remote_path.replace("/workspace", self.root_dir) + os.makedirs(os.path.dirname(dest) or self.root_dir, exist_ok=True) + shutil.copy2(local_path, dest) + + async def pull(self, remote_path: str, local_path: str) -> None: + if not self._alive: + raise ConnectionError("node is dead") + src = remote_path.replace("/workspace", self.root_dir) + if not os.path.exists(src): + raise FileNotFoundError(f"remote file not found: {src}") + os.makedirs(os.path.dirname(local_path) or ".", exist_ok=True) + shutil.copy2(src, local_path) + + async def close(self) -> None: + pass + + def kill(self): + """Simulate node death.""" + self._alive = False + + def revive(self): + """Bring node back (for testing reconnect scenarios).""" + self._alive = True + + +class MockTransport: + """Creates MockConnections backed by temp directories.""" + + def __init__(self): + self._dirs: list[str] = [] + + async def connect(self, host: str, port: int) -> MockConnection: + d = tempfile.mkdtemp(prefix=f"mock_node_{host}_{port}_") + self._dirs.append(d) + return MockConnection(d) + + def cleanup(self): + for d in self._dirs: + shutil.rmtree(d, ignore_errors=True) diff --git a/tests/test_aggregator.py b/tests/test_aggregator.py new file mode 100644 index 0000000..a149662 --- /dev/null +++ b/tests/test_aggregator.py @@ -0,0 +1,57 @@ +from collections import OrderedDict + +import pytest +import torch + +from sched.aggregator import average_weights + + +def _make_state(val: float) -> OrderedDict: + return OrderedDict( + weight=torch.full((3, 4), val), + bias=torch.full((3,), val), + ) + + +def test_average_identical(): + w = _make_state(1.0) + result = average_weights([w, w]) + for key in w: + assert torch.allclose(result[key], w[key]) + + +def test_average_different(): + a = _make_state(0.0) + b = _make_state(2.0) + result = average_weights([a, b]) + expected = _make_state(1.0) + for key in expected: + assert torch.allclose(result[key], expected[key]) + + +def test_average_three(): + a = _make_state(0.0) + b = _make_state(3.0) + c = _make_state(6.0) + result = average_weights([a, b, c]) + expected = _make_state(3.0) + for key in expected: + assert torch.allclose(result[key], expected[key]) + + +def test_single(): + w = _make_state(5.0) + result = average_weights([w]) + for key in w: + assert torch.allclose(result[key], w[key]) + + +def test_empty_raises(): + with pytest.raises(ValueError): + average_weights([]) + + +def test_preserves_keys(): + w = OrderedDict(fc1_weight=torch.ones(2, 2), fc1_bias=torch.zeros(2)) + result = average_weights([w, w]) + assert list(result.keys()) == ["fc1_weight", "fc1_bias"] diff --git a/tests/test_scheduler.py b/tests/test_scheduler.py new file mode 100644 index 0000000..b2a99d7 --- /dev/null +++ b/tests/test_scheduler.py @@ -0,0 +1,261 @@ +"""Integration tests for the scheduler using MockTransport. + +These tests run the actual worker.py as subprocesses in isolated temp dirs, +with real torch training on CPU. No vast.ai instances needed. +""" + +from __future__ import annotations + +import asyncio +import os +import sys + +import pytest +import torch + +from sched.aggregator import average_weights +from sched.job import Job +from sched.scheduler import Node, Scheduler +from tests.mock_transport import MockConnection, 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")) + + +# ------------------------------------------------------------------ +# helpers +# ------------------------------------------------------------------ + + +def make_job(**overrides) -> Job: + defaults = dict( + script=MNIST_SCRIPT, + num_nodes=2, + rounds=3, + local_steps=20, + lr=0.01, + batch_size=32, + ) + defaults.update(overrides) + return Job(**defaults) + + +async def setup_scheduler( + job: Job, num_nodes: int = 2 +) -> tuple[Scheduler, MockTransport]: + """Create a scheduler with mock nodes, deploy code. Ready to run rounds.""" + transport = MockTransport() + sched = Scheduler(job, transport) + + # manually create nodes (skip provision) + 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") + ) + + # deploy worker + job module + await sched.deploy() + return sched, transport + + +# ------------------------------------------------------------------ +# tests +# ------------------------------------------------------------------ + + +@pytest.mark.asyncio +async def test_single_round_completes(): + """One round with 2 workers. Weights should change from initial.""" + job = make_job(rounds=1, local_steps=10) + sched, transport = await setup_scheduler(job) + + try: + job_mod = sched._load_job_module() + model = job_mod.make_model() + init_state = {k: v.clone() for k, v in model.state_dict().items()} + + new_state = await sched.run_round(model.state_dict(), round_num=0) + + # weights should have changed + changed = False + for key in init_state: + if not torch.equal(init_state[key], new_state[key]): + changed = True + break + assert changed, "Weights did not change after training" + finally: + transport.cleanup() + + +@pytest.mark.asyncio +async def test_multiple_rounds(): + """Multiple rounds. Verify the loop completes and returns valid state dict.""" + job = make_job(rounds=3, local_steps=15) + sched, transport = await setup_scheduler(job) + + try: + job_mod = sched._load_job_module() + model = job_mod.make_model() + state = model.state_dict() + + for r in range(job.rounds): + state = await sched.run_round(state, r) + + # final state should be loadable + model.load_state_dict(state) + # quick forward pass shouldn't crash + x = torch.randn(4, 1, 28, 28) + out = model(x) + assert out.shape == (4, 10) + finally: + transport.cleanup() + + +@pytest.mark.asyncio +async def test_worker_dies_midround(): + """One worker dies before a round. Scheduler continues with survivor.""" + job = make_job(rounds=2, local_steps=10) + sched, transport = await setup_scheduler(job) + + try: + job_mod = sched._load_job_module() + model = job_mod.make_model() + state = model.state_dict() + + # round 0 with both workers + state = await sched.run_round(state, 0) + assert len(sched._active_nodes()) == 2 + + # kill worker 1 + sched.nodes[1].conn.kill() + + # round 1 should still work with 1 survivor + state = await sched.run_round(state, 1) + + alive = sched._active_nodes() + assert len(alive) == 1 + assert alive[0].rank == 0 + + # state should be valid + model.load_state_dict(state) + finally: + transport.cleanup() + + +@pytest.mark.asyncio +async def test_all_workers_dead(): + """All workers dead → RuntimeError.""" + job = make_job(rounds=1, local_steps=5) + sched, transport = await setup_scheduler(job) + + try: + job_mod = sched._load_job_module() + state = job_mod.make_model().state_dict() + + for node in sched.nodes: + node.conn.kill() + + with pytest.raises(RuntimeError, match="No.*worker"): + await sched.run_round(state, 0) + finally: + transport.cleanup() + + +@pytest.mark.asyncio +async def test_worker_dies_after_push_before_exec(): + """Worker dies between param push and exec. Round completes with survivor.""" + job = make_job(rounds=1, local_steps=10) + sched, transport = await setup_scheduler(job) + + try: + job_mod = sched._load_job_module() + state = job_mod.make_model().state_dict() + + # wrap node 1's exec to kill it when called + node1_conn = sched.nodes[1].conn + original_exec = node1_conn.exec + + async def dying_exec(cmd, timeout=30.0): + node1_conn.kill() + raise ConnectionError("node is dead") + + node1_conn.exec = dying_exec + + # should still complete with node 0 + new_state = await sched.run_round(state, 0) + assert sched.nodes[1].status == "dead" + assert len([k for k, v in new_state.items()]) > 0 + finally: + transport.cleanup() + + +@pytest.mark.asyncio +async def test_three_nodes(): + """FedAvg with 3 workers.""" + job = make_job(num_nodes=3, rounds=2, local_steps=10) + sched, transport = await setup_scheduler(job, num_nodes=3) + + try: + job_mod = sched._load_job_module() + state = job_mod.make_model().state_dict() + + for r in range(job.rounds): + state = await sched.run_round(state, r) + + assert len(sched._active_nodes()) == 3 + finally: + transport.cleanup() + + +@pytest.mark.asyncio +async def test_deploy_copies_files(): + """Verify deploy actually puts worker.py and job_module.py on nodes.""" + job = make_job() + transport = MockTransport() + sched = Scheduler(job, transport) + + try: + for rank in range(2): + 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() + + for node in sched.nodes: + root = node.conn.root_dir + assert os.path.exists(os.path.join(root, "worker.py")) + assert os.path.exists(os.path.join(root, "job_module.py")) + finally: + transport.cleanup() + + +@pytest.mark.asyncio +async def test_convergence_signal(): + """Run several rounds and check that weight magnitude changes. + + Not a rigorous convergence test (synthetic data is random), but verifies + the training loop is actually doing gradient updates. + """ + job = make_job(rounds=5, local_steps=30) + sched, transport = await setup_scheduler(job) + + try: + job_mod = sched._load_job_module() + model = job_mod.make_model() + state = model.state_dict() + + weight_norms = [] + for r in range(job.rounds): + state = await sched.run_round(state, r) + norm = sum(v.float().norm().item() for v in state.values()) + weight_norms.append(norm) + + # weights should be changing across rounds (not frozen) + assert weight_norms[0] != pytest.approx(weight_norms[-1], abs=1e-3), ( + f"Weights did not change across rounds: {weight_norms}" + ) + finally: + transport.cleanup() diff --git a/tests/test_trace.py b/tests/test_trace.py new file mode 100644 index 0000000..0d734a3 --- /dev/null +++ b/tests/test_trace.py @@ -0,0 +1,265 @@ +"""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 "= args.local_steps: + break + x, y = x.to(device), y.to(device) + opt.zero_grad() + loss = F.cross_entropy(model(x), y) + loss.backward() + opt.step() + step += 1 + + torch.save(model.state_dict(), args.output) + print(f"rank={args.rank} done steps={args.local_steps}") + + +if __name__ == "__main__": + main()