From c33e3952a23dd1850f227a6c844e7c1a6005ac83 Mon Sep 17 00:00:00 2001 From: badtst Date: Fri, 31 Jul 2026 13:08:27 +0000 Subject: [PATCH] Init --- tests/scheduler/test_session_watchdog.py | 107 +++++++++++++++++++++++ warden/api/routes/sessions.py | 29 +----- warden/lib/config/config.py | 2 + warden/lib/config/config.sample.yaml | 7 ++ warden/lib/sessions.py | 42 +++++++++ warden/scheduler/main.py | 32 +++++-- warden/scheduler/session_watchdog.py | 73 ++++++++++++++++ 7 files changed, 261 insertions(+), 31 deletions(-) create mode 100644 tests/scheduler/test_session_watchdog.py create mode 100644 warden/lib/sessions.py create mode 100644 warden/scheduler/session_watchdog.py diff --git a/tests/scheduler/test_session_watchdog.py b/tests/scheduler/test_session_watchdog.py new file mode 100644 index 0000000..c7b79f1 --- /dev/null +++ b/tests/scheduler/test_session_watchdog.py @@ -0,0 +1,107 @@ +"""Testing warden.scheduler.session_watchdog""" + +import asyncio +import contextlib +from datetime import datetime, timedelta, timezone + +import pytest +from sqlalchemy.ext.asyncio import async_sessionmaker + +from warden.lib.config import Config, SchedulerConfig +from warden.lib.models import Job, Session +from warden.scheduler.session_watchdog import session_watchdog + +OLD = datetime.now(timezone.utc) - timedelta(hours=1) + + +def build_watchdog_conf(session_idle_timeout_s: float) -> Config: + return Config( + scheduler=SchedulerConfig( + db_polling_interval_s=0.01, + session_idle_timeout_s=session_idle_timeout_s, + ), + ) + + +async def run_briefly(conf: Config, db_session_maker: async_sessionmaker) -> None: + """Run the watchdog for a few polling intervals, then stop it.""" + task = asyncio.create_task(session_watchdog(conf, db_session_maker)) + try: + await asyncio.sleep(0.2) + finally: + task.cancel() + with contextlib.suppress(asyncio.CancelledError): + await task + + +@pytest.mark.asyncio +async def test_session_watchdog_revokes_idle_session_and_cancels_its_jobs( + db_session_maker: async_sessionmaker, +): + """A session with no recent job activity gets revoked, and its pending job canceled.""" + session_record = Session(user_id="1234", slurm_job_id="1", created_at=OLD) + stale_job = Job( + sequence="{}", + shots=100, + status="PENDING", + session=session_record, + created_at=OLD, + ) + + async with db_session_maker() as session: + session.add_all([session_record, stale_job]) + await session.commit() + + conf = build_watchdog_conf(session_idle_timeout_s=0.05) + await run_briefly(conf, db_session_maker) + + async with db_session_maker() as session: + refreshed_session = await session.get(Session, session_record.id) + refreshed_job = await session.get(Job, stale_job.id) + assert refreshed_session.revoked_at is not None + assert refreshed_job.canceled_at is not None + assert refreshed_job.status == "CANCELED" + + +@pytest.mark.asyncio +async def test_session_watchdog_does_not_revoke_session_with_recent_job( + db_session_maker: async_sessionmaker, +): + """A session whose most recent job is within the idle window stays untouched.""" + session_record = Session(user_id="1234", slurm_job_id="1", created_at=OLD) + recent_job = Job( + sequence="{}", + shots=100, + status="PENDING", + session=session_record, + ) + + async with db_session_maker() as session: + session.add_all([session_record, recent_job]) + await session.commit() + + conf = build_watchdog_conf(session_idle_timeout_s=3600) + await run_briefly(conf, db_session_maker) + + async with db_session_maker() as session: + refreshed_session = await session.get(Session, session_record.id) + assert refreshed_session.revoked_at is None + + +@pytest.mark.asyncio +async def test_session_watchdog_disabled_when_timeout_negative( + db_session_maker: async_sessionmaker, +): + """session_idle_timeout_s = -1 disables the watchdog entirely.""" + session_record = Session(user_id="1234", slurm_job_id="1", created_at=OLD) + + async with db_session_maker() as session: + session.add(session_record) + await session.commit() + + conf = build_watchdog_conf(session_idle_timeout_s=-1) + await run_briefly(conf, db_session_maker) + + async with db_session_maker() as session: + refreshed_session = await session.get(Session, session_record.id) + assert refreshed_session.revoked_at is None diff --git a/warden/api/routes/sessions.py b/warden/api/routes/sessions.py index af1c2e3..e995226 100644 --- a/warden/api/routes/sessions.py +++ b/warden/api/routes/sessions.py @@ -1,4 +1,3 @@ -from datetime import datetime, timezone from logging import getLogger from fastapi import APIRouter, HTTPException @@ -12,7 +11,8 @@ ) from warden.api.routes.dependencies.db import DBSessionDep from warden.api.schemas.sessions import CreateSession, SessionResponse -from warden.lib.models import Job, Session +from warden.lib.models import Session +from warden.lib.sessions import revoke_session_and_cancel_jobs logger = getLogger(__name__) router = APIRouter(prefix="/sessions") @@ -46,30 +46,7 @@ async def revoke_session( session_record = result.scalar_one_or_none() if session_record is None: raise HTTPException(status_code=404, detail="Session not found.") - session_record.revoked_at = datetime.now(timezone.utc) - await db_session.flush() - await db_session.commit() - async with db_session.begin(): - result = await db_session.execute( - select(Job) - .where( - Job.session_id == session_record.id, - Job.status.not_in(("ERROR", "DONE", "CANCELED")), - Job.canceled_at.is_(None), - ) - .with_for_update(of=Job) - ) - jobs_to_cancel = result.scalars() - for job in jobs_to_cancel: - logger.info( - "Cancelling job '%s' attached to session %s", job.id, session_record.id - ) - job.canceled_at = datetime.now(timezone.utc) - # Not yet started by the worker - if job.scheduled_at is None: - # Set job to cancel - job.status = "CANCELED" - # Releases nowait + await revoke_session_and_cancel_jobs(db_session, session_record) return SessionResponse.from_model(session_record) diff --git a/warden/lib/config/config.py b/warden/lib/config/config.py index a9fc30d..d4f3693 100644 --- a/warden/lib/config/config.py +++ b/warden/lib/config/config.py @@ -69,6 +69,8 @@ class SchedulerConfig(BaseSettings): job_polling_interval_s: float = 5 job_polling_timeout_s: float = -1 + session_idle_timeout_s: float = 3600 + class QPUConfig(BaseSettings): uri: str = "http://localhost:8000" diff --git a/warden/lib/config/config.sample.yaml b/warden/lib/config/config.sample.yaml index 0b35b20..8fd8cac 100644 --- a/warden/lib/config/config.sample.yaml +++ b/warden/lib/config/config.sample.yaml @@ -70,6 +70,13 @@ scheduler: # - Set to -1 for no time limit. job_polling_timeout_s: -1 + # Session watchdog policy + # Maximum time in seconds a session can go without a new job being created + # before it is automatically revoked (and its remaining non-terminal jobs + # canceled). + # - Set to -1 to disable the watchdog. + session_idle_timeout_s: 3600 + qpu: # Local Pasqal QPU API configuration uri: http://localhost:8000 diff --git a/warden/lib/sessions.py b/warden/lib/sessions.py new file mode 100644 index 0000000..480706f --- /dev/null +++ b/warden/lib/sessions.py @@ -0,0 +1,42 @@ +"""Session revocation logic shared by the API and the scheduler.""" + +from datetime import datetime, timezone +from logging import getLogger + +from sqlalchemy import select +from sqlalchemy.ext.asyncio import AsyncSession + +from warden.lib.models import Job, Session + +logger = getLogger(__name__) + + +async def revoke_session_and_cancel_jobs( + db_session: AsyncSession, session_record: Session +) -> None: + """Revoke a session and cancel every non-terminal job attached to it.""" + session_record.revoked_at = datetime.now(timezone.utc) + await db_session.flush() + await db_session.commit() + + async with db_session.begin(): + result = await db_session.execute( + select(Job) + .where( + Job.session_id == session_record.id, + Job.status.not_in(("ERROR", "DONE", "CANCELED")), + Job.canceled_at.is_(None), + ) + .with_for_update(of=Job) + ) + jobs_to_cancel = result.scalars() + for job in jobs_to_cancel: + logger.info( + "Cancelling job '%s' attached to session %s", job.id, session_record.id + ) + job.canceled_at = datetime.now(timezone.utc) + # Not yet started by the worker + if job.scheduled_at is None: + # Set job to cancel + job.status = "CANCELED" + # Releases nowait diff --git a/warden/scheduler/main.py b/warden/scheduler/main.py index 90079a4..8db8ce1 100644 --- a/warden/scheduler/main.py +++ b/warden/scheduler/main.py @@ -16,6 +16,7 @@ from warden.lib.models import Job from warden.scheduler.cancellation_worker import cancellation_worker from warden.scheduler.db import job_update_commiter +from warden.scheduler.session_watchdog import session_watchdog from warden.scheduler.strategy import schedulers from warden.scheduler.types import JobUpdateQueue from warden.scheduler.worker import LocalQPUWorker @@ -107,6 +108,19 @@ async def run_cancellation_worker(engine: AsyncEngine, conf: Config): await cancellation_worker(conf=conf, session_factory=session_factory) +async def run_session_watchdog(engine: AsyncEngine, conf: Config): + """Session watchdog main logic + + Runs the session-idle watchdog in an infinite loop. + Gets canceled by `main_async` when stop signal is received. + """ + logger.info("Session watchdog running.") + + session_factory = async_sessionmaker(bind=engine, expire_on_commit=False) + + await session_watchdog(conf=conf, session_factory=session_factory) + + async def shutdown(engine: AsyncEngine): """Cleanup tasks and close DB connections.""" @@ -137,25 +151,33 @@ async def main_async(conf: Config | None = None): try: logger.info( - "Starting scheduler and cancellation worker (Press Ctrl+C to exit)..." + "Starting scheduler, cancellation worker and session watchdog " + "(Press Ctrl+C to exit)..." ) - # Start both scheduler and cancellation worker as separate tasks with same lifetime + # Start scheduler, cancellation worker and session watchdog as separate + # tasks with the same lifetime scheduler_task = loop.create_task(run_scheduler(engine, conf), name="Scheduler") cancellation_task = loop.create_task( run_cancellation_worker(engine, conf), name="Cancellation Worker" ) + watchdog_task = loop.create_task( + run_session_watchdog(engine, conf), name="Session Watchdog" + ) # Wait for stop signal await stop_event.wait() - # Cancel both tasks - logger.info("Stopping scheduler and cancellation worker...") + # Cancel all tasks + logger.info("Stopping scheduler, cancellation worker and session watchdog...") scheduler_task.cancel() cancellation_task.cancel() + watchdog_task.cancel() # Wait for graceful shutdown - await asyncio.gather(scheduler_task, cancellation_task, return_exceptions=True) + await asyncio.gather( + scheduler_task, cancellation_task, watchdog_task, return_exceptions=True + ) finally: await shutdown(engine) diff --git a/warden/scheduler/session_watchdog.py b/warden/scheduler/session_watchdog.py new file mode 100644 index 0000000..5278d71 --- /dev/null +++ b/warden/scheduler/session_watchdog.py @@ -0,0 +1,73 @@ +"""Session-idle watchdog: revokes sessions with no new job in a configurable window.""" + +import asyncio +import logging +from datetime import datetime, timezone + +from sqlalchemy import func, select +from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker + +from warden.lib.config import Config +from warden.lib.models import Job, Session +from warden.lib.sessions import revoke_session_and_cancel_jobs + +logger = logging.getLogger(__name__) + + +async def session_watchdog( + conf: Config, session_factory: async_sessionmaker[AsyncSession] +) -> None: + """Revoke sessions that haven't had a new job created within the idle timeout. + + Infinite loop, meant to run alongside the scheduler and cancellation worker. + A `session_idle_timeout_s` of -1 disables the watchdog. + """ + sleep_interval = conf.scheduler.db_polling_interval_s + timeout_s = conf.scheduler.session_idle_timeout_s + + while True: + if timeout_s < 0: + await asyncio.sleep(sleep_interval) + continue + + async with session_factory() as db_session: + sessions = ( + ( + await db_session.execute( + select(Session).where(Session.revoked_at.is_(None)) + ) + ) + .scalars() + .all() + ) + + if sessions: + last_job_at_by_session = dict( + ( + await db_session.execute( + select(Job.session_id, func.max(Job.created_at)).group_by( + Job.session_id + ) + ) + ).all() + ) + + now = datetime.now(timezone.utc) + for session_record in sessions: + last_activity = last_job_at_by_session.get( + session_record.id, session_record.created_at + ) + if last_activity.tzinfo is None: + # ponytail: sqlite/mariadb drivers hand back naive + # datetimes for DateTime(timezone=True) columns even + # though every write here is UTC; normalize on read. + last_activity = last_activity.replace(tzinfo=timezone.utc) + if (now - last_activity).total_seconds() > timeout_s: + logger.info( + "Revoking session '%s': no new job for over %ss", + session_record.id, + timeout_s, + ) + await revoke_session_and_cancel_jobs(db_session, session_record) + + await asyncio.sleep(sleep_interval)