From 3deb865630570617565af6294820b46ba9473bb5 Mon Sep 17 00:00:00 2001 From: Aleksander Wennersteen Date: Tue, 14 Jul 2026 13:31:40 +0200 Subject: [PATCH 1/2] Add Warden side features for QPU slots --- README.md | 30 ++++- tests/api/conftest.py | 4 + tests/api/test_accessible.py | 30 +++++ tests/api/test_sessions.py | 125 ++++++++++++++++++ tests/scheduler/test_strategy.py | 43 ++++++ ...02__add_session_qpu_slots__8ad0f6a4b2c1.py | 58 ++++++++ warden/api/routes/accessible.py | 33 ++++- warden/api/routes/dependencies/qpu_client.py | 6 + warden/api/routes/sessions.py | 99 +++++++++++--- warden/api/schemas/accessible.py | 3 + warden/api/schemas/sessions.py | 5 +- warden/lib/config/config.py | 1 + warden/lib/config/config.sample.yaml | 5 + warden/lib/models/__init__.py | 4 +- warden/lib/models/sessions.py | 21 ++- warden/scheduler/strategy.py | 14 +- 16 files changed, 450 insertions(+), 31 deletions(-) create mode 100644 warden/api/alembic/versions/2026-07-14T12-10-02__add_session_qpu_slots__8ad0f6a4b2c1.py diff --git a/README.md b/README.md index 9567050..04b3dee 100644 --- a/README.md +++ b/README.md @@ -98,7 +98,6 @@ Warden's access to the PASQAL QPU can be configured through the YAML config or e - `false`: disable verification entirely. **Insecure**, dev/local testing only. - a path (e.g. `/etc/warden/backend-ca.pem`) — verify against a specific CA bundle / certificate file. - The API server can also be configured to only accept new jobs from configured user IDs: | Path | Description | Default | Required | Example Value | @@ -167,3 +166,32 @@ Configure Warden to accept jobs again by configuring: ```bash make set-accessible IS_ACCESSIBLE=true MESSAGE="Maintenance done" ``` + +### External Resource-Manager Polling + +External schedulers can poll `GET /accessible` to decide whether to offer a QPU +resource for early scheduling. The response is: + +```json +{"is_accessible": true, "message": "QPU accessible"} +``` + +`is_accessible=false` means the resource should be treated as unavailable by +that external scheduler. Polling this endpoint is only a readiness hint; Warden +still performs its normal session and job handling. + +If `qpu.qpu_slots_total` is set, sessions may include `qpu_slots` and Warden +rejects new sessions when active sessions would exceed that total. In that +case, `GET /accessible` also returns `qpu_slots_total`, `qpu_slots_used`, and +`qpu_slots_available` for external polling. + +Capacity admission is serialized in the database, so concurrent session +requests cannot oversubscribe the configured total. Warden derives session +idempotency from `(user_id, slurm_job_id)`: repeating an active request for +the same scheduler job returns the existing session, while changing its slot +count returns `409`. + +`qpu_slots` is also the weight for Warden's job-level scheduler. A five-slot +session receives approximately five scheduling turns for every turn received +by a one-slot session, while jobs remain FIFO within a session. Running QPU jobs +are not preempted. diff --git a/tests/api/conftest.py b/tests/api/conftest.py index 582eb50..9f642f9 100644 --- a/tests/api/conftest.py +++ b/tests/api/conftest.py @@ -13,6 +13,7 @@ from warden.api.routes.dependencies.qpu_client import get_qpu_client from warden.lib.config.config import APIConfig, Config, DatabaseConfig, QPUConfig from warden.lib.db.database import Base +from warden.lib.models import QPUCapacityLock from warden.lib.qpu_client.client import AsyncQPUClient @@ -24,6 +25,9 @@ async def app(db_backend_config: DatabaseConfig) -> AsyncGenerator[FastAPI, None # create tables in the test database async with app.state.db_engine.begin() as conn: await conn.run_sync(Base.metadata.create_all) + async with app.state.db_session_factory() as session: + session.add(QPUCapacityLock(id=1)) + await session.commit() yield app async with app.state.db_engine.begin() as conn: await conn.run_sync(Base.metadata.drop_all) diff --git a/tests/api/test_accessible.py b/tests/api/test_accessible.py index da870aa..2f81bcc 100644 --- a/tests/api/test_accessible.py +++ b/tests/api/test_accessible.py @@ -2,6 +2,7 @@ from httpx import AsyncClient from tests.api.conftest import mock_munge_auth +from warden.lib.models import Session @pytest.mark.asyncio @@ -70,3 +71,32 @@ async def test_accessible_auth_update(client: AsyncClient, app): with mock_munge_auth(app, uid=0): response = await client.post("/accessible", json=payload) assert response.status_code == 200 + + +@pytest.mark.asyncio +async def test_accessible_get_contract_for_external_polling(client: AsyncClient): + """Verify GET /accessible is unauthenticated and schema-stable.""" + + response = await client.get("/accessible") + assert response.status_code == 200 + assert {"is_accessible", "message"}.issubset(response.json()) + assert isinstance(response.json()["is_accessible"], bool) + assert isinstance(response.json()["message"], str) + + +@pytest.mark.asyncio +async def test_accessible_reports_configured_qpu_slots(client: AsyncClient, app): + """Verify GET /accessible includes configured QPU slot capacity.""" + + app.state.qpu_config.qpu_slots_total = 10 + async_session = app.state.db_session_factory + async with async_session() as session: + session.add(Session(user_id="1000", slurm_job_id="1", qpu_slots=4)) + await session.commit() + + response = await client.get("/accessible") + + assert response.status_code == 200 + assert response.json()["qpu_slots_total"] == 10 + assert response.json()["qpu_slots_used"] == 4 + assert response.json()["qpu_slots_available"] == 6 diff --git a/tests/api/test_sessions.py b/tests/api/test_sessions.py index d31a8c4..bc3cea6 100644 --- a/tests/api/test_sessions.py +++ b/tests/api/test_sessions.py @@ -1,3 +1,4 @@ +import asyncio from datetime import datetime import pytest @@ -31,6 +32,130 @@ async def test_create_session_success(client, app): assert response.status_code == 200 data = response.json() assert data["user_id"] == payload["user_id"] + assert data["qpu_slots"] == 1 + + +@pytest.mark.asyncio +async def test_create_session_with_qpu_slots(client, app): + """Creating a session stores requested QPU slots.""" + + payload = {"user_id": "1000", "slurm_job_id": "1", "qpu_slots": 5} + with mock_munge_auth(app, uid=0): + response = await client.post("/sessions", json=payload) + assert response.status_code == 200 + assert response.json()["qpu_slots"] == 5 + + +@pytest.mark.asyncio +async def test_create_session_rejects_invalid_qpu_slots(client, app): + """Creating a session rejects non-positive QPU slots.""" + + payload = {"user_id": "1000", "slurm_job_id": "1", "qpu_slots": 0} + with mock_munge_auth(app, uid=0): + response = await client.post("/sessions", json=payload) + assert response.status_code == 422 + + +@pytest.mark.asyncio +async def test_create_session_enforces_configured_qpu_slots(client, app): + """Creating a session fails when active sessions exhaust QPU slots.""" + + app.state.qpu_config.qpu_slots_total = 10 + payload = {"user_id": "1000", "slurm_job_id": "1", "qpu_slots": 5} + with mock_munge_auth(app, uid=0): + assert (await client.post("/sessions", json=payload)).status_code == 200 + assert ( + await client.post( + "/sessions", + json={"user_id": "1000", "slurm_job_id": "2", "qpu_slots": 5}, + ) + ).status_code == 200 + response = await client.post( + "/sessions", json={"user_id": "1000", "slurm_job_id": "3", "qpu_slots": 1} + ) + assert response.status_code == 409 + + +@pytest.mark.asyncio +async def test_create_session_enforces_qpu_slots_concurrently(client, app): + """Concurrent session creation cannot exceed configured QPU slots.""" + + app.state.qpu_config.qpu_slots_total = 10 + with mock_munge_auth(app, uid=0): + responses = await asyncio.gather( + client.post( + "/sessions", + json={"user_id": "1000", "slurm_job_id": "1", "qpu_slots": 6}, + ), + client.post( + "/sessions", + json={"user_id": "1000", "slurm_job_id": "2", "qpu_slots": 6}, + ), + ) + assert sorted(response.status_code for response in responses) == [200, 409] + + +@pytest.mark.asyncio +async def test_create_session_is_idempotent(client, app): + """An active scheduler job returns the original session.""" + + payload = { + "user_id": "1000", + "slurm_job_id": "1", + "qpu_slots": 5, + } + with mock_munge_auth(app, uid=0): + first = await client.post("/sessions", json=payload) + second = await client.post("/sessions", json=payload) + assert first.status_code == second.status_code == 200 + assert first.json()["id"] == second.json()["id"] + + +@pytest.mark.asyncio +async def test_create_session_rejects_job_parameter_change(client, app): + """An active scheduler job cannot be reused with different parameters.""" + + payload = { + "user_id": "1000", + "slurm_job_id": "1", + "qpu_slots": 5, + } + with mock_munge_auth(app, uid=0): + assert (await client.post("/sessions", json=payload)).status_code == 200 + payload["qpu_slots"] = 4 + response = await client.post("/sessions", json=payload) + assert response.status_code == 409 + + +@pytest.mark.asyncio +async def test_revoke_session_frees_qpu_slots(client, app): + """Revoking a session frees its QPU slots for a later session.""" + + app.state.qpu_config.qpu_slots_total = 5 + payload = {"user_id": "1000", "slurm_job_id": "1", "qpu_slots": 5} + with mock_munge_auth(app, uid=0): + response = await client.post("/sessions", json=payload) + assert response.status_code == 200 + session_id = response.json()["id"] + assert (await client.delete(f"/sessions/{session_id}")).status_code == 200 + response = await client.post("/sessions", json=payload) + assert response.status_code == 200 + + +@pytest.mark.asyncio +async def test_revoke_session_is_idempotent(client, app): + """Repeated session revocation preserves the first revocation.""" + + payload = {"user_id": "1000", "slurm_job_id": "1"} + with mock_munge_auth(app, uid=0): + created = await client.post("/sessions", json=payload) + session_id = created.json()["id"] + first = await client.delete(f"/sessions/{session_id}") + second = await client.delete(f"/sessions/{session_id}") + assert first.status_code == second.status_code == 200 + assert first.json()["revoked_at"].rstrip("Z") == second.json()["revoked_at"].rstrip( + "Z" + ) @pytest.mark.asyncio diff --git a/tests/scheduler/test_strategy.py b/tests/scheduler/test_strategy.py index ca7b1aa..56beef0 100644 --- a/tests/scheduler/test_strategy.py +++ b/tests/scheduler/test_strategy.py @@ -198,3 +198,46 @@ async def test_fifo_job_running(db_session_maker): assert schedule[2].id == 3 assert schedule[3].id == 2 assert schedule[4] is None + + +@pytest.mark.asyncio +async def test_fifo_weights_sessions_by_qpu_slots(db_session_maker): + """QPU slots weight job-level scheduling turns across sessions.""" + + scheduler = schedulers[SchedulerStrategy.FIFO] + now = datetime.now() + large = Session(slurm_job_id="large", user_id="1000", qpu_slots=5) + small = Session(slurm_job_id="small", user_id="1001", qpu_slots=1) + jobs = [] + for index in range(12): + jobs.extend( + [ + Job( + session=large, + shots=100, + sequence="{}", + status="PENDING", + created_at=now + timedelta(microseconds=index * 2), + ), + Job( + session=small, + shots=100, + sequence="{}", + status="PENDING", + created_at=now + timedelta(microseconds=index * 2 + 1), + ), + ] + ) + + async with db_session_maker() as session: + session.add_all(jobs) + await session.commit() + scheduled_sessions = [] + for _ in range(6): + job = await scheduler.get_next_job(session) + scheduled_sessions.append(job.session.slurm_job_id) + job.status = "DONE" + await session.commit() + + assert scheduled_sessions.count("large") == 5 + assert scheduled_sessions.count("small") == 1 diff --git a/warden/api/alembic/versions/2026-07-14T12-10-02__add_session_qpu_slots__8ad0f6a4b2c1.py b/warden/api/alembic/versions/2026-07-14T12-10-02__add_session_qpu_slots__8ad0f6a4b2c1.py new file mode 100644 index 0000000..93b086f --- /dev/null +++ b/warden/api/alembic/versions/2026-07-14T12-10-02__add_session_qpu_slots__8ad0f6a4b2c1.py @@ -0,0 +1,58 @@ +"""add QPU capacity and scheduler fields + +Revision ID: 8ad0f6a4b2c1 +Revises: 6c4fad0bfc30 +Create Date: 2026-07-14 12:10:00.000000 + +""" + +from typing import Sequence, Union + +import sqlalchemy as sa +from alembic import op + +# revision identifiers, used by Alembic. +revision: str = "8ad0f6a4b2c1" +down_revision: Union[str, Sequence[str], None] = "6c4fad0bfc30" +branch_labels: Union[str, Sequence[str], None] = None +depends_on: Union[str, Sequence[str], None] = None + + +def upgrade() -> None: + """Upgrade schema.""" + bind = op.get_bind() + op.add_column( + "sessions", + sa.Column("qpu_slots", sa.Integer(), server_default="1", nullable=False), + ) + if bind.dialect.name != "sqlite": + op.alter_column("sessions", "qpu_slots", server_default=None) + + table = op.create_table( + "qpu_capacity_lock", + sa.Column("id", sa.Integer(), nullable=False), + sa.Column("revision", sa.Integer(), nullable=False), + sa.PrimaryKeyConstraint("id"), + ) + op.bulk_insert(table, [{"id": 1, "revision": 0}]) + + op.add_column( + "sessions", + sa.Column( + "scheduler_vruntime", + sa.Float(), + server_default="0", + nullable=False, + ), + ) + if bind.dialect.name != "sqlite": + op.alter_column("sessions", "scheduler_vruntime", server_default=None) + + +def downgrade() -> None: + """Downgrade schema.""" + with op.batch_alter_table("sessions") as batch_op: + batch_op.drop_column("scheduler_vruntime") + op.drop_table("qpu_capacity_lock") + with op.batch_alter_table("sessions") as batch_op: + batch_op.drop_column("qpu_slots") diff --git a/warden/api/routes/accessible.py b/warden/api/routes/accessible.py index a188b93..858df4a 100644 --- a/warden/api/routes/accessible.py +++ b/warden/api/routes/accessible.py @@ -1,26 +1,53 @@ from logging import getLogger -from fastapi import APIRouter +from fastapi import APIRouter, Depends +from sqlalchemy import func, select from warden.api.routes.dependencies.auth import AdminUserDep from warden.api.routes.dependencies.db import DBSessionDep +from warden.api.routes.dependencies.qpu_client import get_qpu_config from warden.api.schemas.accessible import AccessibleResponse, UpdateAccessibleRequest +from warden.lib.config.config import QPUConfig from warden.lib.models.accessible import ( AccessibilitySettings, get_latest_accessibility_settings, ) +from warden.lib.models.sessions import Session, active_session_filter logger = getLogger(__name__) router = APIRouter(prefix="/accessible") @router.get("") -async def is_accessible(db_session: DBSessionDep) -> AccessibleResponse: +async def is_accessible( + db_session: DBSessionDep, + qpu_config: QPUConfig = Depends(get_qpu_config), +) -> AccessibleResponse: """Warden endpoint for qrmi 'is_accessible' interface""" settings = await get_latest_accessibility_settings(db_session) + slot_state = await qpu_slot_state(db_session, qpu_config) return AccessibleResponse( - is_accessible=settings.is_accessible, message=settings.message + is_accessible=settings.is_accessible, + message=settings.message, + **slot_state, + ) + + +async def qpu_slot_state(db_session: DBSessionDep, qpu_config: QPUConfig) -> dict[str, int]: + if qpu_config.qpu_slots_total is None: + return {} + result = await db_session.execute( + select(func.coalesce(func.sum(Session.qpu_slots), 0)).where( + active_session_filter() + ) ) + used = int(result.scalar_one()) + total = qpu_config.qpu_slots_total + return { + "qpu_slots_total": total, + "qpu_slots_used": used, + "qpu_slots_available": max(total - used, 0), + } @router.post("") diff --git a/warden/api/routes/dependencies/qpu_client.py b/warden/api/routes/dependencies/qpu_client.py index dd6c490..84fafd2 100644 --- a/warden/api/routes/dependencies/qpu_client.py +++ b/warden/api/routes/dependencies/qpu_client.py @@ -6,9 +6,15 @@ def init_qpu_client(app: FastAPI, qpu_config: QPUConfig): """Initialize the QPU client.""" + app.state.qpu_config = qpu_config app.state.qpu_client = AsyncQPUClient(qpu_config) def get_qpu_client(request: Request) -> AsyncQPUClient: """Get the initialized http client to interact with the QPU.""" return request.app.state.qpu_client + + +def get_qpu_config(request: Request) -> QPUConfig: + """Get the initialized QPU configuration.""" + return request.app.state.qpu_config diff --git a/warden/api/routes/sessions.py b/warden/api/routes/sessions.py index af1c2e3..fe4e7f5 100644 --- a/warden/api/routes/sessions.py +++ b/warden/api/routes/sessions.py @@ -1,9 +1,9 @@ from datetime import datetime, timezone from logging import getLogger -from fastapi import APIRouter, HTTPException +from fastapi import APIRouter, Depends, HTTPException from pydantic import UUID4 -from sqlalchemy import select +from sqlalchemy import func, select, update from warden.api.routes.dependencies.auth import ( AdminUserDep, @@ -11,8 +11,11 @@ ensure_user_is_authorized, ) from warden.api.routes.dependencies.db import DBSessionDep +from warden.api.routes.dependencies.qpu_client import get_qpu_config from warden.api.schemas.sessions import CreateSession, SessionResponse -from warden.lib.models import Job, Session +from warden.lib.config.config import QPUConfig +from warden.lib.models import Job, QPUCapacityLock, Session +from warden.lib.models.sessions import active_session_filter logger = getLogger(__name__) router = APIRouter(prefix="/sessions") @@ -24,31 +27,93 @@ async def create_session( db_session: DBSessionDep, auth_config: AuthConfigDep, _admin: AdminUserDep, + qpu_config: QPUConfig = Depends(get_qpu_config), ) -> SessionResponse: ensure_user_is_authorized(auth_config, str(payload.user_id)) - new_session = Session( - user_id=str(payload.user_id), - slurm_job_id=payload.slurm_job_id, - ) - db_session.add(new_session) - await db_session.flush() - await db_session.commit() + async with db_session.begin(): + await lock_qpu_capacity(db_session) + existing = await active_session_for_job( + db_session, str(payload.user_id), payload.slurm_job_id + ) + if existing is not None: + if existing.qpu_slots != payload.qpu_slots: + raise HTTPException( + status_code=409, + detail="An active session already exists for this scheduler job with different parameters.", + ) + return SessionResponse.from_model(existing) + if qpu_config.qpu_slots_total is not None: + used = await active_qpu_slots(db_session) + if used + payload.qpu_slots > qpu_config.qpu_slots_total: + raise HTTPException( + status_code=409, + detail="Not enough QPU slots available.", + ) + new_session = Session( + user_id=str(payload.user_id), + slurm_job_id=payload.slurm_job_id, + qpu_slots=payload.qpu_slots, + ) + db_session.add(new_session) + await db_session.flush() return SessionResponse.from_model(new_session) +async def active_session_for_job( + db_session: DBSessionDep, user_id: str, slurm_job_id: str +) -> Session | None: + result = await db_session.execute( + select(Session).where( + Session.user_id == user_id, + Session.slurm_job_id == slurm_job_id, + active_session_filter(), + ) + ) + return result.scalar_one_or_none() + + +async def lock_qpu_capacity(db_session: DBSessionDep) -> None: + result = await db_session.execute( + update(QPUCapacityLock) + .where(QPUCapacityLock.id == 1) + .values(revision=QPUCapacityLock.revision + 1) + ) + if result.rowcount != 1: + raise RuntimeError( + "QPU capacity lock is missing; run the latest Warden database migration." + ) + + +async def active_qpu_slots(db_session: DBSessionDep) -> int: + result = await db_session.execute( + select(func.coalesce(func.sum(Session.qpu_slots), 0)).where( + active_session_filter() + ) + ) + return int(result.scalar_one()) + + @router.delete("/{id}") async def revoke_session( id: UUID4, db_session: DBSessionDep, _admin: AdminUserDep, ) -> SessionResponse: - result = await db_session.execute(select(Session).where(Session.id == id)) - 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() + already_revoked = False + async with db_session.begin(): + await lock_qpu_capacity(db_session) + result = await db_session.execute( + select(Session).where(Session.id == id).with_for_update(of=Session) + ) + session_record = result.scalar_one_or_none() + if session_record is None: + raise HTTPException(status_code=404, detail="Session not found.") + already_revoked = session_record.revoked_at is not None + if not already_revoked: + session_record.revoked_at = datetime.now(timezone.utc) + + if already_revoked: + return SessionResponse.from_model(session_record) async with db_session.begin(): result = await db_session.execute( diff --git a/warden/api/schemas/accessible.py b/warden/api/schemas/accessible.py index 0a7c932..2a1fcc7 100644 --- a/warden/api/schemas/accessible.py +++ b/warden/api/schemas/accessible.py @@ -4,6 +4,9 @@ class AccessibleResponse(BaseModel): is_accessible: bool message: str + qpu_slots_total: int | None = None + qpu_slots_used: int | None = None + qpu_slots_available: int | None = None class UpdateAccessibleRequest(BaseModel): diff --git a/warden/api/schemas/sessions.py b/warden/api/schemas/sessions.py index f1a23f2..e0c2e87 100644 --- a/warden/api/schemas/sessions.py +++ b/warden/api/schemas/sessions.py @@ -1,7 +1,7 @@ from datetime import datetime from uuid import UUID -from pydantic import BaseModel +from pydantic import BaseModel, Field from warden.lib.models.sessions import Session @@ -9,6 +9,7 @@ class CreateSession(BaseModel): user_id: str slurm_job_id: str + qpu_slots: int = Field(default=1, ge=1) class SessionResponse(BaseModel): @@ -16,6 +17,7 @@ class SessionResponse(BaseModel): user_id: str created_at: datetime revoked_at: datetime | None + qpu_slots: int @classmethod def from_model(cls, session: Session) -> "SessionResponse": @@ -24,4 +26,5 @@ def from_model(cls, session: Session) -> "SessionResponse": user_id=session.user_id, created_at=session.created_at, revoked_at=session.revoked_at, + qpu_slots=session.qpu_slots, ) diff --git a/warden/lib/config/config.py b/warden/lib/config/config.py index a9fc30d..8eb57e0 100644 --- a/warden/lib/config/config.py +++ b/warden/lib/config/config.py @@ -72,6 +72,7 @@ class SchedulerConfig(BaseSettings): class QPUConfig(BaseSettings): uri: str = "http://localhost:8000" + qpu_slots_total: int | None = Field(default=None, gt=0) retry_max: int = 10 retry_sleep_s: float = 1 diff --git a/warden/lib/config/config.sample.yaml b/warden/lib/config/config.sample.yaml index 0b35b20..769084e 100644 --- a/warden/lib/config/config.sample.yaml +++ b/warden/lib/config/config.sample.yaml @@ -74,6 +74,11 @@ qpu: # Local Pasqal QPU API configuration uri: http://localhost:8000 + # Optional Warden-enforced QPU slot capacity. + # If unset, session creation stores requested qpu_slots but does not enforce + # an aggregate limit. Set this for local QPU reservation admission. + # qpu_slots_total: 10 + # TLS verification policy for requests to the QPU backend. # Only relevant when 'uri' uses https. Accepts: # system -> verify against the OS trust store, e.g. anchors added to diff --git a/warden/lib/models/__init__.py b/warden/lib/models/__init__.py index 37e9b49..2a1e4e4 100644 --- a/warden/lib/models/__init__.py +++ b/warden/lib/models/__init__.py @@ -1,6 +1,6 @@ from warden.lib.db.database import Base from warden.lib.models.accessible import AccessibilitySettings from warden.lib.models.jobs import Job -from warden.lib.models.sessions import Session +from warden.lib.models.sessions import QPUCapacityLock, Session -__all__ = ["Base", "Job", "Session", "AccessibilitySettings"] +__all__ = ["Base", "Job", "Session", "QPUCapacityLock", "AccessibilitySettings"] diff --git a/warden/lib/models/sessions.py b/warden/lib/models/sessions.py index aca6bd9..6edc00a 100644 --- a/warden/lib/models/sessions.py +++ b/warden/lib/models/sessions.py @@ -5,15 +5,20 @@ from sqlalchemy import ( UUID as UUIDType, ) -from sqlalchemy import ( - DateTime, - String, -) +from sqlalchemy import DateTime, Float, Integer, String from sqlalchemy.orm import Mapped, mapped_column +from sqlalchemy.sql.elements import ColumnElement from warden.lib.db.database import Base +class QPUCapacityLock(Base): + __tablename__ = "qpu_capacity_lock" + + id: Mapped[int] = mapped_column(Integer, primary_key=True) + revision: Mapped[int] = mapped_column(Integer, nullable=False, default=0) + + class Session(Base): __tablename__ = "sessions" @@ -33,3 +38,11 @@ class Session(Base): slurm_job_id: Mapped[str] = mapped_column( String(255), doc="ID of the slurm job which created this session." ) + qpu_slots: Mapped[int] = mapped_column(Integer, nullable=False, default=1) + scheduler_vruntime: Mapped[float] = mapped_column( + Float, nullable=False, default=0.0 + ) + + +def active_session_filter() -> ColumnElement[bool]: + return Session.revoked_at.is_(None) diff --git a/warden/scheduler/strategy.py b/warden/scheduler/strategy.py index bab363b..4d0e4e9 100644 --- a/warden/scheduler/strategy.py +++ b/warden/scheduler/strategy.py @@ -8,7 +8,8 @@ from sqlalchemy.ext.asyncio import AsyncSession from warden.lib.config import SchedulerStrategy -from warden.lib.models import Job +from warden.lib.models import Job, Session +from warden.lib.models.sessions import active_session_filter class Scheduler(ABC): @@ -24,20 +25,27 @@ class FifoScheduler(Scheduler): async def get_next_job(session: AsyncSession) -> Optional[Job]: stmt = ( select(Job) - .where(Job.status.in_(["PENDING", "RUNNING"])) + .join(Session) + .where( + Job.status.in_(["PENDING", "RUNNING"]), + active_session_filter(), + ) .order_by( # Rank jobs with an assigned backend before pending ones without case((Job.backend_id.is_(None), 1), else_=0), + Session.scheduler_vruntime, Job.backend_id.asc(), Job.created_at, Job.id, ) .limit(1) - .with_for_update(of=Job) + .with_for_update(of=[Job, Session]) ) res = await session.execute(stmt) job = res.scalar_one_or_none() if job: + if job.backend_id is None: + job.session.scheduler_vruntime += 1 / job.session.qpu_slots job.scheduled_at = datetime.now(timezone.utc) await session.commit() await session.refresh(job) From ec23bfdfe2ece92dbad16fb225374f9d7b924613 Mon Sep 17 00:00:00 2001 From: Aleksander Wennersteen Date: Tue, 14 Jul 2026 13:43:12 +0200 Subject: [PATCH 2/2] try to fix tests --- tests/api/conftest.py | 2 +- tests/api/test_sessions.py | 10 +++++++--- tests/scheduler/test_strategy.py | 1 + warden/api/routes/accessible.py | 4 +++- warden/api/routes/sessions.py | 4 +++- 5 files changed, 15 insertions(+), 6 deletions(-) diff --git a/tests/api/conftest.py b/tests/api/conftest.py index 9f642f9..f8dc31d 100644 --- a/tests/api/conftest.py +++ b/tests/api/conftest.py @@ -26,7 +26,7 @@ async def app(db_backend_config: DatabaseConfig) -> AsyncGenerator[FastAPI, None async with app.state.db_engine.begin() as conn: await conn.run_sync(Base.metadata.create_all) async with app.state.db_session_factory() as session: - session.add(QPUCapacityLock(id=1)) + await session.merge(QPUCapacityLock(id=1)) await session.commit() yield app async with app.state.db_engine.begin() as conn: diff --git a/tests/api/test_sessions.py b/tests/api/test_sessions.py index bc3cea6..7caf850 100644 --- a/tests/api/test_sessions.py +++ b/tests/api/test_sessions.py @@ -153,9 +153,13 @@ async def test_revoke_session_is_idempotent(client, app): first = await client.delete(f"/sessions/{session_id}") second = await client.delete(f"/sessions/{session_id}") assert first.status_code == second.status_code == 200 - assert first.json()["revoked_at"].rstrip("Z") == second.json()["revoked_at"].rstrip( - "Z" - ) + first_revoked_at = datetime.fromisoformat( + first.json()["revoked_at"].rstrip("Z") + ).replace(microsecond=0) + second_revoked_at = datetime.fromisoformat( + second.json()["revoked_at"].rstrip("Z") + ).replace(microsecond=0) + assert first_revoked_at == second_revoked_at @pytest.mark.asyncio diff --git a/tests/scheduler/test_strategy.py b/tests/scheduler/test_strategy.py index 56beef0..d11c611 100644 --- a/tests/scheduler/test_strategy.py +++ b/tests/scheduler/test_strategy.py @@ -235,6 +235,7 @@ async def test_fifo_weights_sessions_by_qpu_slots(db_session_maker): scheduled_sessions = [] for _ in range(6): job = await scheduler.get_next_job(session) + assert job is not None scheduled_sessions.append(job.session.slurm_job_id) job.status = "DONE" await session.commit() diff --git a/warden/api/routes/accessible.py b/warden/api/routes/accessible.py index 858df4a..5ac93a2 100644 --- a/warden/api/routes/accessible.py +++ b/warden/api/routes/accessible.py @@ -33,7 +33,9 @@ async def is_accessible( ) -async def qpu_slot_state(db_session: DBSessionDep, qpu_config: QPUConfig) -> dict[str, int]: +async def qpu_slot_state( + db_session: DBSessionDep, qpu_config: QPUConfig +) -> dict[str, int]: if qpu_config.qpu_slots_total is None: return {} result = await db_session.execute( diff --git a/warden/api/routes/sessions.py b/warden/api/routes/sessions.py index fe4e7f5..357740a 100644 --- a/warden/api/routes/sessions.py +++ b/warden/api/routes/sessions.py @@ -1,9 +1,11 @@ from datetime import datetime, timezone from logging import getLogger +from typing import Any, cast from fastapi import APIRouter, Depends, HTTPException from pydantic import UUID4 from sqlalchemy import func, select, update +from sqlalchemy.engine import CursorResult from warden.api.routes.dependencies.auth import ( AdminUserDep, @@ -78,7 +80,7 @@ async def lock_qpu_capacity(db_session: DBSessionDep) -> None: .where(QPUCapacityLock.id == 1) .values(revision=QPUCapacityLock.revision + 1) ) - if result.rowcount != 1: + if cast(CursorResult[Any], result).rowcount != 1: raise RuntimeError( "QPU capacity lock is missing; run the latest Warden database migration." )