Add session lifetime expiry

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

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)

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"

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

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

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)

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)