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