fault recovery and watchdog service

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()

Scheduler/app/main.py

@@ -1,10 +1,33 @@
+import asyncio
+import logging
+from contextlib import asynccontextmanager
+
 from fastapi import FastAPI
 from app.db.database import Base, engine
 
 # Import routers
 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

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())

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

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)