diff --git a/Makefile b/Makefile index d54721d..815fba7 100644 --- a/Makefile +++ b/Makefile @@ -5,7 +5,7 @@ include config.mk dev.mk # cluster admin commands to operate Warden .PHONY: set-accessible \ - ping get-logs + ping get-logs get-status # Mock QPU when the actual QPU is not available .PHONY: start-mock-qpu start-qutip-qpu @@ -137,6 +137,11 @@ get-logs: curl -X GET $(URL)/jobs/$(ID)/logs \ -H "X-Munge-Cred: $$(munge -n)" +get-status: + + curl -s -X GET $(URL)/status\ + -H "X-Munge-Cred: $$(munge -n)" + ping: curl $(URL) diff --git a/README.md b/README.md index 3761a79..6a243b1 100644 --- a/README.md +++ b/README.md @@ -136,6 +136,15 @@ and `GET /accounting/jobs` accepts `session_id` and `status` filtering. See FastAPI's `/docs` page for the full request/response schema of each route. +### Status API + +`GET /status` returns a live snapshot of Warden's current activity, restricted +to `api.admin_users`: +- the number of pending jobs +- the job currently running +- and the list of open sessions. +Historical jobs/sessions are not included, see the Accounting API above for reporting. + ### Database Warden supports the following databases: diff --git a/tests/api/test_status.py b/tests/api/test_status.py new file mode 100644 index 0000000..d676933 --- /dev/null +++ b/tests/api/test_status.py @@ -0,0 +1,106 @@ +from datetime import datetime, timedelta, timezone + +import pytest + +from tests.api.conftest import mock_munge_auth +from warden.lib.models import Job, Session + + +async def _seed(app, *, serialized_sequence: str): + async_session = app.state.db_session_factory + + open_session = Session( + user_id="1000", slurm_job_id="1", created_at=datetime.now(timezone.utc) + ) + other_sessions = [ + Session( + user_id="1000", + slurm_job_id="1", + created_at=(datetime.now(timezone.utc) - timedelta(minutes=i)), + ) + for i in range(5, 0, -1) + ] + revoked_session = Session( + user_id="1000", slurm_job_id="2", revoked_at=datetime.now(timezone.utc) + ) + + async with async_session() as session: + session.add_all([open_session, *other_sessions, revoked_session]) + await session.commit() + + pending_jobs = [ + Job(sequence=serialized_sequence, shots=100, session=open_session) + for _ in range(10) + ] + running_job = Job( + sequence=serialized_sequence, + shots=100, + session=open_session, + status="RUNNING", + scheduled_at=datetime.now(), + started_at=datetime.now(), + ) + done_job = Job( + sequence=serialized_sequence, + shots=100, + session=open_session, + status="DONE", + scheduled_at=datetime.now(), + started_at=datetime.now(), + ended_at=datetime.now(), + ) + + async with async_session() as session: + session.add_all([*pending_jobs, running_job, done_job]) + await session.commit() + await session.refresh(running_job) + + return open_session, other_sessions, running_job + + +@pytest.mark.asyncio +async def test_get_status_non_admin(client, app): + with mock_munge_auth(app, uid=1001): + response = await client.get("/status") + assert response.status_code == 403 + + +@pytest.mark.asyncio +async def test_get_status_no_auth(client): + response = await client.get("/status") + assert response.status_code == 401 + + +@pytest.mark.asyncio +async def test_get_status_admin(client, app, serialized_sequence: str): + """A PENDING, RUNNING and DONE job plus several open sessions and a + revoked session are seeded; the snapshot must count only the PENDING job, + surface the RUNNING job as current_job (ignoring DONE), and list the open + sessions ordered by creation time, oldest first.""" + open_session, other_sessions, running_job = await _seed( + app, serialized_sequence=serialized_sequence + ) + + with mock_munge_auth(app, uid=0): + response = await client.get("/status") + assert response.status_code == 200 + data = response.json() + + assert data["pending_jobs_count"] == 10 + assert data["current_job"]["id"] == running_job.id + assert data["current_job"]["session_id"] == str(open_session.id) + + expected_order = [str(s.id) for s in other_sessions] + [str(open_session.id)] + assert [s["id"] for s in data["open_sessions"]] == expected_order + + +@pytest.mark.asyncio +async def test_get_status_no_current_job(client, app): + """With an empty DB, current_job is null and both counts/lists are empty.""" + with mock_munge_auth(app, uid=0): + response = await client.get("/status") + assert response.status_code == 200 + data = response.json() + assert data["pending_jobs_count"] == 0 + assert data["current_job"] is None + assert data["open_sessions"] == [] diff --git a/warden/api/app.py b/warden/api/app.py index 9018b0a..ff3ba1e 100644 --- a/warden/api/app.py +++ b/warden/api/app.py @@ -2,7 +2,7 @@ from fastapi import FastAPI -from warden.api.routes import accessible, acct, jobs, qpu, sessions +from warden.api.routes import accessible, acct, jobs, qpu, sessions, status from warden.api.routes.dependencies.auth import init_auth from warden.api.routes.dependencies.db import init_db from warden.api.routes.dependencies.qpu_client import init_qpu_client @@ -32,6 +32,7 @@ def create_app(config: Config): app.include_router(qpu.router, tags=["qpu"]) app.include_router(accessible.router, tags=["accessible"]) app.include_router(acct.router, tags=["accounting"]) + app.include_router(status.router, tags=["status"]) logger = logging.getLogger(__name__) diff --git a/warden/api/routes/status.py b/warden/api/routes/status.py new file mode 100644 index 0000000..3c27c67 --- /dev/null +++ b/warden/api/routes/status.py @@ -0,0 +1,57 @@ +from fastapi import APIRouter +from sqlalchemy import func, select + +from warden.api.routes.dependencies.auth import AdminUserDep +from warden.api.routes.dependencies.db import DBSessionDep +from warden.api.schemas.status import CurrentJob, OpenSession, StatusResponse +from warden.lib.models import Job, Session + +router = APIRouter(prefix="/status") + + +@router.get("") +async def get_status( + db_session: DBSessionDep, + _admin: AdminUserDep, +) -> StatusResponse: + """ + Provide live snapshot of Warden's current activity: + - count of pending jobs + - the job currently executing + - open sessions + """ + pending_jobs_count = await db_session.scalar( + select(func.count(Job.id)).where(Job.status == "PENDING") + ) + + running_job = await db_session.scalar(select(Job).where(Job.status == "RUNNING")) + current_job = ( + CurrentJob( + id=running_job.id, + session_id=running_job.session_id, + user_id=running_job.user_id, + started_at=running_job.started_at, + backend_id=running_job.backend_id, + ) + if running_job is not None + else None + ) + + open_sessions_result = await db_session.execute( + select(Session).where(Session.revoked_at.is_(None)).order_by(Session.created_at) + ) + open_sessions = [ + OpenSession( + id=session.id, + user_id=session.user_id, + created_at=session.created_at, + slurm_job_id=session.slurm_job_id, + ) + for session in open_sessions_result.scalars() + ] + + return StatusResponse( + pending_jobs_count=pending_jobs_count or 0, + current_job=current_job, + open_sessions=open_sessions, + ) diff --git a/warden/api/schemas/status.py b/warden/api/schemas/status.py new file mode 100644 index 0000000..523da92 --- /dev/null +++ b/warden/api/schemas/status.py @@ -0,0 +1,26 @@ +from datetime import datetime + +from pydantic import BaseModel + +from warden.api.schemas.common import JobID, SessionID, UserID + + +class CurrentJob(BaseModel): + id: JobID + session_id: SessionID + user_id: UserID + started_at: datetime | None + backend_id: str | None + + +class OpenSession(BaseModel): + id: SessionID + user_id: UserID + created_at: datetime + slurm_job_id: str + + +class StatusResponse(BaseModel): + pending_jobs_count: int + current_job: CurrentJob | None + open_sessions: list[OpenSession]