diff --git a/tests/api/test_cancelation_race_condition.py b/tests/api/test_cancelation_race_condition.py new file mode 100644 index 0000000..80afa6e --- /dev/null +++ b/tests/api/test_cancelation_race_condition.py @@ -0,0 +1,271 @@ +""" +Testing possible race conditions between the api cancelling a job and the scheduler +""" + +import asyncio +from typing import cast + +import pytest +from sqlalchemy import Update, select, update +from sqlalchemy.ext.asyncio import AsyncSession + +from tests.api.conftest import mock_munge_auth +from warden.lib.config import SchedulerStrategy +from warden.lib.models.jobs import Job +from warden.lib.models.sessions import Session +from warden.scheduler.strategy import schedulers + +########################################################################## +####################### Repeated behavior testing ######################## +########################################################################## + + +@pytest.mark.asyncio +@pytest.mark.parametrize("strategy", list(SchedulerStrategy)) +async def test_repeated_job_cancel_and_scheduler_pick_race( + client, app, serialized_sequence: str, strategy +): + """Assert racing a real cancel against a real scheduler pick never lets + a job end up CANCELED and scheduled at once + + 1. Create a PENDING job for a given user + 2. Run the API cancel and the scheduler's `get_next_job` concurrently on + that job, with no control over which one the event loop runs first + 3. Assert the job never ends up with status CANCELED and a + `scheduled_at` set. Whoever claims the jobs first wins. + 4. Repeat over fresh jobs, since which side wins isn't controlled here + """ + user_id = 1000 + async_session = app.state.db_session_factory + scheduler = schedulers[strategy] + + for _ in range(40): + job = Job( + session=Session(user_id=str(user_id), slurm_job_id="1"), + sequence=serialized_sequence, + shots=100, + status="PENDING", + ) + async with async_session() as session: + session.add(job) + await session.commit() + await session.refresh(job) + job_id = job.id + + async def do_schedule(): + async with async_session() as session: + return await scheduler.get_next_job(session) + + with mock_munge_auth(app, uid=user_id): + cancel_response, claimed = await asyncio.gather( + client.post(f"/jobs/{job_id}/cancel"), do_schedule() + ) + + async with async_session() as session: + job = ( + await session.execute(select(Job).where(Job.id == job_id)) + ).scalar_one() + + # The property the atomic UPDATEs exist to guarantee: a job the + # scheduler is about to run can't silently flip to CANCELED, since + # the cancellation worker never looks at CANCELED jobs again. + assert not (job.status == "CANCELED" and job.scheduled_at is not None) + + if claimed is not None: + assert job.scheduled_at is not None + if cancel_response.status_code == 200: + assert job.canceled_at is not None + + # Terminal, so the next iteration's candidate pick can't see it. + async with async_session() as session: + await session.execute( + update(Job).where(Job.id == job_id).values(status="DONE") + ) + await session.commit() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("strategy", list(SchedulerStrategy)) +async def test_repeated_session_revoke_and_scheduler_pick_race( + client, app, serialized_sequence: str, strategy +): + """Assert racing a real session revoke against a real scheduler pick + never lets a job end up CANCELED and scheduled at once + """ + user_id = 1000 + async_session = app.state.db_session_factory + scheduler = schedulers[strategy] + + for _ in range(40): + session_record = Session(user_id=str(user_id), slurm_job_id="1") + job = Job( + session=session_record, + sequence=serialized_sequence, + shots=100, + status="PENDING", + ) + async with async_session() as session: + session.add(job) + await session.commit() + await session.refresh(job) + job_id = job.id + session_id = session_record.id + + async def do_schedule(): + async with async_session() as session: + return await scheduler.get_next_job(session) + + with mock_munge_auth(app, uid=0): + revoke_response, claimed = await asyncio.gather( + client.delete(f"/sessions/{session_id}"), do_schedule() + ) + + assert revoke_response.status_code == 200 + + async with async_session() as session: + job = ( + await session.execute(select(Job).where(Job.id == job_id)) + ).scalar_one() + + assert not (job.status == "CANCELED" and job.scheduled_at is not None) + + if claimed is not None: + assert job.scheduled_at is not None + + # Terminal, so the next iteration's candidate pick can't see it. + async with async_session() as session: + await session.execute( + update(Job).where(Job.id == job_id).values(status="DONE") + ) + await session.commit() + + +########################################################################## +############################# Timed testing ############################# +########################################################################## + + +class _HookedUpdateDBSession: + """ + Session proxy running `hook` once, right before the first `UPDATE`. + """ + + def __init__(self, session: AsyncSession, hook): + self._session = session + self._hook = hook + + def __getattr__(self, name): + return getattr(self._session, name) + + async def execute(self, statement, *args, **kwargs): + if self._hook is not None and isinstance(statement, Update): + hook, self._hook = self._hook, None + await hook() + return await self._session.execute(statement, *args, **kwargs) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("strategy", list(SchedulerStrategy)) +async def test_cancel_racing_scheduler_pick_is_not_claimed( + client, app, serialized_sequence: str, strategy +): + """Assert a cancel committed mid-pick stops the scheduler claiming the job + + 1. Create a PENDING job for a given user + 2. Run the scheduler, cancelling the job via the API in the pick -> claim + window (i.e. after the candidate is chosen, before it is claimed) + 3. Assert the scheduler claims nothing: its UPDATE re-checks `status` + against the live row, which is now CANCELED + 4. Assert the job is CANCELED and never got a `scheduled_at` + + A job that is both CANCELED and scheduled is the state this must never + reach: the cancellation worker only picks up PENDING/RUNNING jobs, so the + QPU run would keep going with nothing left to stop it. + """ + user_id = 1000 + job = Job( + session=Session(user_id=str(user_id), slurm_job_id="1"), + sequence=serialized_sequence, + shots=100, + status="PENDING", + ) + async_session = app.state.db_session_factory + + async with async_session() as session: + session.add(job) + await session.commit() + await session.refresh(job) + + job_id = job.id + + async def cancel_mid_pick(): + with mock_munge_auth(app, uid=user_id): + response = await client.post(f"/jobs/{job_id}/cancel") + assert response.status_code == 200 + assert response.json()["status"] == "CANCELED" + + scheduler = schedulers[strategy] + async with async_session() as session: + claimed = await scheduler.get_next_job( + cast(AsyncSession, _HookedUpdateDBSession(session, cancel_mid_pick)) + ) + + assert claimed is None + + async with async_session() as session: + job = (await session.execute(select(Job).where(Job.id == job_id))).scalar_one() + + assert job.status == "CANCELED" + assert job.canceled_at is not None + assert job.scheduled_at is None + + +@pytest.mark.asyncio +@pytest.mark.parametrize("strategy", list(SchedulerStrategy)) +async def test_session_revoke_racing_scheduler_pick_is_not_claimed( + client, app, serialized_sequence: str, strategy +): + """Assert a session revoke committed mid-pick stops the scheduler + claiming the job + + Same forced interleaving as `test_cancel_racing_scheduler_pick_is_not_claimed`, + but through `DELETE /sessions/{id}` (bulk-cancels a session's jobs) + instead of `POST /jobs/{id}/cancel`. + """ + user_id = 1000 + session_record = Session(user_id=str(user_id), slurm_job_id="1") + job = Job( + session=session_record, + sequence=serialized_sequence, + shots=100, + status="PENDING", + ) + async_session = app.state.db_session_factory + + async with async_session() as session: + session.add(job) + await session.commit() + await session.refresh(job) + + job_id = job.id + session_id = session_record.id + + async def revoke_mid_pick(): + with mock_munge_auth(app, uid=0): + response = await client.delete(f"/sessions/{session_id}") + assert response.status_code == 200 + + scheduler = schedulers[strategy] + async with async_session() as session: + claimed = await scheduler.get_next_job( + cast(AsyncSession, _HookedUpdateDBSession(session, revoke_mid_pick)) + ) + + assert claimed is None + + async with async_session() as session: + job = (await session.execute(select(Job).where(Job.id == job_id))).scalar_one() + + assert job.status == "CANCELED" + assert job.canceled_at is not None + assert job.scheduled_at is None diff --git a/tests/api/test_jobs.py b/tests/api/test_jobs.py index 6c56d98..d735300 100644 --- a/tests/api/test_jobs.py +++ b/tests/api/test_jobs.py @@ -1,3 +1,4 @@ +import asyncio import json from datetime import datetime @@ -546,3 +547,44 @@ async def test_cancel_job_twice(client, app, serialized_sequence: str): with mock_munge_auth(app, uid=user_id): response = await client.post(f"/jobs/{job.id}/cancel") assert response.status_code == 409 + + +@pytest.mark.asyncio +async def test_concurrent_cancels_have_a_single_winner( + client, app, serialized_sequence: str +): + """Assert concurrent cancels of the same job produce exactly one cancelation + + 1. Create a PENDING job in db for a given user + 2. Call POST /job/id/cancel four times concurrently + 3. Assert exactly one request gets a 200 and the rest get a 409 + 4. Assert the job is canceled once in DB + """ + user_id = 1000 + job = Job( + session=Session(user_id=str(user_id), slurm_job_id="1"), + sequence=serialized_sequence, + shots=100, + status="PENDING", + ) + async_session = app.state.db_session_factory + + async with async_session() as session: + session.add(job) + await session.commit() + await session.refresh(job) + + job_id = job.id + + with mock_munge_auth(app, uid=user_id): + responses = await asyncio.gather( + *(client.post(f"/jobs/{job_id}/cancel") for _ in range(4)) + ) + + assert sorted(r.status_code for r in responses) == [200, 409, 409, 409] + + async with async_session() as session: + job = (await session.execute(select(Job).where(Job.id == job_id))).scalar_one() + + assert job.status == "CANCELED" + assert job.canceled_at is not None diff --git a/warden/api/routes/jobs.py b/warden/api/routes/jobs.py index 9c3e59c..41702c2 100644 --- a/warden/api/routes/jobs.py +++ b/warden/api/routes/jobs.py @@ -1,9 +1,10 @@ import asyncio from datetime import datetime, timezone from logging import getLogger +from typing import cast from fastapi import APIRouter, Depends, HTTPException -from sqlalchemy import select +from sqlalchemy import CursorResult, case, select, update from warden.api.routes.dependencies.auth import CurrentUserDep, SessionDep from warden.api.routes.dependencies.db import DBSessionDep @@ -98,31 +99,44 @@ async def cancel_job( db_session: DBSessionDep, identity: CurrentUserDep, ) -> JobResponse: - # Start transaction context async with db_session.begin(): - # Lock row/db during transaction - result = await db_session.execute( - select(Job) - .where(Job.user_id == identity.uid, Job.id == id) - .with_for_update(of=Job) + # Atomic claim: ownership and the cancelability guards are + # evaluated by the DB against the live row in one statement, so + # there's no read-then-write gap for a concurrent scheduler pickup + cancel_job_stmt = ( + update(Job) + .where( + Job.id == id, + Job.user_id == identity.uid, + Job.status.not_in(("CANCELED", "DONE", "ERROR")), + Job.canceled_at.is_(None), + ) + .values( + canceled_at=datetime.now(timezone.utc), + status=case((Job.scheduled_at.is_(None), "CANCELED"), else_=Job.status), + ) ) - job = result.scalar_one_or_none() + result = cast( + CursorResult, + await db_session.execute(cancel_job_stmt), + ) + + job = ( + await db_session.execute( + select(Job).where(Job.id == id, Job.user_id == identity.uid) + ) + ).scalar_one_or_none() + if job is None: raise HTTPException(404, detail="Job not found") - elif job.status in ("CANCELED", "DONE", "ERROR"): - raise HTTPException( - 409, detail=f"Job with status '{job.status}' can't be canceled" - ) - elif job.canceled_at is not None: + if result.rowcount == 0: + if job.status in ("CANCELED", "DONE", "ERROR"): + raise HTTPException( + 409, detail=f"Job with status '{job.status}' can't be canceled" + ) raise HTTPException( 409, detail="Job with status was already requested to be stopped" ) - 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 return JobResponse.from_model(job) diff --git a/warden/api/routes/sessions.py b/warden/api/routes/sessions.py index af1c2e3..b6fdca1 100644 --- a/warden/api/routes/sessions.py +++ b/warden/api/routes/sessions.py @@ -1,9 +1,10 @@ from datetime import datetime, timezone from logging import getLogger +from typing import cast from fastapi import APIRouter, HTTPException from pydantic import UUID4 -from sqlalchemy import select +from sqlalchemy import CursorResult, case, select, update from warden.api.routes.dependencies.auth import ( AdminUserDep, @@ -51,25 +52,32 @@ async def revoke_session( 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) + # Atomic claim: same pattern as `jobs.py::cancel_job`, the status + # and cancelability guards are re-evaluated against the live row by + # this single UPDATE, so there's no read-then-write gap for a + # concurrent scheduler pickup. + result = cast( + CursorResult, + await db_session.execute( + update(Job) + .where( + Job.session_id == session_record.id, + Job.status.not_in(("CANCELED", "DONE", "ERROR")), + Job.canceled_at.is_(None), + ) + .values( + canceled_at=datetime.now(timezone.utc), + status=case( + (Job.scheduled_at.is_(None), "CANCELED"), else_=Job.status + ), + ) + ), ) - jobs_to_cancel = result.scalars() - for job in jobs_to_cancel: + if result.rowcount > 0: logger.info( - "Cancelling job '%s' attached to session %s", job.id, session_record.id + "Canceled %d job(s) attached to revoked session '%s'", + result.rowcount, + 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 return SessionResponse.from_model(session_record) diff --git a/warden/scheduler/strategy.py b/warden/scheduler/strategy.py index bab363b..6da4aef 100644 --- a/warden/scheduler/strategy.py +++ b/warden/scheduler/strategy.py @@ -2,29 +2,58 @@ from abc import ABC, abstractmethod from datetime import datetime, timezone -from typing import Optional +from typing import Optional, cast -from sqlalchemy import case, select +from sqlalchemy import CursorResult, case, select, update from sqlalchemy.ext.asyncio import AsyncSession from warden.lib.config import SchedulerStrategy from warden.lib.models import Job +SCHEDULABLE_STATUS = ["PENDING", "RUNNING"] + class Scheduler(ABC): @staticmethod @abstractmethod - async def get_next_job(session: AsyncSession) -> Optional[Job]: - """Return next job to run""" + async def _get_next_job_id(session: AsyncSession) -> Optional[int]: + """Return ID of next job, where each implementatino define its strategy""" pass + async def get_next_job(self, session: AsyncSession) -> Optional[Job]: + """Tries to 'acquire' the job to schedule by setting `scheduler_at` in db. + Then returns job record to schedule on the qpu""" + candidate_id = await self._get_next_job_id(session) + if candidate_id is None: + return None + + # Atomic claim: the status check is re-evaluated against the live + # row by this single UPDATE, so a job canceled concurrently between + # the candidate lookup above and this write can't get claimed here. + result = cast( + CursorResult, + await session.execute( + update(Job) + .where(Job.id == candidate_id, Job.status.in_(SCHEDULABLE_STATUS)) + .values(scheduled_at=datetime.now(timezone.utc)) + ), + ) + await session.commit() + if result.rowcount == 0: + # where clause had 0 match, candidate job was canceled concurrently + return None + + return await session.get(Job, candidate_id) + class FifoScheduler(Scheduler): + """Simple FIFO Queue""" + @staticmethod - async def get_next_job(session: AsyncSession) -> Optional[Job]: - stmt = ( - select(Job) - .where(Job.status.in_(["PENDING", "RUNNING"])) + async def _get_next_job_id(session: AsyncSession) -> Optional[int]: + candidate_stmt = ( + select(Job.id) + .where(Job.status.in_(SCHEDULABLE_STATUS)) .order_by( # Rank jobs with an assigned backend before pending ones without case((Job.backend_id.is_(None), 1), else_=0), @@ -33,15 +62,8 @@ async def get_next_job(session: AsyncSession) -> Optional[Job]: Job.id, ) .limit(1) - .with_for_update(of=Job) ) - res = await session.execute(stmt) - job = res.scalar_one_or_none() - if job: - job.scheduled_at = datetime.now(timezone.utc) - await session.commit() - await session.refresh(job) - return job + return (await session.execute(candidate_stmt)).scalar_one_or_none() schedulers = {SchedulerStrategy.FIFO: FifoScheduler()}