feat: Scheduler with basic metrics, docker
Implementation of a mock scheduler for jobs/tasks that will evolve into a job manager for running training jobs over vast.ai instances.
This commit is contained in:
Zachery Aaron Shores-Chmielewski
parent
b0b04034d7
commit
7d11dfa688
48 changed files with 3187 additions and 0 deletions
12
.dockerignore
Normal file
12
.dockerignore
Normal file
|
|
@ -0,0 +1,12 @@
|
|||
.git
|
||||
.venv
|
||||
.pytest_cache
|
||||
__pycache__
|
||||
*.pyc
|
||||
.python-version
|
||||
uv.lock
|
||||
report.html
|
||||
dtrain
|
||||
DESIGN.md
|
||||
README.md
|
||||
.claude
|
||||
1
.python-version
Normal file
1
.python-version
Normal file
|
|
@ -0,0 +1 @@
|
|||
3.13
|
||||
27
docker/docker-compose.cpu.yml
Normal file
27
docker/docker-compose.cpu.yml
Normal file
|
|
@ -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:
|
||||
27
docker/docker-compose.yml
Normal file
27
docker/docker-compose.yml
Normal file
|
|
@ -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:
|
||||
15
docker/scheduler.Dockerfile
Normal file
15
docker/scheduler.Dockerfile
Normal file
|
|
@ -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"]
|
||||
18
docker/smoke-test.sh
Executable file
18
docker/smoke-test.sh
Executable file
|
|
@ -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"
|
||||
22
docker/worker-cpu.Dockerfile
Normal file
22
docker/worker-cpu.Dockerfile
Normal file
|
|
@ -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"]
|
||||
8
docker/worker-entrypoint.sh
Normal file
8
docker/worker-entrypoint.sh
Normal file
|
|
@ -0,0 +1,8 @@
|
|||
#!/bin/bash
|
||||
set -e
|
||||
|
||||
# Start SSH daemon
|
||||
/usr/sbin/sshd
|
||||
|
||||
# Keep container alive
|
||||
exec tail -f /dev/null
|
||||
20
docker/worker.Dockerfile
Normal file
20
docker/worker.Dockerfile
Normal file
|
|
@ -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"]
|
||||
0
jobs/__init__.py
Normal file
0
jobs/__init__.py
Normal file
BIN
jobs/__pycache__/mnist.cpython-313.pyc
Normal file
BIN
jobs/__pycache__/mnist.cpython-313.pyc
Normal file
Binary file not shown.
54
jobs/mnist.py
Normal file
54
jobs/mnist.py
Normal file
|
|
@ -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)
|
||||
6
main.py
Normal file
6
main.py
Normal file
|
|
@ -0,0 +1,6 @@
|
|||
def main():
|
||||
print("Hello from vastai-utils!")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
14
pyproject.toml
Normal file
14
pyproject.toml
Normal file
|
|
@ -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",
|
||||
]
|
||||
222
report.html
Normal file
222
report.html
Normal file
|
|
@ -0,0 +1,222 @@
|
|||
<!DOCTYPE html>
|
||||
<html lang="en">
|
||||
<head>
|
||||
<meta charset="UTF-8">
|
||||
<title>Training Run Report</title>
|
||||
<style>
|
||||
* { margin: 0; padding: 0; box-sizing: border-box; }
|
||||
body {
|
||||
background: #0f172a; color: #e2e8f0;
|
||||
font-family: 'SF Mono', 'Cascadia Code', 'Consolas', 'Menlo', monospace;
|
||||
padding: 32px; max-width: 960px; margin: 0 auto;
|
||||
}
|
||||
h1 { font-size: 20px; margin-bottom: 24px; }
|
||||
.stats {
|
||||
display: grid; grid-template-columns: repeat(auto-fit, minmax(140px, 1fr));
|
||||
gap: 12px; margin-bottom: 32px;
|
||||
}
|
||||
.stat {
|
||||
background: #1e293b; border-radius: 8px; padding: 16px;
|
||||
border: 1px solid #334155;
|
||||
}
|
||||
.stat-value { font-size: 22px; font-weight: bold; }
|
||||
.stat-label { font-size: 11px; color: #94a3b8; margin-top: 4px; }
|
||||
.legend {
|
||||
display: flex; gap: 20px; margin: 16px 0; flex-wrap: wrap;
|
||||
}
|
||||
.legend-item { display: flex; align-items: center; gap: 6px; font-size: 12px; }
|
||||
.legend-dot { width: 12px; height: 12px; border-radius: 3px; flex-shrink: 0; }
|
||||
.timeline { margin-top: 8px; }
|
||||
.timeline svg { width: 100%; height: auto; }
|
||||
.event-log {
|
||||
margin-top: 32px; background: #1e293b; border-radius: 8px;
|
||||
padding: 16px; border: 1px solid #334155;
|
||||
max-height: 400px; overflow-y: auto;
|
||||
}
|
||||
.event-log pre {
|
||||
font-size: 11px; line-height: 1.6; color: #94a3b8;
|
||||
white-space: pre-wrap;
|
||||
}
|
||||
.event-log .err { color: #ef4444; }
|
||||
</style>
|
||||
</head>
|
||||
<body>
|
||||
<h1>Training Run Report</h1>
|
||||
|
||||
<div class="stats">
|
||||
<div class="stat">
|
||||
<div class="stat-value">3</div>
|
||||
<div class="stat-label">Rounds</div>
|
||||
</div>
|
||||
<div class="stat">
|
||||
<div class="stat-value">2</div>
|
||||
<div class="stat-label">Workers</div>
|
||||
</div>
|
||||
<div class="stat">
|
||||
<div class="stat-value">15.6s</div>
|
||||
<div class="stat-label">Total Time</div>
|
||||
</div>
|
||||
<div class="stat">
|
||||
<div class="stat-value">4.7MB</div>
|
||||
<div class="stat-label">Data Transferred</div>
|
||||
</div>
|
||||
<div class="stat">
|
||||
<div class="stat-value">2.3MB</div>
|
||||
<div class="stat-label">Params Pushed</div>
|
||||
</div>
|
||||
<div class="stat">
|
||||
<div class="stat-value">2.3MB</div>
|
||||
<div class="stat-label">Weights Pulled</div>
|
||||
</div>
|
||||
<div class="stat">
|
||||
<div class="stat-value">8.7</div>
|
||||
<div class="stat-label">Final |W|</div>
|
||||
</div>
|
||||
</div>
|
||||
|
||||
<div class="legend">
|
||||
<div class="legend-item"><div class="legend-dot" style="background:#3b82f6"></div>Push params</div>
|
||||
<div class="legend-item"><div class="legend-dot" style="background:#22c55e"></div>Train</div>
|
||||
<div class="legend-item"><div class="legend-dot" style="background:#f97316"></div>Pull weights</div>
|
||||
<div class="legend-item"><div class="legend-dot" style="background:#8b5cf6"></div>Aggregate</div>
|
||||
<div class="legend-item"><div class="legend-dot" style="background:#ef4444"></div>Dead / Error</div>
|
||||
</div>
|
||||
|
||||
<div class="timeline">
|
||||
<svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 820 466" width="820" height="466">
|
||||
<rect width="820" height="466" fill="#1e293b" rx="8"/>
|
||||
<text x="12" y="26" font-size="13" font-family="monospace" fill="#e2e8f0" font-weight="bold">round 0 [2/2] 4.59s</text>
|
||||
<line x1="100.0" y1="34" x2="100.0" y2="130" stroke="#334155" stroke-width="0.5" stroke-dasharray="4 4"/>
|
||||
<text x="100.0" y="142" font-size="10" font-family="monospace" fill="#94a3b8" text-anchor="middle">0.0s</text>
|
||||
<line x1="240.0" y1="34" x2="240.0" y2="130" stroke="#334155" stroke-width="0.5" stroke-dasharray="4 4"/>
|
||||
<text x="240.0" y="142" font-size="10" font-family="monospace" fill="#94a3b8" text-anchor="middle">0.9s</text>
|
||||
<line x1="380.0" y1="34" x2="380.0" y2="130" stroke="#334155" stroke-width="0.5" stroke-dasharray="4 4"/>
|
||||
<text x="380.0" y="142" font-size="10" font-family="monospace" fill="#94a3b8" text-anchor="middle">1.8s</text>
|
||||
<line x1="520.0" y1="34" x2="520.0" y2="130" stroke="#334155" stroke-width="0.5" stroke-dasharray="4 4"/>
|
||||
<text x="520.0" y="142" font-size="10" font-family="monospace" fill="#94a3b8" text-anchor="middle">2.8s</text>
|
||||
<line x1="660.0" y1="34" x2="660.0" y2="130" stroke="#334155" stroke-width="0.5" stroke-dasharray="4 4"/>
|
||||
<text x="660.0" y="142" font-size="10" font-family="monospace" fill="#94a3b8" text-anchor="middle">3.7s</text>
|
||||
<line x1="800.0" y1="34" x2="800.0" y2="130" stroke="#334155" stroke-width="0.5" stroke-dasharray="4 4"/>
|
||||
<text x="800.0" y="142" font-size="10" font-family="monospace" fill="#94a3b8" text-anchor="middle">4.6s</text>
|
||||
<text x="12" y="52" font-size="11" font-family="monospace" fill="#94a3b8">scheduler</text>
|
||||
<rect x="100" y="38" width="700" height="20" fill="#0f172a" rx="3"/>
|
||||
<rect x="100.00781517173503" y="40" width="3" height="16" fill="#3b82f6" rx="2" opacity="0.9"><title>push 399.7KB → node 0 (0.000s)</title></rect>
|
||||
<rect x="100.0322995388102" y="40" width="3" height="16" fill="#3b82f6" rx="2" opacity="0.9"><title>push 399.7KB → node 1 (0.000s)</title></rect>
|
||||
<rect x="799.338565795855" y="40" width="3" height="16" fill="#f97316" rx="2" opacity="0.9"><title>pull 399.7KB ← node 0 (0.000s)</title></rect>
|
||||
<rect x="799.5730082941284" y="40" width="3" height="16" fill="#f97316" rx="2" opacity="0.9"><title>pull 399.7KB ← node 1 (0.000s)</title></rect>
|
||||
<rect x="799.7174432972781" y="40" width="3" height="16" fill="#8b5cf6" rx="2" opacity="0.9"><title>aggregate: 2 workers, |W|=8.75, Δ=0.0891</title></rect>
|
||||
<text x="12" y="84" font-size="11" font-family="monospace" fill="#94a3b8">worker 0</text>
|
||||
<rect x="100" y="70" width="700" height="20" fill="#0f172a" rx="3"/>
|
||||
<rect x="100.00781517173503" y="72" width="3" height="16" fill="#3b82f6" rx="2" opacity="0.9"><title>push 399.7KB → node 0 (0.000s)</title></rect>
|
||||
<rect x="100.05669380279495" y="72" width="699.27049173338" height="16" fill="#22c55e" rx="2" opacity="0.9"><title>train node 0: 4.59s exit=0
|
||||
cpu = _conversion_method_template(device=torch.device("cpu"))</title></rect>
|
||||
<rect x="799.338565795855" y="72" width="3" height="16" fill="#f97316" rx="2" opacity="0.9"><title>pull 399.7KB ← node 0 (0.000s)</title></rect>
|
||||
<text x="12" y="116" font-size="11" font-family="monospace" fill="#94a3b8">worker 1</text>
|
||||
<rect x="100" y="102" width="700" height="20" fill="#0f172a" rx="3"/>
|
||||
<rect x="100.0322995388102" y="104" width="3" height="16" fill="#3b82f6" rx="2" opacity="0.9"><title>push 399.7KB → node 1 (0.000s)</title></rect>
|
||||
<rect x="100.25768249946535" y="104" width="657.6046307405937" height="16" fill="#22c55e" rx="2" opacity="0.9"><title>train node 1: 4.31s exit=0
|
||||
cpu = _conversion_method_template(device=torch.device("cpu"))</title></rect>
|
||||
<rect x="799.5730082941284" y="104" width="3" height="16" fill="#f97316" rx="2" opacity="0.9"><title>pull 399.7KB ← node 1 (0.000s)</title></rect>
|
||||
<text x="12" y="168" font-size="13" font-family="monospace" fill="#e2e8f0" font-weight="bold">round 1 [2/2] 5.53s</text>
|
||||
<line x1="100.0" y1="176" x2="100.0" y2="272" stroke="#334155" stroke-width="0.5" stroke-dasharray="4 4"/>
|
||||
<text x="100.0" y="284" font-size="10" font-family="monospace" fill="#94a3b8" text-anchor="middle">0.0s</text>
|
||||
<line x1="240.0" y1="176" x2="240.0" y2="272" stroke="#334155" stroke-width="0.5" stroke-dasharray="4 4"/>
|
||||
<text x="240.0" y="284" font-size="10" font-family="monospace" fill="#94a3b8" text-anchor="middle">1.1s</text>
|
||||
<line x1="380.0" y1="176" x2="380.0" y2="272" stroke="#334155" stroke-width="0.5" stroke-dasharray="4 4"/>
|
||||
<text x="380.0" y="284" font-size="10" font-family="monospace" fill="#94a3b8" text-anchor="middle">2.2s</text>
|
||||
<line x1="520.0" y1="176" x2="520.0" y2="272" stroke="#334155" stroke-width="0.5" stroke-dasharray="4 4"/>
|
||||
<text x="520.0" y="284" font-size="10" font-family="monospace" fill="#94a3b8" text-anchor="middle">3.3s</text>
|
||||
<line x1="660.0" y1="176" x2="660.0" y2="272" stroke="#334155" stroke-width="0.5" stroke-dasharray="4 4"/>
|
||||
<text x="660.0" y="284" font-size="10" font-family="monospace" fill="#94a3b8" text-anchor="middle">4.4s</text>
|
||||
<line x1="800.0" y1="176" x2="800.0" y2="272" stroke="#334155" stroke-width="0.5" stroke-dasharray="4 4"/>
|
||||
<text x="800.0" y="284" font-size="10" font-family="monospace" fill="#94a3b8" text-anchor="middle">5.5s</text>
|
||||
<text x="12" y="194" font-size="11" font-family="monospace" fill="#94a3b8">scheduler</text>
|
||||
<rect x="100" y="180" width="700" height="20" fill="#0f172a" rx="3"/>
|
||||
<rect x="100.00834405549375" y="182" width="3" height="16" fill="#3b82f6" rx="2" opacity="0.9"><title>push 399.6KB → node 0 (0.000s)</title></rect>
|
||||
<rect x="100.03682653927218" y="182" width="3" height="16" fill="#3b82f6" rx="2" opacity="0.9"><title>push 399.6KB → node 1 (0.002s)</title></rect>
|
||||
<rect x="799.4986844978123" y="182" width="3" height="16" fill="#f97316" rx="2" opacity="0.9"><title>pull 399.7KB ← node 0 (0.000s)</title></rect>
|
||||
<rect x="799.6683723520547" y="182" width="3" height="16" fill="#f97316" rx="2" opacity="0.9"><title>pull 399.7KB ← node 1 (0.000s)</title></rect>
|
||||
<rect x="799.8614614220614" y="182" width="3" height="16" fill="#8b5cf6" rx="2" opacity="0.9"><title>aggregate: 2 workers, |W|=8.74, Δ=0.0876</title></rect>
|
||||
<text x="12" y="226" font-size="11" font-family="monospace" fill="#94a3b8">worker 0</text>
|
||||
<rect x="100" y="212" width="700" height="20" fill="#0f172a" rx="3"/>
|
||||
<rect x="100.00834405549375" y="214" width="3" height="16" fill="#3b82f6" rx="2" opacity="0.9"><title>push 399.6KB → node 0 (0.000s)</title></rect>
|
||||
<rect x="100.30709872380446" y="214" width="676.9476932782056" height="16" fill="#22c55e" rx="2" opacity="0.9"><title>train node 0: 5.35s exit=0
|
||||
cpu = _conversion_method_template(device=torch.device("cpu"))</title></rect>
|
||||
<rect x="799.4986844978123" y="214" width="3" height="16" fill="#f97316" rx="2" opacity="0.9"><title>pull 399.7KB ← node 0 (0.000s)</title></rect>
|
||||
<text x="12" y="258" font-size="11" font-family="monospace" fill="#94a3b8">worker 1</text>
|
||||
<rect x="100" y="244" width="700" height="20" fill="#0f172a" rx="3"/>
|
||||
<rect x="100.03682653927218" y="246" width="3" height="16" fill="#3b82f6" rx="2" opacity="0.9"><title>push 399.6KB → node 1 (0.002s)</title></rect>
|
||||
<rect x="100.3581074116179" y="246" width="699.1320361372128" height="16" fill="#22c55e" rx="2" opacity="0.9"><title>train node 1: 5.53s exit=0
|
||||
cpu = _conversion_method_template(device=torch.device("cpu"))</title></rect>
|
||||
<rect x="799.6683723520547" y="246" width="3" height="16" fill="#f97316" rx="2" opacity="0.9"><title>pull 399.7KB ← node 1 (0.000s)</title></rect>
|
||||
<text x="12" y="310" font-size="13" font-family="monospace" fill="#e2e8f0" font-weight="bold">round 2 [2/2] 5.50s</text>
|
||||
<line x1="100.0" y1="318" x2="100.0" y2="414" stroke="#334155" stroke-width="0.5" stroke-dasharray="4 4"/>
|
||||
<text x="100.0" y="426" font-size="10" font-family="monospace" fill="#94a3b8" text-anchor="middle">0.0s</text>
|
||||
<line x1="240.0" y1="318" x2="240.0" y2="414" stroke="#334155" stroke-width="0.5" stroke-dasharray="4 4"/>
|
||||
<text x="240.0" y="426" font-size="10" font-family="monospace" fill="#94a3b8" text-anchor="middle">1.1s</text>
|
||||
<line x1="380.0" y1="318" x2="380.0" y2="414" stroke="#334155" stroke-width="0.5" stroke-dasharray="4 4"/>
|
||||
<text x="380.0" y="426" font-size="10" font-family="monospace" fill="#94a3b8" text-anchor="middle">2.2s</text>
|
||||
<line x1="520.0" y1="318" x2="520.0" y2="414" stroke="#334155" stroke-width="0.5" stroke-dasharray="4 4"/>
|
||||
<text x="520.0" y="426" font-size="10" font-family="monospace" fill="#94a3b8" text-anchor="middle">3.3s</text>
|
||||
<line x1="660.0" y1="318" x2="660.0" y2="414" stroke="#334155" stroke-width="0.5" stroke-dasharray="4 4"/>
|
||||
<text x="660.0" y="426" font-size="10" font-family="monospace" fill="#94a3b8" text-anchor="middle">4.4s</text>
|
||||
<line x1="800.0" y1="318" x2="800.0" y2="414" stroke="#334155" stroke-width="0.5" stroke-dasharray="4 4"/>
|
||||
<text x="800.0" y="426" font-size="10" font-family="monospace" fill="#94a3b8" text-anchor="middle">5.5s</text>
|
||||
<text x="12" y="336" font-size="11" font-family="monospace" fill="#94a3b8">scheduler</text>
|
||||
<rect x="100" y="322" width="700" height="20" fill="#0f172a" rx="3"/>
|
||||
<rect x="100.02026148389356" y="324" width="3" height="16" fill="#3b82f6" rx="2" opacity="0.9"><title>push 399.6KB → node 0 (0.000s)</title></rect>
|
||||
<rect x="100.05808530320598" y="324" width="3" height="16" fill="#3b82f6" rx="2" opacity="0.9"><title>push 399.6KB → node 1 (0.000s)</title></rect>
|
||||
<rect x="799.363116457372" y="324" width="3" height="16" fill="#f97316" rx="2" opacity="0.9"><title>pull 399.7KB ← node 0 (0.000s)</title></rect>
|
||||
<rect x="799.5569707971623" y="324" width="3" height="16" fill="#f97316" rx="2" opacity="0.9"><title>pull 399.7KB ← node 1 (0.000s)</title></rect>
|
||||
<rect x="799.7684761253557" y="324" width="3" height="16" fill="#8b5cf6" rx="2" opacity="0.9"><title>aggregate: 2 workers, |W|=8.74, Δ=0.0883</title></rect>
|
||||
<text x="12" y="368" font-size="11" font-family="monospace" fill="#94a3b8">worker 0</text>
|
||||
<rect x="100" y="354" width="700" height="20" fill="#0f172a" rx="3"/>
|
||||
<rect x="100.02026148389356" y="356" width="3" height="16" fill="#3b82f6" rx="2" opacity="0.9"><title>push 399.6KB → node 0 (0.000s)</title></rect>
|
||||
<rect x="100.09750932913751" y="356" width="699.2563629168552" height="16" fill="#22c55e" rx="2" opacity="0.9"><title>train node 0: 5.50s exit=0
|
||||
cpu = _conversion_method_template(device=torch.device("cpu"))</title></rect>
|
||||
<rect x="799.363116457372" y="356" width="3" height="16" fill="#f97316" rx="2" opacity="0.9"><title>pull 399.7KB ← node 0 (0.000s)</title></rect>
|
||||
<text x="12" y="400" font-size="11" font-family="monospace" fill="#94a3b8">worker 1</text>
|
||||
<rect x="100" y="386" width="700" height="20" fill="#0f172a" rx="3"/>
|
||||
<rect x="100.05808530320598" y="388" width="3" height="16" fill="#3b82f6" rx="2" opacity="0.9"><title>push 399.6KB → node 1 (0.000s)</title></rect>
|
||||
<rect x="100.15763886743558" y="388" width="684.0486400269059" height="16" fill="#22c55e" rx="2" opacity="0.9"><title>train node 1: 5.38s exit=0
|
||||
cpu = _conversion_method_template(device=torch.device("cpu"))</title></rect>
|
||||
<rect x="799.5569707971623" y="388" width="3" height="16" fill="#f97316" rx="2" opacity="0.9"><title>pull 399.7KB ← node 1 (0.000s)</title></rect>
|
||||
</svg>
|
||||
</div>
|
||||
|
||||
<div class="event-log">
|
||||
<pre>[ 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) ────────</pre>
|
||||
</div>
|
||||
|
||||
</body>
|
||||
</html>
|
||||
31
reports/log.txt
Normal file
31
reports/log.txt
Normal file
|
|
@ -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
|
||||
140
reports/report.html
Normal file
140
reports/report.html
Normal file
|
|
@ -0,0 +1,140 @@
|
|||
<!DOCTYPE html>
|
||||
<html lang="en">
|
||||
<head>
|
||||
<meta charset="UTF-8">
|
||||
<title>Training Run Report</title>
|
||||
<style>
|
||||
* { margin: 0; padding: 0; box-sizing: border-box; }
|
||||
body {
|
||||
background: #0f172a; color: #e2e8f0;
|
||||
font-family: 'SF Mono', 'Cascadia Code', 'Consolas', 'Menlo', monospace;
|
||||
padding: 32px; max-width: 960px; margin: 0 auto;
|
||||
}
|
||||
h1 { font-size: 20px; margin-bottom: 24px; }
|
||||
.stats {
|
||||
display: grid; grid-template-columns: repeat(auto-fit, minmax(140px, 1fr));
|
||||
gap: 12px; margin-bottom: 32px;
|
||||
}
|
||||
.stat {
|
||||
background: #1e293b; border-radius: 8px; padding: 16px;
|
||||
border: 1px solid #334155;
|
||||
}
|
||||
.stat-value { font-size: 22px; font-weight: bold; }
|
||||
.stat-label { font-size: 11px; color: #94a3b8; margin-top: 4px; }
|
||||
.legend {
|
||||
display: flex; gap: 20px; margin: 16px 0; flex-wrap: wrap;
|
||||
}
|
||||
.legend-item { display: flex; align-items: center; gap: 6px; font-size: 12px; }
|
||||
.legend-dot { width: 12px; height: 12px; border-radius: 3px; flex-shrink: 0; }
|
||||
.timeline { margin-top: 8px; }
|
||||
.timeline svg { width: 100%; height: auto; }
|
||||
.event-log {
|
||||
margin-top: 32px; background: #1e293b; border-radius: 8px;
|
||||
padding: 16px; border: 1px solid #334155;
|
||||
max-height: 400px; overflow-y: auto;
|
||||
}
|
||||
.event-log pre {
|
||||
font-size: 11px; line-height: 1.6; color: #94a3b8;
|
||||
white-space: pre-wrap;
|
||||
}
|
||||
.event-log .err { color: #ef4444; }
|
||||
</style>
|
||||
</head>
|
||||
<body>
|
||||
<h1>Training Run Report</h1>
|
||||
|
||||
<div class="stats">
|
||||
<div class="stat">
|
||||
<div class="stat-value">1</div>
|
||||
<div class="stat-label">Rounds</div>
|
||||
</div>
|
||||
<div class="stat">
|
||||
<div class="stat-value">2</div>
|
||||
<div class="stat-label">Workers</div>
|
||||
</div>
|
||||
<div class="stat">
|
||||
<div class="stat-value">9.9s</div>
|
||||
<div class="stat-label">Total Time</div>
|
||||
</div>
|
||||
<div class="stat">
|
||||
<div class="stat-value">1.6MB</div>
|
||||
<div class="stat-label">Data Transferred</div>
|
||||
</div>
|
||||
<div class="stat">
|
||||
<div class="stat-value">799.4KB</div>
|
||||
<div class="stat-label">Params Pushed</div>
|
||||
</div>
|
||||
<div class="stat">
|
||||
<div class="stat-value">799.4KB</div>
|
||||
<div class="stat-label">Weights Pulled</div>
|
||||
</div>
|
||||
<div class="stat">
|
||||
<div class="stat-value">8.7</div>
|
||||
<div class="stat-label">Final |W|</div>
|
||||
</div>
|
||||
</div>
|
||||
|
||||
<div class="legend">
|
||||
<div class="legend-item"><div class="legend-dot" style="background:#3b82f6"></div>Push params</div>
|
||||
<div class="legend-item"><div class="legend-dot" style="background:#22c55e"></div>Train</div>
|
||||
<div class="legend-item"><div class="legend-dot" style="background:#f97316"></div>Pull weights</div>
|
||||
<div class="legend-item"><div class="legend-dot" style="background:#8b5cf6"></div>Aggregate</div>
|
||||
<div class="legend-item"><div class="legend-dot" style="background:#ef4444"></div>Dead / Error</div>
|
||||
</div>
|
||||
|
||||
<div class="timeline">
|
||||
<svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 820 182" width="820" height="182">
|
||||
<rect width="820" height="182" fill="#1e293b" rx="8"/>
|
||||
<text x="12" y="26" font-size="13" font-family="monospace" fill="#e2e8f0" font-weight="bold">round 0 [2/2] 9.91s</text>
|
||||
<line x1="100.0" y1="34" x2="100.0" y2="130" stroke="#334155" stroke-width="0.5" stroke-dasharray="4 4"/>
|
||||
<text x="100.0" y="142" font-size="10" font-family="monospace" fill="#94a3b8" text-anchor="middle">0.0s</text>
|
||||
<line x1="240.0" y1="34" x2="240.0" y2="130" stroke="#334155" stroke-width="0.5" stroke-dasharray="4 4"/>
|
||||
<text x="240.0" y="142" font-size="10" font-family="monospace" fill="#94a3b8" text-anchor="middle">2.0s</text>
|
||||
<line x1="380.0" y1="34" x2="380.0" y2="130" stroke="#334155" stroke-width="0.5" stroke-dasharray="4 4"/>
|
||||
<text x="380.0" y="142" font-size="10" font-family="monospace" fill="#94a3b8" text-anchor="middle">4.0s</text>
|
||||
<line x1="520.0" y1="34" x2="520.0" y2="130" stroke="#334155" stroke-width="0.5" stroke-dasharray="4 4"/>
|
||||
<text x="520.0" y="142" font-size="10" font-family="monospace" fill="#94a3b8" text-anchor="middle">5.9s</text>
|
||||
<line x1="660.0" y1="34" x2="660.0" y2="130" stroke="#334155" stroke-width="0.5" stroke-dasharray="4 4"/>
|
||||
<text x="660.0" y="142" font-size="10" font-family="monospace" fill="#94a3b8" text-anchor="middle">7.9s</text>
|
||||
<line x1="800.0" y1="34" x2="800.0" y2="130" stroke="#334155" stroke-width="0.5" stroke-dasharray="4 4"/>
|
||||
<text x="800.0" y="142" font-size="10" font-family="monospace" fill="#94a3b8" text-anchor="middle">9.9s</text>
|
||||
<text x="12" y="52" font-size="11" font-family="monospace" fill="#94a3b8">scheduler</text>
|
||||
<rect x="100" y="38" width="700" height="20" fill="#0f172a" rx="3"/>
|
||||
<rect x="100.00695051486753" y="40" width="10.269547854213515" height="16" fill="#3b82f6" rx="2" opacity="0.9"><title>push 399.7KB → node 0 (0.145s)</title></rect>
|
||||
<rect x="100.03894679062752" y="40" width="10.298896778860655" height="16" fill="#3b82f6" rx="2" opacity="0.9"><title>push 399.7KB → node 1 (0.146s)</title></rect>
|
||||
<rect x="775.4392780576894" y="40" width="12.318390925589076" height="16" fill="#f97316" rx="2" opacity="0.9"><title>pull 399.7KB ← node 0 (0.174s)</title></rect>
|
||||
<rect x="787.9259294950273" y="40" width="11.67198867050213" height="16" fill="#f97316" rx="2" opacity="0.9"><title>pull 399.7KB ← node 1 (0.165s)</title></rect>
|
||||
<rect x="799.7428815775497" y="40" width="3" height="16" fill="#8b5cf6" rx="2" opacity="0.9"><title>aggregate: 2 workers, |W|=8.72, Δ=0.3177</title></rect>
|
||||
<text x="12" y="84" font-size="11" font-family="monospace" fill="#94a3b8">worker 0</text>
|
||||
<rect x="100" y="70" width="700" height="20" fill="#0f172a" rx="3"/>
|
||||
<rect x="100.00695051486753" y="72" width="10.269547854213515" height="16" fill="#3b82f6" rx="2" opacity="0.9"><title>push 399.7KB → node 0 (0.145s)</title></rect>
|
||||
<rect x="110.35350371003277" y="72" width="665.0670345270505" height="16" fill="#22c55e" rx="2" opacity="0.9"><title>train node 0: 9.42s exit=0
|
||||
cpu = _conversion_method_template(device=torch.device("cpu"))</title></rect>
|
||||
<rect x="775.4392780576894" y="72" width="12.318390925589076" height="16" fill="#f97316" rx="2" opacity="0.9"><title>pull 399.7KB ← node 0 (0.174s)</title></rect>
|
||||
<text x="12" y="116" font-size="11" font-family="monospace" fill="#94a3b8">worker 1</text>
|
||||
<rect x="100" y="102" width="700" height="20" fill="#0f172a" rx="3"/>
|
||||
<rect x="100.03894679062752" y="104" width="10.298896778860655" height="16" fill="#3b82f6" rx="2" opacity="0.9"><title>push 399.7KB → node 1 (0.146s)</title></rect>
|
||||
<rect x="110.39965258944382" y="104" width="664.8660109438434" height="16" fill="#22c55e" rx="2" opacity="0.9"><title>train node 1: 9.42s exit=0
|
||||
cpu = _conversion_method_template(device=torch.device("cpu"))</title></rect>
|
||||
<rect x="787.9259294950273" y="104" width="11.67198867050213" height="16" fill="#f97316" rx="2" opacity="0.9"><title>pull 399.7KB ← node 1 (0.165s)</title></rect>
|
||||
</svg>
|
||||
</div>
|
||||
|
||||
<div class="event-log">
|
||||
<pre>[ 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) ────────</pre>
|
||||
</div>
|
||||
|
||||
</body>
|
||||
</html>
|
||||
0
sched/__init__.py
Normal file
0
sched/__init__.py
Normal file
122
sched/__main__.py
Normal file
122
sched/__main__.py
Normal file
|
|
@ -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())
|
||||
BIN
sched/__pycache__/__init__.cpython-313.pyc
Normal file
BIN
sched/__pycache__/__init__.cpython-313.pyc
Normal file
Binary file not shown.
BIN
sched/__pycache__/__main__.cpython-313.pyc
Normal file
BIN
sched/__pycache__/__main__.cpython-313.pyc
Normal file
Binary file not shown.
BIN
sched/__pycache__/aggregator.cpython-313.pyc
Normal file
BIN
sched/__pycache__/aggregator.cpython-313.pyc
Normal file
Binary file not shown.
BIN
sched/__pycache__/job.cpython-313.pyc
Normal file
BIN
sched/__pycache__/job.cpython-313.pyc
Normal file
Binary file not shown.
BIN
sched/__pycache__/report.cpython-313.pyc
Normal file
BIN
sched/__pycache__/report.cpython-313.pyc
Normal file
Binary file not shown.
BIN
sched/__pycache__/scheduler.cpython-313.pyc
Normal file
BIN
sched/__pycache__/scheduler.cpython-313.pyc
Normal file
Binary file not shown.
BIN
sched/__pycache__/trace.cpython-313.pyc
Normal file
BIN
sched/__pycache__/trace.cpython-313.pyc
Normal file
Binary file not shown.
BIN
sched/__pycache__/transport.cpython-313.pyc
Normal file
BIN
sched/__pycache__/transport.cpython-313.pyc
Normal file
Binary file not shown.
BIN
sched/__pycache__/vastai.cpython-313.pyc
Normal file
BIN
sched/__pycache__/vastai.cpython-313.pyc
Normal file
Binary file not shown.
18
sched/aggregator.py
Normal file
18
sched/aggregator.py
Normal file
|
|
@ -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
|
||||
17
sched/job.py
Normal file
17
sched/job.py
Normal file
|
|
@ -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)
|
||||
403
sched/report.py
Normal file
403
sched/report.py
Normal file
|
|
@ -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 '<svg xmlns="http://www.w3.org/2000/svg"></svg>'
|
||||
|
||||
# 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'<svg xmlns="http://www.w3.org/2000/svg" '
|
||||
f'viewBox="0 0 {total_w} {total_h}" '
|
||||
f'width="{total_w}" height="{total_h}">'
|
||||
)
|
||||
parts.append(
|
||||
f'<rect width="{total_w}" height="{total_h}" fill="{CARD}" rx="8"/>'
|
||||
)
|
||||
|
||||
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'<text x="12" y="{y_cursor + 14}" '
|
||||
f'font-size="13" font-family="monospace" fill="{TEXT}" font-weight="bold">'
|
||||
f"round {rnum} [{surv}/{total_nodes}] {dur:.2f}s</text>"
|
||||
)
|
||||
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'<line x1="{gx}" y1="{y_cursor}" '
|
||||
f'x2="{gx}" y2="{y_cursor + rows_per_round * (row_h + row_gap)}" '
|
||||
f'stroke="{GRID}" stroke-width="0.5" stroke-dasharray="4 4"/>'
|
||||
)
|
||||
t_label = frac * dur
|
||||
parts.append(
|
||||
f'<text x="{gx}" y="{y_cursor + rows_per_round * (row_h + row_gap) + 12}" '
|
||||
f'font-size="10" font-family="monospace" fill="{MUTED}" text-anchor="middle">'
|
||||
f"{t_label:.1f}s</text>"
|
||||
)
|
||||
|
||||
# --- scheduler row ---
|
||||
row_y = y_cursor
|
||||
parts.append(
|
||||
f'<text x="12" y="{row_y + 18}" '
|
||||
f'font-size="11" font-family="monospace" fill="{MUTED}">scheduler</text>'
|
||||
)
|
||||
# background bar
|
||||
parts.append(
|
||||
f'<rect x="{label_w}" y="{row_y + 4}" width="{chart_w}" height="{row_h - 8}" '
|
||||
f'fill="{BG}" rx="3"/>'
|
||||
)
|
||||
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'<rect x="{bx}" y="{row_y + 6}" width="{bw}" height="{row_h - 12}" '
|
||||
f'fill="{color}" rx="2" opacity="0.9">'
|
||||
f"<title>{html.escape(tooltip)}</title></rect>"
|
||||
)
|
||||
y_cursor += row_h + row_gap
|
||||
|
||||
# --- worker rows ---
|
||||
for rank in range(max_rank + 1):
|
||||
row_y = y_cursor
|
||||
parts.append(
|
||||
f'<text x="12" y="{row_y + 18}" '
|
||||
f'font-size="11" font-family="monospace" fill="{MUTED}">worker {rank}</text>'
|
||||
)
|
||||
parts.append(
|
||||
f'<rect x="{label_w}" y="{row_y + 4}" width="{chart_w}" height="{row_h - 8}" '
|
||||
f'fill="{BG}" rx="3"/>'
|
||||
)
|
||||
|
||||
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'<rect x="{bx}" y="{row_y + 6}" width="{bw}" height="{row_h - 12}" '
|
||||
f'fill="{color}" rx="2" opacity="0.9">'
|
||||
f"<title>{html.escape(tooltip)}</title></rect>"
|
||||
)
|
||||
|
||||
# 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'<rect x="{dead_x}" y="{row_y + 6}" '
|
||||
f'width="{dw}" height="{row_h - 12}" '
|
||||
f'fill="{COLORS["dead"]}" rx="2" opacity="0.25"/>'
|
||||
)
|
||||
|
||||
# no events at all — full dead bar
|
||||
if not rank_evs and rank < rd.get("num_workers", 0):
|
||||
parts.append(
|
||||
f'<rect x="{label_w}" y="{row_y + 6}" '
|
||||
f'width="{chart_w}" height="{row_h - 12}" '
|
||||
f'fill="{COLORS["dead"]}" rx="2" opacity="0.3">'
|
||||
f"<title>worker {rank}: no response</title></rect>"
|
||||
)
|
||||
|
||||
y_cursor += row_h + row_gap
|
||||
|
||||
y_cursor += round_gap
|
||||
|
||||
parts.append("</svg>")
|
||||
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"""<!DOCTYPE html>
|
||||
<html lang="en">
|
||||
<head>
|
||||
<meta charset="UTF-8">
|
||||
<title>Training Run Report</title>
|
||||
<style>
|
||||
* {{ margin: 0; padding: 0; box-sizing: border-box; }}
|
||||
body {{
|
||||
background: {BG}; color: {TEXT};
|
||||
font-family: 'SF Mono', 'Cascadia Code', 'Consolas', 'Menlo', monospace;
|
||||
padding: 32px; max-width: 960px; margin: 0 auto;
|
||||
}}
|
||||
h1 {{ font-size: 20px; margin-bottom: 24px; }}
|
||||
.stats {{
|
||||
display: grid; grid-template-columns: repeat(auto-fit, minmax(140px, 1fr));
|
||||
gap: 12px; margin-bottom: 32px;
|
||||
}}
|
||||
.stat {{
|
||||
background: {CARD}; border-radius: 8px; padding: 16px;
|
||||
border: 1px solid {GRID};
|
||||
}}
|
||||
.stat-value {{ font-size: 22px; font-weight: bold; }}
|
||||
.stat-label {{ font-size: 11px; color: {MUTED}; margin-top: 4px; }}
|
||||
.legend {{
|
||||
display: flex; gap: 20px; margin: 16px 0; flex-wrap: wrap;
|
||||
}}
|
||||
.legend-item {{ display: flex; align-items: center; gap: 6px; font-size: 12px; }}
|
||||
.legend-dot {{ width: 12px; height: 12px; border-radius: 3px; flex-shrink: 0; }}
|
||||
.timeline {{ margin-top: 8px; }}
|
||||
.timeline svg {{ width: 100%; height: auto; }}
|
||||
.event-log {{
|
||||
margin-top: 32px; background: {CARD}; border-radius: 8px;
|
||||
padding: 16px; border: 1px solid {GRID};
|
||||
max-height: 400px; overflow-y: auto;
|
||||
}}
|
||||
.event-log pre {{
|
||||
font-size: 11px; line-height: 1.6; color: {MUTED};
|
||||
white-space: pre-wrap;
|
||||
}}
|
||||
.event-log .err {{ color: {COLORS["dead"]}; }}
|
||||
</style>
|
||||
</head>
|
||||
<body>
|
||||
<h1>Training Run Report</h1>
|
||||
|
||||
<div class="stats">
|
||||
<div class="stat">
|
||||
<div class="stat-value">{n_rounds}</div>
|
||||
<div class="stat-label">Rounds</div>
|
||||
</div>
|
||||
<div class="stat">
|
||||
<div class="stat-value">{max_workers}</div>
|
||||
<div class="stat-label">Workers</div>
|
||||
</div>
|
||||
<div class="stat">
|
||||
<div class="stat-value">{total_time:.1f}s</div>
|
||||
<div class="stat-label">Total Time</div>
|
||||
</div>
|
||||
<div class="stat">
|
||||
<div class="stat-value">{_fmt_bytes(total_push + total_pull)}</div>
|
||||
<div class="stat-label">Data Transferred</div>
|
||||
</div>
|
||||
<div class="stat">
|
||||
<div class="stat-value">{_fmt_bytes(total_push)}</div>
|
||||
<div class="stat-label">Params Pushed</div>
|
||||
</div>
|
||||
<div class="stat">
|
||||
<div class="stat-value">{_fmt_bytes(total_pull)}</div>
|
||||
<div class="stat-label">Weights Pulled</div>
|
||||
</div>
|
||||
<div class="stat">
|
||||
<div class="stat-value">{final_wnorm:.1f}</div>
|
||||
<div class="stat-label">Final |W|</div>
|
||||
</div>
|
||||
</div>
|
||||
|
||||
<div class="legend">
|
||||
<div class="legend-item"><div class="legend-dot" style="background:{COLORS['push']}"></div>Push params</div>
|
||||
<div class="legend-item"><div class="legend-dot" style="background:{COLORS['exec']}"></div>Train</div>
|
||||
<div class="legend-item"><div class="legend-dot" style="background:{COLORS['pull']}"></div>Pull weights</div>
|
||||
<div class="legend-item"><div class="legend-dot" style="background:{COLORS['aggregate']}"></div>Aggregate</div>
|
||||
<div class="legend-item"><div class="legend-dot" style="background:{COLORS['dead']}"></div>Dead / Error</div>
|
||||
</div>
|
||||
|
||||
<div class="timeline">
|
||||
{svg}
|
||||
</div>
|
||||
|
||||
<div class="event-log">
|
||||
<pre>{_render_event_log(tracer)}</pre>
|
||||
</div>
|
||||
|
||||
</body>
|
||||
</html>"""
|
||||
|
||||
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} <span class="err">!! {html.escape(err)}</span>'
|
||||
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} <span class="err">!! {html.escape(err)}</span>'
|
||||
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)
|
||||
300
sched/scheduler.py
Normal file
300
sched/scheduler.py
Normal file
|
|
@ -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
|
||||
424
sched/trace.py
Normal file
424
sched/trace.py
Normal file
|
|
@ -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()
|
||||
84
sched/transport.py
Normal file
84
sched/transport.py
Normal file
|
|
@ -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)
|
||||
47
sched/vastai.py
Normal file
47
sched/vastai.py
Normal file
|
|
@ -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))
|
||||
0
tests/__init__.py
Normal file
0
tests/__init__.py
Normal file
BIN
tests/__pycache__/__init__.cpython-313.pyc
Normal file
BIN
tests/__pycache__/__init__.cpython-313.pyc
Normal file
Binary file not shown.
BIN
tests/__pycache__/mock_transport.cpython-313.pyc
Normal file
BIN
tests/__pycache__/mock_transport.cpython-313.pyc
Normal file
Binary file not shown.
BIN
tests/__pycache__/test_aggregator.cpython-313-pytest-9.0.2.pyc
Normal file
BIN
tests/__pycache__/test_aggregator.cpython-313-pytest-9.0.2.pyc
Normal file
Binary file not shown.
BIN
tests/__pycache__/test_scheduler.cpython-313-pytest-9.0.2.pyc
Normal file
BIN
tests/__pycache__/test_scheduler.cpython-313-pytest-9.0.2.pyc
Normal file
Binary file not shown.
BIN
tests/__pycache__/test_trace.cpython-313-pytest-9.0.2.pyc
Normal file
BIN
tests/__pycache__/test_trace.cpython-313-pytest-9.0.2.pyc
Normal file
Binary file not shown.
85
tests/mock_transport.py
Normal file
85
tests/mock_transport.py
Normal file
|
|
@ -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)
|
||||
57
tests/test_aggregator.py
Normal file
57
tests/test_aggregator.py
Normal file
|
|
@ -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"]
|
||||
261
tests/test_scheduler.py
Normal file
261
tests/test_scheduler.py
Normal file
|
|
@ -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()
|
||||
265
tests/test_trace.py
Normal file
265
tests/test_trace.py
Normal file
|
|
@ -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 "<!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()
|
||||
425
uv.lock
Normal file
425
uv.lock
Normal file
|
|
@ -0,0 +1,425 @@
|
|||
version = 1
|
||||
revision = 3
|
||||
requires-python = ">=3.13"
|
||||
resolution-markers = [
|
||||
"sys_platform != 'darwin'",
|
||||
"sys_platform == 'darwin'",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "colorama"
|
||||
version = "0.4.6"
|
||||
source = { registry = "https://pypi.org/simple" }
|
||||
sdist = { url = "https://files.pythonhosted.org/packages/d8/53/6f443c9a4a8358a93a6792e2acffb9d9d5cb0a5cfd8802644b7b1c9a02e4/colorama-0.4.6.tar.gz", hash = "sha256:08695f5cb7ed6e0531a20572697297273c47b8cae5a63ffc6d6ed5c201be6e44", size = 27697, upload-time = "2022-10-25T02:36:22.414Z" }
|
||||
wheels = [
|
||||
{ url = "https://files.pythonhosted.org/packages/d1/d6/3965ed04c63042e047cb6a3e6ed1a63a35087b6a609aa3a15ed8ac56c221/colorama-0.4.6-py2.py3-none-any.whl", hash = "sha256:4f1d9991f5acc0ca119f9d443620b77f9d6b33703e51011c16baf57afb285fc6", size = 25335, upload-time = "2022-10-25T02:36:20.889Z" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "cuda-bindings"
|
||||
version = "12.9.4"
|
||||
source = { registry = "https://pypi.org/simple" }
|
||||
dependencies = [
|
||||
{ name = "cuda-pathfinder", marker = "sys_platform != 'darwin'" },
|
||||
]
|
||||
wheels = [
|
||||
{ url = "https://files.pythonhosted.org/packages/63/56/e465c31dc9111be3441a9ba7df1941fe98f4aa6e71e8788a3fb4534ce24d/cuda_bindings-12.9.4-cp313-cp313-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:32bdc5a76906be4c61eb98f546a6786c5773a881f3b166486449b5d141e4a39f", size = 11906628, upload-time = "2025-10-21T14:51:49.905Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/a3/84/1e6be415e37478070aeeee5884c2022713c1ecc735e6d82d744de0252eee/cuda_bindings-12.9.4-cp313-cp313t-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:56e0043c457a99ac473ddc926fe0dc4046694d99caef633e92601ab52cbe17eb", size = 11925991, upload-time = "2025-10-21T14:51:56.535Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/d1/af/6dfd8f2ed90b1d4719bc053ff8940e494640fe4212dc3dd72f383e4992da/cuda_bindings-12.9.4-cp314-cp314-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:8b72ee72a9cc1b531db31eebaaee5c69a8ec3500e32c6933f2d3b15297b53686", size = 11922703, upload-time = "2025-10-21T14:52:03.585Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/6c/19/90ac264acc00f6df8a49378eedec9fd2db3061bf9263bf9f39fd3d8377c3/cuda_bindings-12.9.4-cp314-cp314t-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:d80bffc357df9988dca279734bc9674c3934a654cab10cadeed27ce17d8635ee", size = 11924658, upload-time = "2025-10-21T14:52:10.411Z" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "cuda-pathfinder"
|
||||
version = "1.3.3"
|
||||
source = { registry = "https://pypi.org/simple" }
|
||||
wheels = [
|
||||
{ url = "https://files.pythonhosted.org/packages/0b/02/4dbe7568a42e46582248942f54dc64ad094769532adbe21e525e4edf7bc4/cuda_pathfinder-1.3.3-py3-none-any.whl", hash = "sha256:9984b664e404f7c134954a771be8775dfd6180ea1e1aef4a5a37d4be05d9bbb1", size = 27154, upload-time = "2025-12-04T22:35:08.996Z" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "filelock"
|
||||
version = "3.20.0"
|
||||
source = { registry = "https://pypi.org/simple" }
|
||||
sdist = { url = "https://files.pythonhosted.org/packages/58/46/0028a82567109b5ef6e4d2a1f04a583fb513e6cf9527fcdd09afd817deeb/filelock-3.20.0.tar.gz", hash = "sha256:711e943b4ec6be42e1d4e6690b48dc175c822967466bb31c0c293f34334c13f4", size = 18922, upload-time = "2025-10-08T18:03:50.056Z" }
|
||||
wheels = [
|
||||
{ url = "https://files.pythonhosted.org/packages/76/91/7216b27286936c16f5b4d0c530087e4a54eead683e6b0b73dd0c64844af6/filelock-3.20.0-py3-none-any.whl", hash = "sha256:339b4732ffda5cd79b13f4e2711a31b0365ce445d95d243bb996273d072546a2", size = 16054, upload-time = "2025-10-08T18:03:48.35Z" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "fsspec"
|
||||
version = "2025.12.0"
|
||||
source = { registry = "https://pypi.org/simple" }
|
||||
sdist = { url = "https://files.pythonhosted.org/packages/b6/27/954057b0d1f53f086f681755207dda6de6c660ce133c829158e8e8fe7895/fsspec-2025.12.0.tar.gz", hash = "sha256:c505de011584597b1060ff778bb664c1bc022e87921b0e4f10cc9c44f9635973", size = 309748, upload-time = "2025-12-03T15:23:42.687Z" }
|
||||
wheels = [
|
||||
{ url = "https://files.pythonhosted.org/packages/51/c7/b64cae5dba3a1b138d7123ec36bb5ccd39d39939f18454407e5468f4763f/fsspec-2025.12.0-py3-none-any.whl", hash = "sha256:8bf1fe301b7d8acfa6e8571e3b1c3d158f909666642431cc78a1b7b4dbc5ec5b", size = 201422, upload-time = "2025-12-03T15:23:41.434Z" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "iniconfig"
|
||||
version = "2.3.0"
|
||||
source = { registry = "https://pypi.org/simple" }
|
||||
sdist = { url = "https://files.pythonhosted.org/packages/72/34/14ca021ce8e5dfedc35312d08ba8bf51fdd999c576889fc2c24cb97f4f10/iniconfig-2.3.0.tar.gz", hash = "sha256:c76315c77db068650d49c5b56314774a7804df16fee4402c1f19d6d15d8c4730", size = 20503, upload-time = "2025-10-18T21:55:43.219Z" }
|
||||
wheels = [
|
||||
{ url = "https://files.pythonhosted.org/packages/cb/b1/3846dd7f199d53cb17f49cba7e651e9ce294d8497c8c150530ed11865bb8/iniconfig-2.3.0-py3-none-any.whl", hash = "sha256:f631c04d2c48c52b84d0d0549c99ff3859c98df65b3101406327ecc7d53fbf12", size = 7484, upload-time = "2025-10-18T21:55:41.639Z" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "jinja2"
|
||||
version = "3.1.6"
|
||||
source = { registry = "https://pypi.org/simple" }
|
||||
dependencies = [
|
||||
{ name = "markupsafe" },
|
||||
]
|
||||
sdist = { url = "https://files.pythonhosted.org/packages/df/bf/f7da0350254c0ed7c72f3e33cef02e048281fec7ecec5f032d4aac52226b/jinja2-3.1.6.tar.gz", hash = "sha256:0137fb05990d35f1275a587e9aee6d56da821fc83491a0fb838183be43f66d6d", size = 245115, upload-time = "2025-03-05T20:05:02.478Z" }
|
||||
wheels = [
|
||||
{ url = "https://files.pythonhosted.org/packages/62/a1/3d680cbfd5f4b8f15abc1d571870c5fc3e594bb582bc3b64ea099db13e56/jinja2-3.1.6-py3-none-any.whl", hash = "sha256:85ece4451f492d0c13c5dd7c13a64681a86afae63a5f347908daf103ce6d2f67", size = 134899, upload-time = "2025-03-05T20:05:00.369Z" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "markupsafe"
|
||||
version = "3.0.2"
|
||||
source = { registry = "https://pypi.org/simple" }
|
||||
sdist = { url = "https://files.pythonhosted.org/packages/b2/97/5d42485e71dfc078108a86d6de8fa46db44a1a9295e89c5d6d4a06e23a62/markupsafe-3.0.2.tar.gz", hash = "sha256:ee55d3edf80167e48ea11a923c7386f4669df67d7994554387f84e7d8b0a2bf0", size = 20537, upload-time = "2024-10-18T15:21:54.129Z" }
|
||||
wheels = [
|
||||
{ url = "https://files.pythonhosted.org/packages/83/0e/67eb10a7ecc77a0c2bbe2b0235765b98d164d81600746914bebada795e97/MarkupSafe-3.0.2-cp313-cp313-macosx_10_13_universal2.whl", hash = "sha256:ba9527cdd4c926ed0760bc301f6728ef34d841f405abf9d4f959c478421e4efd", size = 14274, upload-time = "2024-10-18T15:21:24.577Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/2b/6d/9409f3684d3335375d04e5f05744dfe7e9f120062c9857df4ab490a1031a/MarkupSafe-3.0.2-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:f8b3d067f2e40fe93e1ccdd6b2e1d16c43140e76f02fb1319a05cf2b79d99430", size = 12352, upload-time = "2024-10-18T15:21:25.382Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/d2/f5/6eadfcd3885ea85fe2a7c128315cc1bb7241e1987443d78c8fe712d03091/MarkupSafe-3.0.2-cp313-cp313-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:569511d3b58c8791ab4c2e1285575265991e6d8f8700c7be0e88f86cb0672094", size = 24122, upload-time = "2024-10-18T15:21:26.199Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/0c/91/96cf928db8236f1bfab6ce15ad070dfdd02ed88261c2afafd4b43575e9e9/MarkupSafe-3.0.2-cp313-cp313-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:15ab75ef81add55874e7ab7055e9c397312385bd9ced94920f2802310c930396", size = 23085, upload-time = "2024-10-18T15:21:27.029Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/c2/cf/c9d56af24d56ea04daae7ac0940232d31d5a8354f2b457c6d856b2057d69/MarkupSafe-3.0.2-cp313-cp313-manylinux_2_5_i686.manylinux1_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:f3818cb119498c0678015754eba762e0d61e5b52d34c8b13d770f0719f7b1d79", size = 22978, upload-time = "2024-10-18T15:21:27.846Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/2a/9f/8619835cd6a711d6272d62abb78c033bda638fdc54c4e7f4272cf1c0962b/MarkupSafe-3.0.2-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:cdb82a876c47801bb54a690c5ae105a46b392ac6099881cdfb9f6e95e4014c6a", size = 24208, upload-time = "2024-10-18T15:21:28.744Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/f9/bf/176950a1792b2cd2102b8ffeb5133e1ed984547b75db47c25a67d3359f77/MarkupSafe-3.0.2-cp313-cp313-musllinux_1_2_i686.whl", hash = "sha256:cabc348d87e913db6ab4aa100f01b08f481097838bdddf7c7a84b7575b7309ca", size = 23357, upload-time = "2024-10-18T15:21:29.545Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/ce/4f/9a02c1d335caabe5c4efb90e1b6e8ee944aa245c1aaaab8e8a618987d816/MarkupSafe-3.0.2-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:444dcda765c8a838eaae23112db52f1efaf750daddb2d9ca300bcae1039adc5c", size = 23344, upload-time = "2024-10-18T15:21:30.366Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/ee/55/c271b57db36f748f0e04a759ace9f8f759ccf22b4960c270c78a394f58be/MarkupSafe-3.0.2-cp313-cp313-win32.whl", hash = "sha256:bcf3e58998965654fdaff38e58584d8937aa3096ab5354d493c77d1fdd66d7a1", size = 15101, upload-time = "2024-10-18T15:21:31.207Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/29/88/07df22d2dd4df40aba9f3e402e6dc1b8ee86297dddbad4872bd5e7b0094f/MarkupSafe-3.0.2-cp313-cp313-win_amd64.whl", hash = "sha256:e6a2a455bd412959b57a172ce6328d2dd1f01cb2135efda2e4576e8a23fa3b0f", size = 15603, upload-time = "2024-10-18T15:21:32.032Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/62/6a/8b89d24db2d32d433dffcd6a8779159da109842434f1dd2f6e71f32f738c/MarkupSafe-3.0.2-cp313-cp313t-macosx_10_13_universal2.whl", hash = "sha256:b5a6b3ada725cea8a5e634536b1b01c30bcdcd7f9c6fff4151548d5bf6b3a36c", size = 14510, upload-time = "2024-10-18T15:21:33.625Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/7a/06/a10f955f70a2e5a9bf78d11a161029d278eeacbd35ef806c3fd17b13060d/MarkupSafe-3.0.2-cp313-cp313t-macosx_11_0_arm64.whl", hash = "sha256:a904af0a6162c73e3edcb969eeeb53a63ceeb5d8cf642fade7d39e7963a22ddb", size = 12486, upload-time = "2024-10-18T15:21:34.611Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/34/cf/65d4a571869a1a9078198ca28f39fba5fbb910f952f9dbc5220afff9f5e6/MarkupSafe-3.0.2-cp313-cp313t-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:4aa4e5faecf353ed117801a068ebab7b7e09ffb6e1d5e412dc852e0da018126c", size = 25480, upload-time = "2024-10-18T15:21:35.398Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/0c/e3/90e9651924c430b885468b56b3d597cabf6d72be4b24a0acd1fa0e12af67/MarkupSafe-3.0.2-cp313-cp313t-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:c0ef13eaeee5b615fb07c9a7dadb38eac06a0608b41570d8ade51c56539e509d", size = 23914, upload-time = "2024-10-18T15:21:36.231Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/66/8c/6c7cf61f95d63bb866db39085150df1f2a5bd3335298f14a66b48e92659c/MarkupSafe-3.0.2-cp313-cp313t-manylinux_2_5_i686.manylinux1_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:d16a81a06776313e817c951135cf7340a3e91e8c1ff2fac444cfd75fffa04afe", size = 23796, upload-time = "2024-10-18T15:21:37.073Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/bb/35/cbe9238ec3f47ac9a7c8b3df7a808e7cb50fe149dc7039f5f454b3fba218/MarkupSafe-3.0.2-cp313-cp313t-musllinux_1_2_aarch64.whl", hash = "sha256:6381026f158fdb7c72a168278597a5e3a5222e83ea18f543112b2662a9b699c5", size = 25473, upload-time = "2024-10-18T15:21:37.932Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/e6/32/7621a4382488aa283cc05e8984a9c219abad3bca087be9ec77e89939ded9/MarkupSafe-3.0.2-cp313-cp313t-musllinux_1_2_i686.whl", hash = "sha256:3d79d162e7be8f996986c064d1c7c817f6df3a77fe3d6859f6f9e7be4b8c213a", size = 24114, upload-time = "2024-10-18T15:21:39.799Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/0d/80/0985960e4b89922cb5a0bac0ed39c5b96cbc1a536a99f30e8c220a996ed9/MarkupSafe-3.0.2-cp313-cp313t-musllinux_1_2_x86_64.whl", hash = "sha256:131a3c7689c85f5ad20f9f6fb1b866f402c445b220c19fe4308c0b147ccd2ad9", size = 24098, upload-time = "2024-10-18T15:21:40.813Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/82/78/fedb03c7d5380df2427038ec8d973587e90561b2d90cd472ce9254cf348b/MarkupSafe-3.0.2-cp313-cp313t-win32.whl", hash = "sha256:ba8062ed2cf21c07a9e295d5b8a2a5ce678b913b45fdf68c32d95d6c1291e0b6", size = 15208, upload-time = "2024-10-18T15:21:41.814Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/4f/65/6079a46068dfceaeabb5dcad6d674f5f5c61a6fa5673746f42a9f4c233b3/MarkupSafe-3.0.2-cp313-cp313t-win_amd64.whl", hash = "sha256:e444a31f8db13eb18ada366ab3cf45fd4b31e4db1236a4448f68778c1d1a5a2f", size = 15739, upload-time = "2024-10-18T15:21:42.784Z" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "mpmath"
|
||||
version = "1.3.0"
|
||||
source = { registry = "https://pypi.org/simple" }
|
||||
sdist = { url = "https://files.pythonhosted.org/packages/e0/47/dd32fa426cc72114383ac549964eecb20ecfd886d1e5ccf5340b55b02f57/mpmath-1.3.0.tar.gz", hash = "sha256:7a28eb2a9774d00c7bc92411c19a89209d5da7c4c9a9e227be8330a23a25b91f", size = 508106, upload-time = "2023-03-07T16:47:11.061Z" }
|
||||
wheels = [
|
||||
{ url = "https://files.pythonhosted.org/packages/43/e3/7d92a15f894aa0c9c4b49b8ee9ac9850d6e63b03c9c32c0367a13ae62209/mpmath-1.3.0-py3-none-any.whl", hash = "sha256:a0b2b9fe80bbcd81a6647ff13108738cfb482d481d826cc0e02f5b35e5c88d2c", size = 536198, upload-time = "2023-03-07T16:47:09.197Z" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "networkx"
|
||||
version = "3.6.1"
|
||||
source = { registry = "https://pypi.org/simple" }
|
||||
sdist = { url = "https://files.pythonhosted.org/packages/6a/51/63fe664f3908c97be9d2e4f1158eb633317598cfa6e1fc14af5383f17512/networkx-3.6.1.tar.gz", hash = "sha256:26b7c357accc0c8cde558ad486283728b65b6a95d85ee1cd66bafab4c8168509", size = 2517025, upload-time = "2025-12-08T17:02:39.908Z" }
|
||||
wheels = [
|
||||
{ url = "https://files.pythonhosted.org/packages/9e/c9/b2622292ea83fbb4ec318f5b9ab867d0a28ab43c5717bb85b0a5f6b3b0a4/networkx-3.6.1-py3-none-any.whl", hash = "sha256:d47fbf302e7d9cbbb9e2555a0d267983d2aa476bac30e90dfbe5669bd57f3762", size = 2068504, upload-time = "2025-12-08T17:02:38.159Z" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "nvidia-cublas-cu12"
|
||||
version = "12.8.4.1"
|
||||
source = { registry = "https://pypi.org/simple" }
|
||||
wheels = [
|
||||
{ url = "https://files.pythonhosted.org/packages/dc/61/e24b560ab2e2eaeb3c839129175fb330dfcfc29e5203196e5541a4c44682/nvidia_cublas_cu12-12.8.4.1-py3-none-manylinux_2_27_x86_64.whl", hash = "sha256:8ac4e771d5a348c551b2a426eda6193c19aa630236b418086020df5ba9667142", size = 594346921, upload-time = "2025-03-07T01:44:31.254Z" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "nvidia-cuda-cupti-cu12"
|
||||
version = "12.8.90"
|
||||
source = { registry = "https://pypi.org/simple" }
|
||||
wheels = [
|
||||
{ url = "https://files.pythonhosted.org/packages/f8/02/2adcaa145158bf1a8295d83591d22e4103dbfd821bcaf6f3f53151ca4ffa/nvidia_cuda_cupti_cu12-12.8.90-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:ea0cb07ebda26bb9b29ba82cda34849e73c166c18162d3913575b0c9db9a6182", size = 10248621, upload-time = "2025-03-07T01:40:21.213Z" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "nvidia-cuda-nvrtc-cu12"
|
||||
version = "12.8.93"
|
||||
source = { registry = "https://pypi.org/simple" }
|
||||
wheels = [
|
||||
{ url = "https://files.pythonhosted.org/packages/05/6b/32f747947df2da6994e999492ab306a903659555dddc0fbdeb9d71f75e52/nvidia_cuda_nvrtc_cu12-12.8.93-py3-none-manylinux2010_x86_64.manylinux_2_12_x86_64.whl", hash = "sha256:a7756528852ef889772a84c6cd89d41dfa74667e24cca16bb31f8f061e3e9994", size = 88040029, upload-time = "2025-03-07T01:42:13.562Z" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "nvidia-cuda-runtime-cu12"
|
||||
version = "12.8.90"
|
||||
source = { registry = "https://pypi.org/simple" }
|
||||
wheels = [
|
||||
{ url = "https://files.pythonhosted.org/packages/0d/9b/a997b638fcd068ad6e4d53b8551a7d30fe8b404d6f1804abf1df69838932/nvidia_cuda_runtime_cu12-12.8.90-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:adade8dcbd0edf427b7204d480d6066d33902cab2a4707dcfc48a2d0fd44ab90", size = 954765, upload-time = "2025-03-07T01:40:01.615Z" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "nvidia-cudnn-cu12"
|
||||
version = "9.10.2.21"
|
||||
source = { registry = "https://pypi.org/simple" }
|
||||
dependencies = [
|
||||
{ name = "nvidia-cublas-cu12", marker = "sys_platform != 'darwin'" },
|
||||
]
|
||||
wheels = [
|
||||
{ url = "https://files.pythonhosted.org/packages/ba/51/e123d997aa098c61d029f76663dedbfb9bc8dcf8c60cbd6adbe42f76d049/nvidia_cudnn_cu12-9.10.2.21-py3-none-manylinux_2_27_x86_64.whl", hash = "sha256:949452be657fa16687d0930933f032835951ef0892b37d2d53824d1a84dc97a8", size = 706758467, upload-time = "2025-06-06T21:54:08.597Z" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "nvidia-cufft-cu12"
|
||||
version = "11.3.3.83"
|
||||
source = { registry = "https://pypi.org/simple" }
|
||||
dependencies = [
|
||||
{ name = "nvidia-nvjitlink-cu12", marker = "sys_platform != 'darwin'" },
|
||||
]
|
||||
wheels = [
|
||||
{ url = "https://files.pythonhosted.org/packages/1f/13/ee4e00f30e676b66ae65b4f08cb5bcbb8392c03f54f2d5413ea99a5d1c80/nvidia_cufft_cu12-11.3.3.83-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:4d2dd21ec0b88cf61b62e6b43564355e5222e4a3fb394cac0db101f2dd0d4f74", size = 193118695, upload-time = "2025-03-07T01:45:27.821Z" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "nvidia-cufile-cu12"
|
||||
version = "1.13.1.3"
|
||||
source = { registry = "https://pypi.org/simple" }
|
||||
wheels = [
|
||||
{ url = "https://files.pythonhosted.org/packages/bb/fe/1bcba1dfbfb8d01be8d93f07bfc502c93fa23afa6fd5ab3fc7c1df71038a/nvidia_cufile_cu12-1.13.1.3-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:1d069003be650e131b21c932ec3d8969c1715379251f8d23a1860554b1cb24fc", size = 1197834, upload-time = "2025-03-07T01:45:50.723Z" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "nvidia-curand-cu12"
|
||||
version = "10.3.9.90"
|
||||
source = { registry = "https://pypi.org/simple" }
|
||||
wheels = [
|
||||
{ url = "https://files.pythonhosted.org/packages/fb/aa/6584b56dc84ebe9cf93226a5cde4d99080c8e90ab40f0c27bda7a0f29aa1/nvidia_curand_cu12-10.3.9.90-py3-none-manylinux_2_27_x86_64.whl", hash = "sha256:b32331d4f4df5d6eefa0554c565b626c7216f87a06a4f56fab27c3b68a830ec9", size = 63619976, upload-time = "2025-03-07T01:46:23.323Z" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "nvidia-cusolver-cu12"
|
||||
version = "11.7.3.90"
|
||||
source = { registry = "https://pypi.org/simple" }
|
||||
dependencies = [
|
||||
{ name = "nvidia-cublas-cu12", marker = "sys_platform != 'darwin'" },
|
||||
{ name = "nvidia-cusparse-cu12", marker = "sys_platform != 'darwin'" },
|
||||
{ name = "nvidia-nvjitlink-cu12", marker = "sys_platform != 'darwin'" },
|
||||
]
|
||||
wheels = [
|
||||
{ url = "https://files.pythonhosted.org/packages/85/48/9a13d2975803e8cf2777d5ed57b87a0b6ca2cc795f9a4f59796a910bfb80/nvidia_cusolver_cu12-11.7.3.90-py3-none-manylinux_2_27_x86_64.whl", hash = "sha256:4376c11ad263152bd50ea295c05370360776f8c3427b30991df774f9fb26c450", size = 267506905, upload-time = "2025-03-07T01:47:16.273Z" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "nvidia-cusparse-cu12"
|
||||
version = "12.5.8.93"
|
||||
source = { registry = "https://pypi.org/simple" }
|
||||
dependencies = [
|
||||
{ name = "nvidia-nvjitlink-cu12", marker = "sys_platform != 'darwin'" },
|
||||
]
|
||||
wheels = [
|
||||
{ url = "https://files.pythonhosted.org/packages/c2/f5/e1854cb2f2bcd4280c44736c93550cc300ff4b8c95ebe370d0aa7d2b473d/nvidia_cusparse_cu12-12.5.8.93-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:1ec05d76bbbd8b61b06a80e1eaf8cf4959c3d4ce8e711b65ebd0443bb0ebb13b", size = 288216466, upload-time = "2025-03-07T01:48:13.779Z" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "nvidia-cusparselt-cu12"
|
||||
version = "0.7.1"
|
||||
source = { registry = "https://pypi.org/simple" }
|
||||
wheels = [
|
||||
{ url = "https://files.pythonhosted.org/packages/56/79/12978b96bd44274fe38b5dde5cfb660b1d114f70a65ef962bcbbed99b549/nvidia_cusparselt_cu12-0.7.1-py3-none-manylinux2014_x86_64.whl", hash = "sha256:f1bb701d6b930d5a7cea44c19ceb973311500847f81b634d802b7b539dc55623", size = 287193691, upload-time = "2025-02-26T00:15:44.104Z" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "nvidia-nccl-cu12"
|
||||
version = "2.27.5"
|
||||
source = { registry = "https://pypi.org/simple" }
|
||||
wheels = [
|
||||
{ url = "https://files.pythonhosted.org/packages/6e/89/f7a07dc961b60645dbbf42e80f2bc85ade7feb9a491b11a1e973aa00071f/nvidia_nccl_cu12-2.27.5-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:ad730cf15cb5d25fe849c6e6ca9eb5b76db16a80f13f425ac68d8e2e55624457", size = 322348229, upload-time = "2025-06-26T04:11:28.385Z" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "nvidia-nvjitlink-cu12"
|
||||
version = "12.8.93"
|
||||
source = { registry = "https://pypi.org/simple" }
|
||||
wheels = [
|
||||
{ url = "https://files.pythonhosted.org/packages/f6/74/86a07f1d0f42998ca31312f998bd3b9a7eff7f52378f4f270c8679c77fb9/nvidia_nvjitlink_cu12-12.8.93-py3-none-manylinux2010_x86_64.manylinux_2_12_x86_64.whl", hash = "sha256:81ff63371a7ebd6e6451970684f916be2eab07321b73c9d244dc2b4da7f73b88", size = 39254836, upload-time = "2025-03-07T01:49:55.661Z" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "nvidia-nvshmem-cu12"
|
||||
version = "3.4.5"
|
||||
source = { registry = "https://pypi.org/simple" }
|
||||
wheels = [
|
||||
{ url = "https://files.pythonhosted.org/packages/b5/09/6ea3ea725f82e1e76684f0708bbedd871fc96da89945adeba65c3835a64c/nvidia_nvshmem_cu12-3.4.5-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:042f2500f24c021db8a06c5eec2539027d57460e1c1a762055a6554f72c369bd", size = 139103095, upload-time = "2025-09-06T00:32:31.266Z" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "nvidia-nvtx-cu12"
|
||||
version = "12.8.90"
|
||||
source = { registry = "https://pypi.org/simple" }
|
||||
wheels = [
|
||||
{ url = "https://files.pythonhosted.org/packages/a2/eb/86626c1bbc2edb86323022371c39aa48df6fd8b0a1647bc274577f72e90b/nvidia_nvtx_cu12-12.8.90-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:5b17e2001cc0d751a5bc2c6ec6d26ad95913324a4adb86788c944f8ce9ba441f", size = 89954, upload-time = "2025-03-07T01:42:44.131Z" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "packaging"
|
||||
version = "26.0"
|
||||
source = { registry = "https://pypi.org/simple" }
|
||||
sdist = { url = "https://files.pythonhosted.org/packages/65/ee/299d360cdc32edc7d2cf530f3accf79c4fca01e96ffc950d8a52213bd8e4/packaging-26.0.tar.gz", hash = "sha256:00243ae351a257117b6a241061796684b084ed1c516a08c48a3f7e147a9d80b4", size = 143416, upload-time = "2026-01-21T20:50:39.064Z" }
|
||||
wheels = [
|
||||
{ url = "https://files.pythonhosted.org/packages/b7/b9/c538f279a4e237a006a2c98387d081e9eb060d203d8ed34467cc0f0b9b53/packaging-26.0-py3-none-any.whl", hash = "sha256:b36f1fef9334a5588b4166f8bcd26a14e521f2b55e6b9de3aaa80d3ff7a37529", size = 74366, upload-time = "2026-01-21T20:50:37.788Z" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "pluggy"
|
||||
version = "1.6.0"
|
||||
source = { registry = "https://pypi.org/simple" }
|
||||
sdist = { url = "https://files.pythonhosted.org/packages/f9/e2/3e91f31a7d2b083fe6ef3fa267035b518369d9511ffab804f839851d2779/pluggy-1.6.0.tar.gz", hash = "sha256:7dcc130b76258d33b90f61b658791dede3486c3e6bfb003ee5c9bfb396dd22f3", size = 69412, upload-time = "2025-05-15T12:30:07.975Z" }
|
||||
wheels = [
|
||||
{ url = "https://files.pythonhosted.org/packages/54/20/4d324d65cc6d9205fabedc306948156824eb9f0ee1633355a8f7ec5c66bf/pluggy-1.6.0-py3-none-any.whl", hash = "sha256:e920276dd6813095e9377c0bc5566d94c932c33b27a3e3945d8389c374dd4746", size = 20538, upload-time = "2025-05-15T12:30:06.134Z" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "pygments"
|
||||
version = "2.19.2"
|
||||
source = { registry = "https://pypi.org/simple" }
|
||||
sdist = { url = "https://files.pythonhosted.org/packages/b0/77/a5b8c569bf593b0140bde72ea885a803b82086995367bf2037de0159d924/pygments-2.19.2.tar.gz", hash = "sha256:636cb2477cec7f8952536970bc533bc43743542f70392ae026374600add5b887", size = 4968631, upload-time = "2025-06-21T13:39:12.283Z" }
|
||||
wheels = [
|
||||
{ url = "https://files.pythonhosted.org/packages/c7/21/705964c7812476f378728bdf590ca4b771ec72385c533964653c68e86bdc/pygments-2.19.2-py3-none-any.whl", hash = "sha256:86540386c03d588bb81d44bc3928634ff26449851e99741617ecb9037ee5ec0b", size = 1225217, upload-time = "2025-06-21T13:39:07.939Z" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "pytest"
|
||||
version = "9.0.2"
|
||||
source = { registry = "https://pypi.org/simple" }
|
||||
dependencies = [
|
||||
{ name = "colorama", marker = "sys_platform == 'win32'" },
|
||||
{ name = "iniconfig" },
|
||||
{ name = "packaging" },
|
||||
{ name = "pluggy" },
|
||||
{ name = "pygments" },
|
||||
]
|
||||
sdist = { url = "https://files.pythonhosted.org/packages/d1/db/7ef3487e0fb0049ddb5ce41d3a49c235bf9ad299b6a25d5780a89f19230f/pytest-9.0.2.tar.gz", hash = "sha256:75186651a92bd89611d1d9fc20f0b4345fd827c41ccd5c299a868a05d70edf11", size = 1568901, upload-time = "2025-12-06T21:30:51.014Z" }
|
||||
wheels = [
|
||||
{ url = "https://files.pythonhosted.org/packages/3b/ab/b3226f0bd7cdcf710fbede2b3548584366da3b19b5021e74f5bde2a8fa3f/pytest-9.0.2-py3-none-any.whl", hash = "sha256:711ffd45bf766d5264d487b917733b453d917afd2b0ad65223959f59089f875b", size = 374801, upload-time = "2025-12-06T21:30:49.154Z" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "pytest-asyncio"
|
||||
version = "1.3.0"
|
||||
source = { registry = "https://pypi.org/simple" }
|
||||
dependencies = [
|
||||
{ name = "pytest" },
|
||||
]
|
||||
sdist = { url = "https://files.pythonhosted.org/packages/90/2c/8af215c0f776415f3590cac4f9086ccefd6fd463befeae41cd4d3f193e5a/pytest_asyncio-1.3.0.tar.gz", hash = "sha256:d7f52f36d231b80ee124cd216ffb19369aa168fc10095013c6b014a34d3ee9e5", size = 50087, upload-time = "2025-11-10T16:07:47.256Z" }
|
||||
wheels = [
|
||||
{ url = "https://files.pythonhosted.org/packages/e5/35/f8b19922b6a25bc0880171a2f1a003eaeb93657475193ab516fd87cac9da/pytest_asyncio-1.3.0-py3-none-any.whl", hash = "sha256:611e26147c7f77640e6d0a92a38ed17c3e9848063698d5c93d5aa7aa11cebff5", size = 15075, upload-time = "2025-11-10T16:07:45.537Z" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "setuptools"
|
||||
version = "70.2.0"
|
||||
source = { registry = "https://pypi.org/simple" }
|
||||
sdist = { url = "https://files.pythonhosted.org/packages/8d/e6/2fc95aec377988ff3ca882aa58d4f6ab35ff59a12b1611a9fe3075eb3019/setuptools-70.2.0.tar.gz", hash = "sha256:bd63e505105011b25c3c11f753f7e3b8465ea739efddaccef8f0efac2137bac1", size = 2332711, upload-time = "2024-07-01T16:32:41.666Z" }
|
||||
wheels = [
|
||||
{ url = "https://files.pythonhosted.org/packages/42/54/2a8ecfcc9a714a6fbf86559a4b0f50b126a4ac4269ea8134f2c75c3e73de/setuptools-70.2.0-py3-none-any.whl", hash = "sha256:b8b8060bb426838fbe942479c90296ce976249451118ef566a5a0b7d8b78fb05", size = 930834, upload-time = "2024-07-01T16:32:34.354Z" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "sympy"
|
||||
version = "1.14.0"
|
||||
source = { registry = "https://pypi.org/simple" }
|
||||
dependencies = [
|
||||
{ name = "mpmath" },
|
||||
]
|
||||
sdist = { url = "https://files.pythonhosted.org/packages/83/d3/803453b36afefb7c2bb238361cd4ae6125a569b4db67cd9e79846ba2d68c/sympy-1.14.0.tar.gz", hash = "sha256:d3d3fe8df1e5a0b42f0e7bdf50541697dbe7d23746e894990c030e2b05e72517", size = 7793921, upload-time = "2025-04-27T18:05:01.611Z" }
|
||||
wheels = [
|
||||
{ url = "https://files.pythonhosted.org/packages/a2/09/77d55d46fd61b4a135c444fc97158ef34a095e5681d0a6c10b75bf356191/sympy-1.14.0-py3-none-any.whl", hash = "sha256:e091cc3e99d2141a0ba2847328f5479b05d94a6635cb96148ccb3f34671bd8f5", size = 6299353, upload-time = "2025-04-27T18:04:59.103Z" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "torch"
|
||||
version = "2.10.0"
|
||||
source = { registry = "https://pypi.org/simple" }
|
||||
dependencies = [
|
||||
{ name = "cuda-bindings", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" },
|
||||
{ name = "filelock" },
|
||||
{ name = "fsspec" },
|
||||
{ name = "jinja2" },
|
||||
{ name = "networkx" },
|
||||
{ name = "nvidia-cublas-cu12", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" },
|
||||
{ name = "nvidia-cuda-cupti-cu12", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" },
|
||||
{ name = "nvidia-cuda-nvrtc-cu12", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" },
|
||||
{ name = "nvidia-cuda-runtime-cu12", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" },
|
||||
{ name = "nvidia-cudnn-cu12", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" },
|
||||
{ name = "nvidia-cufft-cu12", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" },
|
||||
{ name = "nvidia-cufile-cu12", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" },
|
||||
{ name = "nvidia-curand-cu12", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" },
|
||||
{ name = "nvidia-cusolver-cu12", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" },
|
||||
{ name = "nvidia-cusparse-cu12", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" },
|
||||
{ name = "nvidia-cusparselt-cu12", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" },
|
||||
{ name = "nvidia-nccl-cu12", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" },
|
||||
{ name = "nvidia-nvjitlink-cu12", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" },
|
||||
{ name = "nvidia-nvshmem-cu12", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" },
|
||||
{ name = "nvidia-nvtx-cu12", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" },
|
||||
{ name = "setuptools" },
|
||||
{ name = "sympy" },
|
||||
{ name = "triton", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" },
|
||||
{ name = "typing-extensions" },
|
||||
]
|
||||
wheels = [
|
||||
{ url = "https://files.pythonhosted.org/packages/c9/6f/f2e91e34e3fcba2e3fc8d8f74e7d6c22e74e480bbd1db7bc8900fdf3e95c/torch-2.10.0-cp313-cp313-manylinux_2_28_aarch64.whl", hash = "sha256:5c4d217b14741e40776dd7074d9006fd28b8a97ef5654db959d8635b2fe5f29b", size = 146004247, upload-time = "2026-01-21T16:24:29.335Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/98/fb/5160261aeb5e1ee12ee95fe599d0541f7c976c3701d607d8fc29e623229f/torch-2.10.0-cp313-cp313-manylinux_2_28_x86_64.whl", hash = "sha256:6b71486353fce0f9714ca0c9ef1c850a2ae766b409808acd58e9678a3edb7738", size = 915716445, upload-time = "2026-01-21T16:22:45.353Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/6a/16/502fb1b41e6d868e8deb5b0e3ae926bbb36dab8ceb0d1b769b266ad7b0c3/torch-2.10.0-cp313-cp313-win_amd64.whl", hash = "sha256:c2ee399c644dc92ef7bc0d4f7e74b5360c37cdbe7c5ba11318dda49ffac2bc57", size = 113757050, upload-time = "2026-01-21T16:24:19.204Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/1a/0b/39929b148f4824bc3ad6f9f72a29d4ad865bcf7ebfc2fa67584773e083d2/torch-2.10.0-cp313-cp313t-macosx_14_0_arm64.whl", hash = "sha256:3202429f58309b9fa96a614885eace4b7995729f44beb54d3e4a47773649d382", size = 79851305, upload-time = "2026-01-21T16:24:09.209Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/d8/14/21fbce63bc452381ba5f74a2c0a959fdf5ad5803ccc0c654e752e0dbe91a/torch-2.10.0-cp313-cp313t-manylinux_2_28_aarch64.whl", hash = "sha256:aae1b29cd68e50a9397f5ee897b9c24742e9e306f88a807a27d617f07adb3bd8", size = 146005472, upload-time = "2026-01-21T16:22:29.022Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/54/fd/b207d1c525cb570ef47f3e9f836b154685011fce11a2f444ba8a4084d042/torch-2.10.0-cp313-cp313t-manylinux_2_28_x86_64.whl", hash = "sha256:6021db85958db2f07ec94e1bc77212721ba4920c12a18dc552d2ae36a3eb163f", size = 915612644, upload-time = "2026-01-21T16:21:47.019Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/36/53/0197f868c75f1050b199fe58f9bf3bf3aecac9b4e85cc9c964383d745403/torch-2.10.0-cp313-cp313t-win_amd64.whl", hash = "sha256:ff43db38af76fda183156153983c9a096fc4c78d0cd1e07b14a2314c7f01c2c8", size = 113997015, upload-time = "2026-01-21T16:23:00.767Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/0e/13/e76b4d9c160e89fff48bf16b449ea324bda84745d2ab30294c37c2434c0d/torch-2.10.0-cp313-none-macosx_11_0_arm64.whl", hash = "sha256:cdf2a523d699b70d613243211ecaac14fe9c5df8a0b0a9c02add60fb2a413e0f", size = 79498248, upload-time = "2026-01-21T16:23:09.315Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/4f/93/716b5ac0155f1be70ed81bacc21269c3ece8dba0c249b9994094110bfc51/torch-2.10.0-cp314-cp314-macosx_14_0_arm64.whl", hash = "sha256:bf0d9ff448b0218e0433aeb198805192346c4fd659c852370d5cc245f602a06a", size = 79464992, upload-time = "2026-01-21T16:23:05.162Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/69/2b/51e663ff190c9d16d4a8271203b71bc73a16aa7619b9f271a69b9d4a936b/torch-2.10.0-cp314-cp314-manylinux_2_28_aarch64.whl", hash = "sha256:233aed0659a2503b831d8a67e9da66a62c996204c0bba4f4c442ccc0c68a3f60", size = 146018567, upload-time = "2026-01-21T16:22:23.393Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/5e/cd/4b95ef7f293b927c283db0b136c42be91c8ec6845c44de0238c8c23bdc80/torch-2.10.0-cp314-cp314-manylinux_2_28_x86_64.whl", hash = "sha256:682497e16bdfa6efeec8cde66531bc8d1fbbbb4d8788ec6173c089ed3cc2bfe5", size = 915721646, upload-time = "2026-01-21T16:21:16.983Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/56/97/078a007208f8056d88ae43198833469e61a0a355abc0b070edd2c085eb9a/torch-2.10.0-cp314-cp314-win_amd64.whl", hash = "sha256:6528f13d2a8593a1a412ea07a99812495bec07e9224c28b2a25c0a30c7da025c", size = 113752373, upload-time = "2026-01-21T16:22:13.471Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/d8/94/71994e7d0d5238393df9732fdab607e37e2b56d26a746cb59fdb415f8966/torch-2.10.0-cp314-cp314t-macosx_14_0_arm64.whl", hash = "sha256:f5ab4ba32383061be0fb74bda772d470140a12c1c3b58a0cfbf3dae94d164c28", size = 79850324, upload-time = "2026-01-21T16:22:09.494Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/e2/65/1a05346b418ea8ccd10360eef4b3e0ce688fba544e76edec26913a8d0ee0/torch-2.10.0-cp314-cp314t-manylinux_2_28_aarch64.whl", hash = "sha256:716b01a176c2a5659c98f6b01bf868244abdd896526f1c692712ab36dbaf9b63", size = 146006482, upload-time = "2026-01-21T16:22:18.42Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/1d/b9/5f6f9d9e859fc3235f60578fa64f52c9c6e9b4327f0fe0defb6de5c0de31/torch-2.10.0-cp314-cp314t-manylinux_2_28_x86_64.whl", hash = "sha256:d8f5912ba938233f86361e891789595ff35ca4b4e2ac8fe3670895e5976731d6", size = 915613050, upload-time = "2026-01-21T16:20:49.035Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/66/4d/35352043ee0eaffdeff154fad67cd4a31dbed7ff8e3be1cc4549717d6d51/torch-2.10.0-cp314-cp314t-win_amd64.whl", hash = "sha256:71283a373f0ee2c89e0f0d5f446039bdabe8dbc3c9ccf35f0f784908b0acd185", size = 113995816, upload-time = "2026-01-21T16:22:05.312Z" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "triton"
|
||||
version = "3.6.0"
|
||||
source = { registry = "https://pypi.org/simple" }
|
||||
wheels = [
|
||||
{ url = "https://files.pythonhosted.org/packages/f9/0b/37d991d8c130ce81a8728ae3c25b6e60935838e9be1b58791f5997b24a54/triton-3.6.0-cp313-cp313-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:10c7f76c6e72d2ef08df639e3d0d30729112f47a56b0c81672edc05ee5116ac9", size = 188289450, upload-time = "2026-01-20T16:00:49.136Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/35/f8/9c66bfc55361ec6d0e4040a0337fb5924ceb23de4648b8a81ae9d33b2b38/triton-3.6.0-cp313-cp313t-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:d002e07d7180fd65e622134fbd980c9a3d4211fb85224b56a0a0efbd422ab72f", size = 188400296, upload-time = "2026-01-20T16:00:56.042Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/df/3d/9e7eee57b37c80cec63322c0231bb6da3cfe535a91d7a4d64896fcb89357/triton-3.6.0-cp314-cp314-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:a17a5d5985f0ac494ed8a8e54568f092f7057ef60e1b0fa09d3fd1512064e803", size = 188273063, upload-time = "2026-01-20T16:01:07.278Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/f6/56/6113c23ff46c00aae423333eb58b3e60bdfe9179d542781955a5e1514cb3/triton-3.6.0-cp314-cp314t-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:46bd1c1af4b6704e554cad2eeb3b0a6513a980d470ccfa63189737340c7746a7", size = 188397994, upload-time = "2026-01-20T16:01:14.236Z" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "typing-extensions"
|
||||
version = "4.15.0"
|
||||
source = { registry = "https://pypi.org/simple" }
|
||||
sdist = { url = "https://files.pythonhosted.org/packages/72/94/1a15dd82efb362ac84269196e94cf00f187f7ed21c242792a923cdb1c61f/typing_extensions-4.15.0.tar.gz", hash = "sha256:0cea48d173cc12fa28ecabc3b837ea3cf6f38c6d1136f85cbaaf598984861466", size = 109391, upload-time = "2025-08-25T13:49:26.313Z" }
|
||||
wheels = [
|
||||
{ url = "https://files.pythonhosted.org/packages/18/67/36e9267722cc04a6b9f15c7f3441c2363321a3ea07da7ae0c0707beb2a9c/typing_extensions-4.15.0-py3-none-any.whl", hash = "sha256:f0fa19c6845758ab08074a0cfa8b7aecb71c999ca73d62883bc25cc018c4e548", size = 44614, upload-time = "2025-08-25T13:49:24.86Z" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "vastai-utils"
|
||||
version = "0.1.0"
|
||||
source = { virtual = "." }
|
||||
dependencies = [
|
||||
{ name = "torch" },
|
||||
]
|
||||
|
||||
[package.dev-dependencies]
|
||||
dev = [
|
||||
{ name = "pytest" },
|
||||
{ name = "pytest-asyncio" },
|
||||
]
|
||||
|
||||
[package.metadata]
|
||||
requires-dist = [{ name = "torch", specifier = ">=2.10.0" }]
|
||||
|
||||
[package.metadata.requires-dev]
|
||||
dev = [
|
||||
{ name = "pytest", specifier = ">=9.0.2" },
|
||||
{ name = "pytest-asyncio", specifier = ">=1.3.0" },
|
||||
]
|
||||
0
worker/__init__.py
Normal file
0
worker/__init__.py
Normal file
62
worker/worker.py
Normal file
62
worker/worker.py
Normal file
|
|
@ -0,0 +1,62 @@
|
|||
#!/usr/bin/env python3
|
||||
"""Generic FedAvg worker.
|
||||
|
||||
Deployed to /workspace/ by the scheduler.
|
||||
Loads model params, trains locally for K steps, saves updated weights.
|
||||
The job-specific model and data come from job_module.py (deployed alongside).
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import os
|
||||
import sys
|
||||
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("--params", required=True)
|
||||
parser.add_argument("--output", required=True)
|
||||
parser.add_argument("--rank", type=int, required=True)
|
||||
parser.add_argument("--world-size", type=int, required=True)
|
||||
parser.add_argument("--local-steps", type=int, required=True)
|
||||
parser.add_argument("--lr", type=float, default=0.01)
|
||||
parser.add_argument("--batch-size", type=int, default=64)
|
||||
args = parser.parse_args()
|
||||
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
|
||||
# import job module from same directory
|
||||
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
|
||||
from job_module import make_dataloader, make_model
|
||||
|
||||
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
||||
|
||||
model = make_model().to(device)
|
||||
model.load_state_dict(
|
||||
torch.load(args.params, map_location=device, weights_only=True)
|
||||
)
|
||||
model.train()
|
||||
|
||||
loader = make_dataloader(args.rank, args.world_size, args.batch_size)
|
||||
opt = torch.optim.SGD(model.parameters(), lr=args.lr)
|
||||
|
||||
step = 0
|
||||
while step < args.local_steps:
|
||||
for x, y in loader:
|
||||
if step >= 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()
|
||||
Loading…
Reference in a new issue