From dc322288d6c5fc7c4ee368f8c03b700c6e637752 Mon Sep 17 00:00:00 2001 From: badtst Date: Thu, 3 Sep 2026 08:18:49 +0000 Subject: [PATCH 1/6] Move scheduling and cancelation db operation to atomic update requests --- warden/api/routes/jobs.py | 52 +++++++++++++++++++++++------------- warden/scheduler/strategy.py | 37 ++++++++++++++++--------- 2 files changed, 58 insertions(+), 31 deletions(-) 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/scheduler/strategy.py b/warden/scheduler/strategy.py index bab363b..1be8db7 100644 --- a/warden/scheduler/strategy.py +++ b/warden/scheduler/strategy.py @@ -2,9 +2,9 @@ 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 @@ -22,8 +22,8 @@ async def get_next_job(session: AsyncSession) -> Optional[Job]: class FifoScheduler(Scheduler): @staticmethod async def get_next_job(session: AsyncSession) -> Optional[Job]: - stmt = ( - select(Job) + candidate_stmt = ( + select(Job.id) .where(Job.status.in_(["PENDING", "RUNNING"])) .order_by( # Rank jobs with an assigned backend before pending ones without @@ -33,15 +33,28 @@ 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 + candidate_id = (await session.execute(candidate_stmt)).scalar_one_or_none() + 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_(["PENDING", "RUNNING"])) + .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) schedulers = {SchedulerStrategy.FIFO: FifoScheduler()} From 091bcc7321b892684f554fb9bddc409ab59d1cf7 Mon Sep 17 00:00:00 2001 From: badtst Date: Thu, 3 Sep 2026 13:28:35 +0000 Subject: [PATCH 2/6] small scheduler stratehy refactor --- warden/scheduler/strategy.py | 49 +++++++++++++++++++++--------------- 1 file changed, 29 insertions(+), 20 deletions(-) diff --git a/warden/scheduler/strategy.py b/warden/scheduler/strategy.py index 1be8db7..03fd8fa 100644 --- a/warden/scheduler/strategy.py +++ b/warden/scheduler/strategy.py @@ -10,31 +10,20 @@ 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 it's strategy""" pass - -class FifoScheduler(Scheduler): - @staticmethod - async def get_next_job(session: AsyncSession) -> Optional[Job]: - candidate_stmt = ( - select(Job.id) - .where(Job.status.in_(["PENDING", "RUNNING"])) - .order_by( - # Rank jobs with an assigned backend before pending ones without - case((Job.backend_id.is_(None), 1), else_=0), - Job.backend_id.asc(), - Job.created_at, - Job.id, - ) - .limit(1) - ) - candidate_id = (await session.execute(candidate_stmt)).scalar_one_or_none() + 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 @@ -45,7 +34,7 @@ async def get_next_job(session: AsyncSession) -> Optional[Job]: CursorResult, await session.execute( update(Job) - .where(Job.id == candidate_id, Job.status.in_(["PENDING", "RUNNING"])) + .where(Job.id == candidate_id, Job.status.in_(SCHEDULABLE_STATUS)) .values(scheduled_at=datetime.now(timezone.utc)) ), ) @@ -57,4 +46,24 @@ async def get_next_job(session: AsyncSession) -> Optional[Job]: return await session.get(Job, candidate_id) +class FifoScheduler(Scheduler): + """Simple FIFO Queue""" + + @staticmethod + 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), + Job.backend_id.asc(), + Job.created_at, + Job.id, + ) + .limit(1) + ) + return (await session.execute(candidate_stmt)).scalar_one_or_none() + + schedulers = {SchedulerStrategy.FIFO: FifoScheduler()} From f29294f4e0af21ce80306b1b64701b5d011d1cbe Mon Sep 17 00:00:00 2001 From: badtst Date: Thu, 3 Sep 2026 14:35:26 +0000 Subject: [PATCH 3/6] Working state --- tests/api/test_cancelation_race_condition.py | 207 +++++++++++++++++++ tests/api/test_jobs.py | 45 ++++ warden/api/routes/sessions.py | 44 ++-- 3 files changed, 278 insertions(+), 18 deletions(-) create mode 100644 tests/api/test_cancelation_race_condition.py diff --git a/tests/api/test_cancelation_race_condition.py b/tests/api/test_cancelation_race_condition.py new file mode 100644 index 0000000..20f8388 --- /dev/null +++ b/tests/api/test_cancelation_race_condition.py @@ -0,0 +1,207 @@ +import asyncio +from typing import cast + +import pytest +from sqlalchemy import 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 + + +@pytest.mark.asyncio +async def test_empiric_cancel_and_scheduler_pick_race_no_double_claim( + client, app, serialized_sequence: str +): + """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: whichever side wins, the other must lose cleanly + 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[SchedulerStrategy.FIFO] + + 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() + + +class _HookedSession: + """Session proxy running `hook` once, right after the first `execute`. + + Passed to `get_next_job`, it lands the hook in the candidate-pick -> + claim window: the one read-then-write gap the scheduler still has. + """ + + 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, *args, **kwargs): + result = await self._session.execute(*args, **kwargs) + if self._hook is not None: + hook, self._hook = self._hook, None + await hook() + return result + + +@pytest.mark.asyncio +async def test_cancel_racing_scheduler_pick_is_not_claimed( + client, app, serialized_sequence: str +): + """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[SchedulerStrategy.FIFO] + async with async_session() as session: + claimed = await scheduler.get_next_job( + cast(AsyncSession, _HookedSession(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 +async def test_empiric_session_revoke_and_scheduler_pick_race_no_double_claim( + client, app, serialized_sequence: str +): + """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[SchedulerStrategy.FIFO] + + 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() diff --git a/tests/api/test_jobs.py b/tests/api/test_jobs.py index 6c56d98..c77f17a 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,47 @@ 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 winner + + 1. Create a PENDING job 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 + + Each request runs in its own session/transaction, so the `canceled_at IS + NULL` guard has to be enforced by the DB, not by a prior read. + """ + 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/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) From ed9038c6f72f10d125db493b7dc00b47ef9f42ec Mon Sep 17 00:00:00 2001 From: badtst Date: Thu, 3 Sep 2026 15:21:02 +0000 Subject: [PATCH 4/6] Improve testing --- tests/api/test_cancelation_race_condition.py | 167 +++++++++++++------ tests/api/test_jobs.py | 7 +- 2 files changed, 118 insertions(+), 56 deletions(-) diff --git a/tests/api/test_cancelation_race_condition.py b/tests/api/test_cancelation_race_condition.py index 20f8388..7c59014 100644 --- a/tests/api/test_cancelation_race_condition.py +++ b/tests/api/test_cancelation_race_condition.py @@ -1,8 +1,12 @@ +""" +Testing possible race conditions between the api cancelling a job and the scheduler +""" + import asyncio from typing import cast import pytest -from sqlalchemy import select, update +from sqlalchemy import Update, select, update from sqlalchemy.ext.asyncio import AsyncSession from tests.api.conftest import mock_munge_auth @@ -11,9 +15,13 @@ from warden.lib.models.sessions import Session from warden.scheduler.strategy import schedulers +########################################################################## +####################### Repeated behavior testing ######################## +########################################################################## + @pytest.mark.asyncio -async def test_empiric_cancel_and_scheduler_pick_race_no_double_claim( +async def test_repeated_job_cancel_and_scheduler_pick_race( client, app, serialized_sequence: str ): """Assert racing a real cancel against a real scheduler pick never lets @@ -23,7 +31,7 @@ async def test_empiric_cancel_and_scheduler_pick_race_no_double_claim( 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: whichever side wins, the other must lose cleanly + `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 @@ -75,11 +83,74 @@ async def do_schedule(): await session.commit() -class _HookedSession: - """Session proxy running `hook` once, right after the first `execute`. +@pytest.mark.asyncio +async def test_repeated_session_revoke_and_scheduler_pick_race( + client, app, serialized_sequence: str +): + """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[SchedulerStrategy.FIFO] + + 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) - Passed to `get_next_job`, it lands the hook in the candidate-pick -> - claim window: the one read-then-write gap the scheduler still has. + 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`. + + Passed to `get_next_job`, this lands the hook in the candidate-pick -> + claim window: whatever reads a strategy does to pick a candidate, the + claim itself has to be a write, so triggering on the first `UPDATE` + (rather than on a fixed call count) keeps this test working regardless + of how many reads a strategy's candidate pick does. """ def __init__(self, session: AsyncSession, hook): @@ -89,12 +160,11 @@ def __init__(self, session: AsyncSession, hook): def __getattr__(self, name): return getattr(self._session, name) - async def execute(self, *args, **kwargs): - result = await self._session.execute(*args, **kwargs) - if self._hook is not None: + 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 result + return await self._session.execute(statement, *args, **kwargs) @pytest.mark.asyncio @@ -139,7 +209,7 @@ async def cancel_mid_pick(): scheduler = schedulers[SchedulerStrategy.FIFO] async with async_session() as session: claimed = await scheduler.get_next_job( - cast(AsyncSession, _HookedSession(session, cancel_mid_pick)) + cast(AsyncSession, _HookedUpdateDBSession(session, cancel_mid_pick)) ) assert claimed is None @@ -153,55 +223,50 @@ async def cancel_mid_pick(): @pytest.mark.asyncio -async def test_empiric_session_revoke_and_scheduler_pick_race_no_double_claim( +async def test_session_revoke_racing_scheduler_pick_is_not_claimed( client, app, serialized_sequence: str ): - """Assert racing a real session revoke against a real scheduler pick - never lets a job end up CANCELED and scheduled at once + """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 - scheduler = schedulers[SchedulerStrategy.FIFO] - 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 with async_session() as session: + session.add(job) + await session.commit() + await session.refresh(job) - async def do_schedule(): - async with async_session() as session: - return await scheduler.get_next_job(session) + job_id = job.id + session_id = session_record.id + async def revoke_mid_pick(): 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 + response = await client.delete(f"/sessions/{session_id}") + assert response.status_code == 200 - async with async_session() as session: - job = ( - await session.execute(select(Job).where(Job.id == job_id)) - ).scalar_one() + scheduler = schedulers[SchedulerStrategy.FIFO] + async with async_session() as session: + claimed = await scheduler.get_next_job( + cast(AsyncSession, _HookedUpdateDBSession(session, revoke_mid_pick)) + ) - assert not (job.status == "CANCELED" and job.scheduled_at is not None) + assert claimed is None - if claimed is not None: - assert job.scheduled_at is not None + async with async_session() as session: + job = (await session.execute(select(Job).where(Job.id == job_id))).scalar_one() - # 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() + 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 c77f17a..d735300 100644 --- a/tests/api/test_jobs.py +++ b/tests/api/test_jobs.py @@ -553,15 +553,12 @@ async def test_cancel_job_twice(client, app, serialized_sequence: str): async def test_concurrent_cancels_have_a_single_winner( client, app, serialized_sequence: str ): - """Assert concurrent cancels of the same job produce exactly one winner + """Assert concurrent cancels of the same job produce exactly one cancelation - 1. Create a PENDING job for a given user + 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 - - Each request runs in its own session/transaction, so the `canceled_at IS - NULL` guard has to be enforced by the DB, not by a prior read. """ user_id = 1000 job = Job( From 90579fcb6fa2d2eee48b5dbbc462770bedf0071a Mon Sep 17 00:00:00 2001 From: badtst Date: Thu, 3 Sep 2026 16:09:05 +0000 Subject: [PATCH 5/6] minro fix --- tests/api/test_cancelation_race_condition.py | 20 ++++++++++++-------- warden/scheduler/strategy.py | 2 +- 2 files changed, 13 insertions(+), 9 deletions(-) diff --git a/tests/api/test_cancelation_race_condition.py b/tests/api/test_cancelation_race_condition.py index 7c59014..992e7c0 100644 --- a/tests/api/test_cancelation_race_condition.py +++ b/tests/api/test_cancelation_race_condition.py @@ -21,8 +21,9 @@ @pytest.mark.asyncio +@pytest.mark.parametrize("strategy", list(SchedulerStrategy)) async def test_repeated_job_cancel_and_scheduler_pick_race( - client, app, serialized_sequence: str + 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 @@ -36,7 +37,7 @@ async def test_repeated_job_cancel_and_scheduler_pick_race( """ user_id = 1000 async_session = app.state.db_session_factory - scheduler = schedulers[SchedulerStrategy.FIFO] + scheduler = schedulers[strategy] for _ in range(40): job = Job( @@ -84,15 +85,16 @@ async def do_schedule(): @pytest.mark.asyncio +@pytest.mark.parametrize("strategy", list(SchedulerStrategy)) async def test_repeated_session_revoke_and_scheduler_pick_race( - client, app, serialized_sequence: str + 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[SchedulerStrategy.FIFO] + scheduler = schedulers[strategy] for _ in range(40): session_record = Session(user_id=str(user_id), slurm_job_id="1") @@ -168,8 +170,9 @@ async def execute(self, 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 + client, app, serialized_sequence: str, strategy ): """Assert a cancel committed mid-pick stops the scheduler claiming the job @@ -206,7 +209,7 @@ async def cancel_mid_pick(): assert response.status_code == 200 assert response.json()["status"] == "CANCELED" - scheduler = schedulers[SchedulerStrategy.FIFO] + scheduler = schedulers[strategy] async with async_session() as session: claimed = await scheduler.get_next_job( cast(AsyncSession, _HookedUpdateDBSession(session, cancel_mid_pick)) @@ -223,8 +226,9 @@ async def cancel_mid_pick(): @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 + client, app, serialized_sequence: str, strategy ): """Assert a session revoke committed mid-pick stops the scheduler claiming the job @@ -256,7 +260,7 @@ async def revoke_mid_pick(): response = await client.delete(f"/sessions/{session_id}") assert response.status_code == 200 - scheduler = schedulers[SchedulerStrategy.FIFO] + scheduler = schedulers[strategy] async with async_session() as session: claimed = await scheduler.get_next_job( cast(AsyncSession, _HookedUpdateDBSession(session, revoke_mid_pick)) diff --git a/warden/scheduler/strategy.py b/warden/scheduler/strategy.py index 03fd8fa..6da4aef 100644 --- a/warden/scheduler/strategy.py +++ b/warden/scheduler/strategy.py @@ -17,7 +17,7 @@ class Scheduler(ABC): @staticmethod @abstractmethod async def _get_next_job_id(session: AsyncSession) -> Optional[int]: - """Return ID of next job, where each implementatino define it's strategy""" + """Return ID of next job, where each implementatino define its strategy""" pass async def get_next_job(self, session: AsyncSession) -> Optional[Job]: From 945fc879841c831be13b7e61244e5de0e518f34f Mon Sep 17 00:00:00 2001 From: badtst Date: Tue, 8 Sep 2026 11:10:53 +0000 Subject: [PATCH 6/6] Update Hooked session comment --- tests/api/test_cancelation_race_condition.py | 9 ++------- 1 file changed, 2 insertions(+), 7 deletions(-) diff --git a/tests/api/test_cancelation_race_condition.py b/tests/api/test_cancelation_race_condition.py index 992e7c0..80afa6e 100644 --- a/tests/api/test_cancelation_race_condition.py +++ b/tests/api/test_cancelation_race_condition.py @@ -146,13 +146,8 @@ async def do_schedule(): class _HookedUpdateDBSession: - """Session proxy running `hook` once, right before the first `UPDATE`. - - Passed to `get_next_job`, this lands the hook in the candidate-pick -> - claim window: whatever reads a strategy does to pick a candidate, the - claim itself has to be a write, so triggering on the first `UPDATE` - (rather than on a fixed call count) keeps this test working regardless - of how many reads a strategy's candidate pick does. + """ + Session proxy running `hook` once, right before the first `UPDATE`. """ def __init__(self, session: AsyncSession, hook):