| """One worker thread per GPU, pulling jobs off the shared queue.""" |
|
|
| import threading |
| import time |
| import traceback |
|
|
| from . import db |
| from .trainer import run_training |
|
|
| _cancel = set() |
| _cancel_lock = threading.Lock() |
| _workers = [] |
|
|
|
|
| def request_cancel(job_id): |
| with _cancel_lock: |
| _cancel.add(job_id) |
| j = db.get_job(job_id) |
| if j and j["status"] == "queued": |
| db.update_job(job_id, status="cancelled", finished_at=time.time()) |
| return True |
| return bool(j) |
|
|
|
|
| def _is_cancelled(job_id): |
| with _cancel_lock: |
| return job_id in _cancel |
|
|
|
|
| _devices_cache = None |
|
|
|
|
| def devices(): |
| """Never raises and never blocks twice — CUDA probing is done once.""" |
| global _devices_cache |
| if _devices_cache is not None: |
| return _devices_cache |
| found = ["cpu"] |
| try: |
| import torch |
| if torch.cuda.is_available() and torch.cuda.device_count() > 0: |
| found = [f"cuda:{i}" for i in range(torch.cuda.device_count())] |
| except Exception as exc: |
| print(f"[worker] no CUDA ({exc!r}) — falling back to CPU", flush=True) |
| _devices_cache = found |
| return found |
|
|
|
|
| def _loop(gpu_index, device, tokenizer_dir): |
| while True: |
| try: |
| job = db.claim_next_queued(gpu_index) |
| except Exception: |
| job = None |
| if not job: |
| time.sleep(2) |
| continue |
|
|
| jid = job["id"] |
|
|
| def log(msg, _jid=jid, _dev=device): |
| print(f"[{_dev}][{_jid}] {msg}", flush=True) |
| try: |
| db.add_log(_jid, msg) |
| except Exception: |
| pass |
|
|
| |
| |
| try: |
| from . import data as ndata |
| need = min(int(job["target_tokens"]), ndata.TARGET_TOKENS) |
| last_note = 0.0 |
| while ndata.available_tokens() < need and ndata.state()["status"] == "building": |
| if time.time() - last_note > 60: |
| log(f"waiting for the fineweb-edu cache to reach " |
| f"{need / 1e6:.0f}M tokens " |
| f"({ndata.available_tokens() / 1e6:.0f}M ready so far — " |
| f"this is a one-time cost per container)") |
| last_note = time.time() |
| if _is_cancelled(jid): |
| break |
| time.sleep(5) |
| except Exception: |
| pass |
|
|
| if _is_cancelled(jid): |
| db.update_job(jid, status="cancelled", finished_at=time.time()) |
| with _cancel_lock: |
| _cancel.discard(jid) |
| continue |
|
|
| log(f"starting on {device}") |
| try: |
| run_training(job, device, tokenizer_dir, lambda: _is_cancelled(jid), log) |
| except Exception as exc: |
| tb = traceback.format_exc(limit=6) |
| print(tb, flush=True) |
| log(f"FAILED: {exc!r}") |
| db.update_job(jid, status="failed", error=str(exc)[:500], |
| finished_at=time.time()) |
| finally: |
| with _cancel_lock: |
| _cancel.discard(jid) |
| try: |
| import torch |
| if device.startswith("cuda"): |
| torch.cuda.empty_cache() |
| except Exception: |
| pass |
|
|
|
|
| def start_workers(tokenizer_dir): |
| global _workers |
| if _workers: |
| return _workers |
| for i, dev in enumerate(devices()): |
| if dev.startswith("cuda"): |
| try: |
| import torch |
| name = torch.cuda.get_device_name(i) |
| gb = torch.cuda.get_device_properties(i).total_memory / 1024**3 |
| print(f"[worker] {dev}: {name} ({gb:.0f} GB)", flush=True) |
| except Exception: |
| pass |
| t = threading.Thread(target=_loop, args=(i, dev, tokenizer_dir), |
| name=f"worker-{dev}", daemon=True) |
| t.start() |
| _workers.append(t) |
| print(f"[worker] started on {dev}", flush=True) |
| return _workers |
|
|
|
|
| def worker_info(): |
| devs = devices() |
| running = {j["gpu_index"]: j for j in db.list_jobs(status=["running"], limit=64)} |
| return [{"index": i, "device": d, "job": running.get(i)} for i, d in enumerate(devs)] |
|
|
