Add Warden side features for QPU slots

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.

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:
+        await session.merge(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)

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

tests/api/test_sessions.py

@@ -1,3 +1,4 @@
+import asyncio
 from datetime import datetime
 
 import pytest
@@ -31,6 +32,134 @@ 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
+    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

tests/scheduler/test_strategy.py

@@ -198,3 +198,47 @@ 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)
+            assert job is not None
+            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

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")