Add Warden side features for QPU slots - #73
Conversation
SBOM vulnerability checkTotal vulnerabilities observed: 13
|
| Vulnerability | Severity | Component(s) | Reference |
|---|---|---|---|
CVE-2025-65896 |
CRITICAL | asyncmy@0.2.11 |
link |
warden-sbom-pg.cdx.json
- Severity breakdown: HIGH: 2, MEDIUM: 2
warden-sbom.cdx.json
- Severity breakdown: HIGH: 2, MEDIUM: 2
|
Nit: I'd suggest adding lifetimes to sessions before merging if we base the available slots on non-revoked sessions: #74 |
Good point @badtst. Do you have any idea for how to best handle this? One concern of mine will be that if a user job is extended by the cluster admin, we cannot e.g. simply use the initial requested walltime of the job as the lifetime. |
|
We could set an absurdly long lifetime as a catch all fallback, then have a policy to close a session if a job has not been submitted in a long time (e.g. 1 hour ?). Also maybe add an admin endpoint to close a specific session if the spank plugin session release failed and we want to free QPU slots quickly ? |
| 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) | ||
|
|
There was a problem hiding this comment.
I don't really understand the purpose of this QPUCapacityLock table
There was a problem hiding this comment.
I didn't like it. but since we have a DB driven architecture I saw no option.
The purpose is to ensure that the the QPU capacity is only changed sequentially. so that we don't e.g. accept a job, but then accept another before the slots is reduced.
There was a problem hiding this comment.
I see, maybe for this purpose we could use a asyncio.lock in the session creation route (https://docs.python.org/3/library/asyncio-sync.html#asyncio.Lock) that way sessions creation are mutex (if we stay single-threaded for the API) ?
| class FifoScheduler(Scheduler): | ||
| @staticmethod | ||
| 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) |
There was a problem hiding this comment.
Shouldn't this be the opportunity to define a new strategy here ? Leave this as a true FIFO and define a new WeightedFIFO scheduling strategy ?
| status_code=409, | ||
| detail="An active session already exists for this scheduler job with different parameters.", | ||
| ) | ||
| return SessionResponse.from_model(existing) |
There was a problem hiding this comment.
Only an admin user is allowed to create a new session, in the spank plugin case I don't see how there would be a request to update the session parameters (qpu_slots) or have duplicate requests (user_id and slurm_job_id) for a new session ? Would there be other workload implementations that would allow this (Gridware ? )
| return { | ||
| "qpu_slots_total": total, | ||
| "qpu_slots_used": used, | ||
| "qpu_slots_available": max(total - used, 0), | ||
| } |
There was a problem hiding this comment.
If this is information dedicated to an external daemon (i.e. for syncing dynamic licenses on slurm) maybe this could have it's dedicated route/endpoint ?
| `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. |
There was a problem hiding this comment.
This section is a bit confusing to me. Since is_accessible=false actually stops the spank plugin execution flow and has no meaning besides this since the API can still have all the qpu_slots be used and still return is_accessible=true. I'd move qpu_slot info to another route ?
| 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 cast(CursorResult[Any], result).rowcount != 1: | ||
| raise RuntimeError( | ||
| "QPU capacity lock is missing; run the latest Warden database migration." | ||
| ) |
There was a problem hiding this comment.
Is this actually blocking on all QPU backends ?
| await db_session.commit() | ||
| already_revoked = False | ||
| async with db_session.begin(): | ||
| await lock_qpu_capacity(db_session) |
There was a problem hiding this comment.
Is this actually locking? For sqlite we have unique concurrent writes so this will work but for other DBs I don't see the code actually locking. We could use "for update" for a db lock. Alternatively since we have a single API running, an in memory lock should do the trick for now too IIUC
|
This looks very promising to use with licenses! I have a suggestion for a future improvement. Here we are using the number of QPU jobs as an indicator of the usage of the QPU. However in practice QPU jobs runtime can vary greatly depending on nb of qubits, choice of layout, nb of shots, ... To get something more accurate we could instead compute If you like the idea, we can tackle it in a separate MR it shouldn't be too difficult to setup with what we already have here |
|
|
||
| @router.get("") | ||
| async def is_accessible(db_session: DBSessionDep) -> AccessibleResponse: | ||
| async def is_accessible( |
There was a problem hiding this comment.
is_accessible was really meant for the QRMI which is expecting a simple bool. I would recommend introducing a new endpoint to expose the qpu slot data.
3 participants
Add optional Warden-owned QPU capacity tracking.
- Accept and validate qpu_slots on sessions.
- Enforce configured capacity with serialized admission.
- Report total, used, and available slots via /accessible.
- Weight scheduler turns by requested slot count.
- Add tests, configuration, documentation, and reversible migrations
Let's discuss this offline before merging. But the idea is to use this with external schedulers, such as Open Cluster Scheduler or via a dynamic licence manager with Slurm that will help with the multi-tenenacy requirements at CINECA.