From e415b7beb7051de98cc304c9182dc9774601fe70 Mon Sep 17 00:00:00 2001 From: muttakin Date: Tue, 28 Jul 2026 14:16:36 +0600 Subject: [PATCH] fault recovery and watchdog service --- Object_store/init_buckets.py | 3 +- Scheduler/app/main.py | 27 +++- Scheduler/app/models/job_model.py | 7 +- Scheduler/app/services/job_service.py | 111 +++++++++++-- Scheduler/app/services/watchdog_service.py | 86 ++++++++++ Worker/worker.py | 175 +++++++++++++++++++-- 6 files changed, 382 insertions(+), 27 deletions(-) create mode 100644 Scheduler/app/services/watchdog_service.py diff --git a/Object_store/init_buckets.py b/Object_store/init_buckets.py index 79aa1f09..4da6688d 100644 --- a/Object_store/init_buckets.py +++ b/Object_store/init_buckets.py @@ -11,6 +11,7 @@ BUCKETS = ( os.environ.get("OBJECT_STORE_BUCKET", "uploads"), os.environ.get("OBJECT_OUTPUT_BUCKET", "outputs"), + os.environ.get("OBJECT_STORE_CHECKPOINT_BUCKET", "checkpoints"), ) @@ -40,4 +41,4 @@ def ensure_buckets(client: Minio | None = None) -> None: if __name__ == "__main__": - ensure_buckets() + ensure_buckets() \ No newline at end of file diff --git a/Scheduler/app/main.py b/Scheduler/app/main.py index f8be7512..e27324f2 100644 --- a/Scheduler/app/main.py +++ b/Scheduler/app/main.py @@ -1,3 +1,7 @@ +import asyncio +import logging +from contextlib import asynccontextmanager + from fastapi import FastAPI from app.db.database import Base, engine @@ -5,6 +9,25 @@ from app.api.jobs_route import router as jobs_router from app.api.scheduler_route import router as scheduler_router from app.api.worker_route import router as workers_router # include worker registration +from app.services.watchdog_service import run_watchdog_loop + + +@asynccontextmanager +async def lifespan(app: FastAPI): + # Create all tables (for development; in production use Alembic migrations) + Base.metadata.create_all(bind=engine) + + watchdog_task = asyncio.create_task(run_watchdog_loop()) + try: + yield + finally: + watchdog_task.cancel() + try: + await watchdog_task + except asyncio.CancelledError: + pass + + # Create FastAPI app app = FastAPI( @@ -16,8 +39,8 @@ openapi_url="/openapi.json", ) -# Create all tables (for development; in production use Alembic migrations) -Base.metadata.create_all(bind=engine) +# # Create all tables (for development; in production use Alembic migrations) +# Base.metadata.create_all(bind=engine) # Allow the Vite development server to request the API directly. from fastapi.middleware.cors import CORSMiddleware diff --git a/Scheduler/app/models/job_model.py b/Scheduler/app/models/job_model.py index 41a8d23c..e0a1c665 100644 --- a/Scheduler/app/models/job_model.py +++ b/Scheduler/app/models/job_model.py @@ -1,4 +1,4 @@ -from sqlalchemy import Column, String, DateTime, Float, Enum, JSON +from sqlalchemy import Column, String, DateTime, Float, Integer, Enum, JSON from sqlalchemy.sql import func import uuid from app.db.database import Base @@ -10,6 +10,7 @@ class JobStatus(enum.Enum): VRAM_ESTIMATION_PENDING = "VRAM_ESTIMATION_PENDING" RUNNABLE = "RUNNABLE" IN_PROGRESS = "IN_PROGRESS" + RETRY_PENDING = "RETRY_PENDING" # worker died mid-training; needs reassignment COMPLETED = "COMPLETED" FAILED = "FAILED" @@ -32,5 +33,9 @@ class Job(Base): vram_required = Column(Float, nullable=True) # in GB step_time = Column(Float, nullable=True) # in seconds per step + # Fault recovery + assigned_worker_id = Column(String, nullable=True) # worker currently (or last) running this job + retry_count = Column(Integer, nullable=False, default=0) # number of times reassigned after worker death + created_at = Column(DateTime(timezone=True), server_default=func.now()) updated_at = Column(DateTime(timezone=True), onupdate=func.now()) \ No newline at end of file diff --git a/Scheduler/app/services/job_service.py b/Scheduler/app/services/job_service.py index eef3c3a4..7e997cf0 100644 --- a/Scheduler/app/services/job_service.py +++ b/Scheduler/app/services/job_service.py @@ -3,6 +3,7 @@ from sqlalchemy import or_ from app.models.job_model import Job, JobStatus from app.schemas.worker_schema import WorkerResource +import os def create_job(db: Session, job_data: dict): @@ -164,15 +165,15 @@ async def _check_vram_estimation_strategy(db: Session, request: WorkerResource) return _format_job_response(job, flag="vram_estimation") -def _check_training_job_strategy(db: Session, request: WorkerResource) -> dict | None: - """ - Strategy 2: Find runnable training job where (vram_required + 1.0) <= available vram of pulling worker. - Selects the job with largest vram_required. - """ - job = ( +def _find_matching_job(db: Session, status: JobStatus, request: WorkerResource): + """Find the best job in `status` whose vram_required (+1.0 GB safety margin) + fits the pulling worker's free VRAM. Jobs with unknown vram_required (not yet + estimated) are also eligible. Prefers the largest vram_required first (best + packing), then oldest first.""" + return ( db.query(Job) .filter( - Job.status == JobStatus.RUNNABLE, + Job.status == status, or_( Job.vram_required.is_(None), (Job.vram_required + 1.0) <= request.free_vram, @@ -182,19 +183,51 @@ def _check_training_job_strategy(db: Session, request: WorkerResource) -> dict | .first() ) - if not job: - return None +def _assign_job_to_worker(db: Session, job: Job, worker_id: str, flag: str) -> dict: job.status = JobStatus.IN_PROGRESS + job.assigned_worker_id = worker_id db.commit() db.refresh(job) + return _format_job_response(job, flag=flag) + + +def _check_training_job_strategy(db: Session, request: WorkerResource) -> dict | None: + """ + Strategy 2: Find a fresh RUNNABLE job (never started before) that fits. + """ + job = _find_matching_job(db, JobStatus.RUNNABLE, request) + if not job: + return None - return _format_job_response(job, flag="training") + return _assign_job_to_worker(db, job, request.worker_id, flag="training") + + +def _check_retry_job_strategy(db: Session, request: WorkerResource) -> dict | None: + """ + Strategy 3: Find a job whose previous worker died mid-training + (RETRY_PENDING, set by the watchdog). The worker resumes it from its last + checkpoint, so it's matched using the same VRAM-fit logic as a fresh job. + Placed ahead of fresh jobs in SCHEDULING_STRATEGIES since it represents + work already in flight. + """ + job = _find_matching_job(db, JobStatus.RETRY_PENDING, request) + if not job: + return None + + return _assign_job_to_worker(db, job, request.worker_id, flag="retry") # List of scheduling strategies in priority order. Easy to extend with new strategies. +# Order rationale: +# 1. VRAM estimation jobs are cheap, quick, and unblock scheduling for everything +# else, so they always jump the queue. +# 2. Retries represent work that's already partially done (has a checkpoint) and +# whose recovery is time-sensitive, so they're preferred over brand-new jobs. +# 3. Fresh jobs run last. SCHEDULING_STRATEGIES = [ _check_vram_estimation_strategy, + _check_retry_job_strategy, _check_training_job_strategy, ] @@ -235,3 +268,61 @@ def set_to_completed(db: Session, job_id: str): db.refresh(job) return job + + +# --------------------------------------------------------------------------- +# Fault recovery (watchdog support) +# --------------------------------------------------------------------------- + +MAX_JOB_RETRIES = int(os.getenv("MAX_JOB_RETRIES", "3")) + + +def get_in_progress_assignments(db: Session) -> list[tuple[str, str]]: + """Return (job_id, assigned_worker_id) for every job currently IN_PROGRESS. + Used by the watchdog to check which of these workers are still alive.""" + rows = ( + db.query(Job.id, Job.assigned_worker_id) + .filter( + Job.status == JobStatus.IN_PROGRESS, + Job.assigned_worker_id.isnot(None), + ) + .all() + ) + return [(row[0], row[1]) for row in rows] + + +def requeue_job_after_worker_death(db: Session, job_id: str, dead_worker_id: str) -> Job | None: + """Called by the watchdog when a job's assigned worker has stopped sending + heartbeats. Re-checks the job is still assigned to that same worker (avoids + racing a legitimate completion/reassignment that happened concurrently), + then either requeues it as RETRY_PENDING (picked up again via the retry + scheduling strategy, resuming from its last checkpoint) or marks it FAILED + if it has already exhausted MAX_JOB_RETRIES -- guards against a job that + keeps crashing workers (e.g. a checkpoint that itself corrupts) looping + forever. + """ + job = ( + db.query(Job) + .filter( + Job.id == job_id, + Job.status == JobStatus.IN_PROGRESS, + Job.assigned_worker_id == dead_worker_id, + ) + .first() + ) + if not job: + # Job already moved on (completed / reassigned) since the watchdog + # took its snapshot -- nothing to do. + return None + + job.assigned_worker_id = None + + if job.retry_count >= MAX_JOB_RETRIES: + job.status = JobStatus.FAILED + else: + job.retry_count += 1 + job.status = JobStatus.RETRY_PENDING + + db.commit() + db.refresh(job) + return job \ No newline at end of file diff --git a/Scheduler/app/services/watchdog_service.py b/Scheduler/app/services/watchdog_service.py new file mode 100644 index 00000000..bcd5b6a2 --- /dev/null +++ b/Scheduler/app/services/watchdog_service.py @@ -0,0 +1,86 @@ +""" +Fault-detection watchdog. + +Design: polling over Redis keyspace-notifications. +----------------------------------------------------- +Redis *could* push an "expired" pub/sub event the instant a worker's heartbeat +key times out (`notify-keyspace-events Ex`), which sounds more elegant. It's +deliberately not used here: + - it requires non-default Redis server config, which is easy to forget when + the deployment/infra changes (docker-compose image swap, managed Redis, etc.) + - pub/sub delivery isn't durable -- if the scheduler process is briefly down + or the subscriber connection drops, expiry events are lost forever and the + job would hang in IN_PROGRESS with nobody watching it. + - it needs its own reconnect/backoff logic to be reliable, at which point + it's more moving parts than a poll loop for the same guarantee. + +A simple poll loop is stateless and self-healing: every tick it recomputes +"which IN_PROGRESS jobs have a worker that's actually still alive" from +scratch, straight from Redis. A missed tick just means one extra +WATCHDOG_INTERVAL of delay before a dead job is noticed, not a permanently +stuck job. +""" + +import asyncio +import logging +import os + +from app.core.redis import redis_client +from app.db.database import SessionLocal +from app.services.job_service import get_in_progress_assignments, requeue_job_after_worker_death + +logger = logging.getLogger("watchdog") + +WATCHDOG_INTERVAL = int(os.getenv("WATCHDOG_INTERVAL", "5")) # seconds + + +async def _sweep_once(): + db = SessionLocal() + try: + assignments = get_in_progress_assignments(db) + if not assignments: + return + + # One round-trip for all worker keys instead of one per job. + worker_ids = {worker_id for _job_id, worker_id in assignments} + pipe = redis_client.pipeline() + for worker_id in worker_ids: + pipe.exists(f"worker:{worker_id}") + alive_flags = await pipe.execute() + alive = { + worker_id: bool(flag) + for worker_id, flag in zip(worker_ids, alive_flags) + } + + for job_id, worker_id in assignments: + if alive.get(worker_id): + continue # heartbeat still active, job is fine + + job = requeue_job_after_worker_death(db, job_id, worker_id) + if job is None: + continue # already moved on (completed/reassigned) concurrently + + if job.status.value == "FAILED": + logger.error( + "Job %s exhausted retries after worker %s died; marking FAILED.", + job_id, worker_id, + ) + else: + logger.warning( + "Worker %s missed its heartbeat while running job %s " + "(attempt %d/%s) -- requeued as RETRY_PENDING.", + worker_id, job_id, job.retry_count, + os.getenv("MAX_JOB_RETRIES", "3"), + ) + except Exception: + # Never let one bad sweep kill the loop -- log and try again next tick. + logger.exception("Watchdog sweep failed") + finally: + db.close() + + +async def run_watchdog_loop(): + logger.info("Fault-recovery watchdog started (interval=%ss).", WATCHDOG_INTERVAL) + while True: + await _sweep_once() + await asyncio.sleep(WATCHDOG_INTERVAL) \ No newline at end of file diff --git a/Worker/worker.py b/Worker/worker.py index 0e07265a..b9e23fbd 100644 --- a/Worker/worker.py +++ b/Worker/worker.py @@ -3,7 +3,10 @@ import uuid import json import shlex +import shutil import logging +import zipfile +import threading import subprocess import tempfile import requests @@ -33,6 +36,13 @@ VRAM_ESTIMATION_SCRIPT = os.path.join(os.path.dirname(os.path.abspath(__file__)), "vram_estimation.py") os.makedirs(OUTPUT_DIR, exist_ok=True) +# Object store (checkpoint persistence) +OBJECT_STORE_URL = os.getenv("OBJECT_STORE_URL", "http://localhost:8010").rstrip("/") +CHECKPOINT_BUCKET = os.getenv("OBJECT_STORE_CHECKPOINT_BUCKET", "checkpoints") +CHECKPOINT_UPLOAD_URL = f"{OBJECT_STORE_URL}/objects/upload" +CHECKPOINT_CONTAINER_PATH = "/checkpoints" # contract: user code checkpoints here +CHECKPOINT_SYNC_INTERVAL = int(os.getenv("CHECKPOINT_SYNC_INTERVAL", "30")) # seconds + # Logger setup logging.basicConfig( level=logging.INFO, @@ -207,17 +217,149 @@ def handle_vram_estimation(job_id: str, image_name: str, command: str): save_vram_estimation(job_id, report) -def handle_training(job_id: str, image_name: str): - """Handle training job.""" - logger.info("Training job received for job %s.", job_id) +def _checkpoint_object_key(job_id: str) -> str: + return f"{job_id}/checkpoint.zip" + + +def _zip_dir(src_dir: str, zip_path: str): + """Zip the contents of src_dir (relative paths, no top-level folder) into zip_path.""" + with zipfile.ZipFile(zip_path, "w", zipfile.ZIP_DEFLATED) as zf: + for root, _dirs, files in os.walk(src_dir): + for name in files: + file_path = os.path.join(root, name) + arcname = os.path.relpath(file_path, src_dir) + zf.write(file_path, arcname) - cmd = ["docker", "run", "--rm", "--gpus", "all", image_name] - logger.info("Running training container with command: %s", " ".join(cmd)) + +def upload_checkpoint(job_id: str, checkpoint_dir: str) -> bool: + """Zip the host-side checkpoint dir and push it to the object store, overwriting + any previous checkpoint for this job (object key is stable per job_id).""" + if not os.path.isdir(checkpoint_dir) or not os.listdir(checkpoint_dir): + return False + + with tempfile.NamedTemporaryFile(suffix=".zip", delete=False) as tmp: + zip_path = tmp.name try: - res = subprocess.run(cmd, capture_output=True, text=True) - log_file = os.path.join(OUTPUT_DIR, f"{job_id}.txt") + _zip_dir(checkpoint_dir, zip_path) + object_key = _checkpoint_object_key(job_id) + with open(zip_path, "rb") as f: + files = {"file": ("checkpoint.zip", f, "application/zip")} + data = {"bucket": CHECKPOINT_BUCKET, "object_key": object_key} + resp = requests.post(CHECKPOINT_UPLOAD_URL, data=data, files=files, timeout=60) + if resp.status_code >= 400: + logger.error( + "Checkpoint upload failed for job %s: %s %s", job_id, resp.status_code, resp.text + ) + return False + logger.info("Checkpoint synced for job %s.", job_id) + return True + except Exception as e: + logger.error("Checkpoint upload error for job %s: %s", job_id, e) + return False + finally: + try: + os.remove(zip_path) + except OSError: + pass + + +def download_checkpoint(job_id: str, dest_dir: str) -> bool: + """Fetch the latest checkpoint zip for job_id from the object store and extract + it into dest_dir. Returns False (no error) if no checkpoint exists yet.""" + object_key = _checkpoint_object_key(job_id) + url = f"{OBJECT_STORE_URL}/objects/{CHECKPOINT_BUCKET}/{object_key}" + try: + resp = requests.get(url, timeout=60) + if resp.status_code == 404: + logger.warning("No existing checkpoint found for job %s.", job_id) + return False + resp.raise_for_status() + except Exception as e: + logger.error("Failed to download checkpoint for job %s: %s", job_id, e) + return False + with tempfile.NamedTemporaryFile(suffix=".zip", delete=False) as tmp: + tmp.write(resp.content) + zip_path = tmp.name + + try: + os.makedirs(dest_dir, exist_ok=True) + with zipfile.ZipFile(zip_path, "r") as zf: + zf.extractall(dest_dir) + logger.info("Checkpoint restored for job %s into %s.", job_id, dest_dir) + return True + except Exception as e: + logger.error("Failed to extract checkpoint for job %s: %s", job_id, e) + return False + finally: + try: + os.remove(zip_path) + except OSError: + pass + + +class _CheckpointSyncer: + """Background thread that periodically uploads a host checkpoint dir while a + training container is running. Started once training begins, stopped (with one + final sync) once the container exits, so we never lose more than one interval + of progress if the worker dies mid-run.""" + + def __init__(self, job_id: str, checkpoint_dir: str, interval: int = CHECKPOINT_SYNC_INTERVAL): + self.job_id = job_id + self.checkpoint_dir = checkpoint_dir + self.interval = interval + self._stop_event = threading.Event() + self._thread = threading.Thread(target=self._run, daemon=True) + + def _run(self): + while not self._stop_event.wait(self.interval): + upload_checkpoint(self.job_id, self.checkpoint_dir) + + def start(self): + self._thread.start() + + def stop_and_final_sync(self): + self._stop_event.set() + self._thread.join() + upload_checkpoint(self.job_id, self.checkpoint_dir) + + +def _run_training_container(job_id: str, image_name: str, resume: bool): + """Shared execution path for fresh training runs and retries. Mounts a host + checkpoint dir into the container at CHECKPOINT_CONTAINER_PATH, optionally + pre-populating it from the last known checkpoint, and periodically syncs it + back to the object store while the container runs so a mid-training worker + failure loses at most CHECKPOINT_SYNC_INTERVAL seconds of progress. + """ + with tempfile.TemporaryDirectory(prefix=f"ckpt_{job_id}_") as checkpoint_dir: + if resume: + download_checkpoint(job_id, checkpoint_dir) + + cmd = [ + "docker", "run", "--rm", "--gpus", "all", + "-v", f"{checkpoint_dir}:{CHECKPOINT_CONTAINER_PATH}", + image_name, + ] + logger.info( + "Running %s training container for job %s: %s", + "resumed" if resume else "fresh", job_id, " ".join(cmd), + ) + + syncer = _CheckpointSyncer(job_id, checkpoint_dir) + syncer.start() + + try: + res = subprocess.run(cmd, capture_output=True, text=True) + except Exception as e: + logger.error("Execution error for job %s: %s", job_id, e) + syncer.stop_and_final_sync() + return + finally: + if syncer._thread.is_alive(): + syncer.stop_and_final_sync() + + log_file = os.path.join(OUTPUT_DIR, f"{job_id}.txt") with open(log_file, "w", encoding="utf-8") as f: f.write(res.stdout + "\n" + res.stderr) @@ -228,13 +370,20 @@ def handle_training(job_id: str, image_name: str): upload_output_file(log_file, job_id) - except Exception as e: - logger.error("Execution error for job %s: %s", job_id, e) + +def handle_training(job_id: str, image_name: str): + """Handle a fresh (first-attempt) training job.""" + logger.info("Training job received for job %s.", job_id) + _run_training_container(job_id, image_name, resume=False) -def handle_retry(job_id: str): - """Handle retry job.""" +def handle_retry(job_id: str, image_name: str): + """Handle a retry job: a previous worker died mid-training (missed heartbeats), + and the scheduler has re-assigned this job to us. Resume from the last + checkpoint synced to the object store instead of starting from scratch. + """ logger.info("Retry job received for job %s.", job_id) + _run_training_container(job_id, image_name, resume=True) def process_job(): @@ -256,7 +405,7 @@ def process_job(): elif flag == "training": handle_training(job_id, image_name) elif flag == "retry": - handle_retry(job_id) + handle_retry(job_id, image_name) else: logger.warning("Unknown job flag '%s' for job %s.", flag, job_id) @@ -280,4 +429,4 @@ def process_job(): except Exception as e: logger.error("Error processing job: %s", e) - time.sleep(JOB_POLL_INTERVAL) + time.sleep(JOB_POLL_INTERVAL) \ No newline at end of file