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