diff --git a/tests/api/conftest.py b/tests/api/conftest.py index 582eb50..ca35909 100644 --- a/tests/api/conftest.py +++ b/tests/api/conftest.py @@ -1,5 +1,6 @@ import json from contextlib import contextmanager +from datetime import datetime, timedelta from typing import AsyncGenerator, Callable, Generator import pytest @@ -13,6 +14,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 Job, Session from warden.lib.qpu_client.client import AsyncQPUClient @@ -249,3 +251,54 @@ def qpu_specs() -> dict: specs = json.loads(DigitalAnalogDevice.to_abstract_repr()) specs["name"] = "FRESNEL_CAN1" return specs + + +async def acct_populate_db( + app, + serialized_sequence: str, + n_users: int, + first_session_start: datetime = datetime(2026, 1, 1, 0, 0, 0), + session_duration: timedelta = timedelta(hours=1), + user_time_offset: timedelta = timedelta(hours=1), +) -> tuple[list[str], list[Job], list[Session]]: + """Creates mock data for accounting testing in DB""" + BASE_START_DATETIME = first_session_start + BASE_END_DATETIME = BASE_START_DATETIME + session_duration + user_uids = [str(i) for i in range(1000, 1000 + n_users)] + + # Create at least one session and job per user + sessions = [] + jobs = [] + for i, uid in enumerate(user_uids): + start_session = BASE_START_DATETIME + i * user_time_offset + end_session = BASE_END_DATETIME + i * user_time_offset + + sessions.append( + Session( + created_at=start_session, + revoked_at=end_session, + user_id=uid, + slurm_job_id=str(i), + ) + ) + jobs.append( + Job( + status="DONE", + logs="", + shots=100, + sequence=serialized_sequence, + created_at=start_session, + scheduled_at=start_session, + started_at=start_session, + ended_at=end_session, + # Relationship + session=sessions[-1], + ) + ) + + async_session_factory = app.state.db_session_factory + async with async_session_factory() as db_session: + db_session.add_all(sessions) + db_session.add_all(jobs) + await db_session.commit() + return user_uids, jobs, sessions diff --git a/tests/api/test_acct.py b/tests/api/test_acct.py new file mode 100644 index 0000000..7218bce --- /dev/null +++ b/tests/api/test_acct.py @@ -0,0 +1,161 @@ +from datetime import datetime, timedelta + +import pytest +from httpx import AsyncClient + +from tests.api.conftest import acct_populate_db, mock_munge_auth + +ACCT_ENDPOINT = "/acct" + + +@pytest.mark.asyncio +async def test_acct_required_start(client: AsyncClient, app): + """Assert that 'start_datetime' is required in the accounting data query.""" + + with mock_munge_auth(app, uid=0): + response = await client.get(ACCT_ENDPOINT) + assert response.status_code == 422 + + with mock_munge_auth(app, uid=0): + response = await client.get(ACCT_ENDPOINT + "?end_datetime=2022-01-01T00:00:00") + assert response.status_code == 422 + + with mock_munge_auth(app, uid=0): + response = await client.get( + ACCT_ENDPOINT + "?start_datetime=2022-01-01T00:00:00" + ) + assert response.status_code == 200 + + +@pytest.mark.asyncio +async def test_acct_nominal_pagination_filter(client, app, serialized_sequence): + """Assert that the accounting data query returns the right number of rows.""" + N_USERS = 10 + + user_uids, _, _ = await acct_populate_db(app, serialized_sequence, N_USERS) + + # Test default + with mock_munge_auth(app, uid=0): + response = await client.get( + ACCT_ENDPOINT + "?start_datetime=2020-01-01T00:00:00&limit=100" + ) + assert response.status_code == 200 + assert response.json()["pagination"]["total"] == N_USERS + assert response.json()["pagination"]["start"] == 0 + assert response.json()["pagination"]["end"] == N_USERS + assert len(response.json()["data"]) == N_USERS + assert response.json()["data"][0]["user_id"] == user_uids[0] + + # Test pagination limit + PAGINATION_LIMIT = 5 + assert PAGINATION_LIMIT < N_USERS + + with mock_munge_auth(app, uid=0): + response = await client.get( + ACCT_ENDPOINT + + f"?start_datetime=2020-01-01T00:00:00&limit={PAGINATION_LIMIT}" + ) + assert response.status_code == 200 + assert response.json()["pagination"]["total"] == N_USERS + assert response.json()["pagination"]["start"] == 0 + assert response.json()["pagination"]["end"] == PAGINATION_LIMIT + assert len(response.json()["data"]) == PAGINATION_LIMIT + assert response.json()["data"][0]["user_id"] == user_uids[0] + + # Test offset + PAGINATION_LIMIT = 5 + PAGINATION_OFFSET = 4 + assert PAGINATION_LIMIT + PAGINATION_OFFSET < N_USERS + + with mock_munge_auth(app, uid=0): + response = await client.get( + ACCT_ENDPOINT + + f"?start_datetime=2020-01-01T00:00:00&limit={PAGINATION_LIMIT}&offset={PAGINATION_OFFSET}" + ) + assert response.status_code == 200 + assert response.json()["pagination"]["total"] == N_USERS + assert response.json()["pagination"]["start"] == PAGINATION_OFFSET + assert response.json()["pagination"]["end"] == PAGINATION_LIMIT + PAGINATION_OFFSET + assert len(response.json()["data"]) == PAGINATION_LIMIT + assert response.json()["data"][0]["user_id"] == user_uids[PAGINATION_OFFSET] + + # Test offset end of data length + with mock_munge_auth(app, uid=0): + response = await client.get( + ACCT_ENDPOINT + f"?start_datetime=2020-01-01T00:00:00&offset={N_USERS - 2}" + ) + assert response.status_code == 200 + assert response.json()["pagination"]["total"] == N_USERS + assert response.json()["pagination"]["start"] == N_USERS - 2 + assert response.json()["pagination"]["end"] == N_USERS + assert len(response.json()["data"]) == 2 + assert response.json()["data"][0]["user_id"] == user_uids[N_USERS - 2] + + +@pytest.mark.asyncio +async def test_acct_nominal_datetime_filter(client, app, serialized_sequence): + """Assert that the accounting data query correctly filters according to start/end datetime.""" + + N_USERS = 10 + FIRST_SESSION_START = datetime(2026, 1, 1, 0, 0, 0) + + SESSION_DURATION = timedelta(minutes=45) + USER_TIME_OFFSET = timedelta(hours=1) + + user_ids, _, sessions = await acct_populate_db( + app, + serialized_sequence, + first_session_start=FIRST_SESSION_START, + n_users=N_USERS, + session_duration=SESSION_DURATION, + user_time_offset=USER_TIME_OFFSET, + ) + + # Test default get all + with mock_munge_auth(app, uid=0): + response = await client.get( + ACCT_ENDPOINT + "?start_datetime=" + FIRST_SESSION_START.isoformat() + ) + assert response.status_code == 200 + assert len(response.json()["data"]) == N_USERS + assert response.json()["data"][0]["user_id"] == user_ids[0] + + # Test filtering only first user session + first_session = sessions[0] + assert first_session.revoked_at is not None + with mock_munge_auth(app, uid=0): + response = await client.get( + ACCT_ENDPOINT + + "?start_datetime=" + + FIRST_SESSION_START.isoformat() + + "&end_datetime=" + + (first_session.revoked_at + timedelta(seconds=1)).isoformat() + ) + assert response.status_code == 200 + assert len(response.json()["data"]) == 1 + assert response.json()["data"][0]["user_id"] == user_ids[0] + + # Test time filtering is total partition of data + # We split the filtering exactly at the revocation time of the 2nd user's session + second_session = sessions[1] + assert second_session.revoked_at is not None + with mock_munge_auth(app, uid=0): + response_before = await client.get( + ACCT_ENDPOINT + + "?start_datetime=" + + FIRST_SESSION_START.isoformat() + + "&end_datetime=" + + second_session.revoked_at.isoformat() + + "&limit=100" + ) + response_after = await client.get( + ACCT_ENDPOINT + + "?start_datetime=" + + second_session.revoked_at.isoformat() + + "&limit=100" + ) + assert response_after.status_code == 200 + assert response_before.status_code == 200 + + assert len(response_before.json()["data"]) == 1 + assert len(response_after.json()["data"]) == 9 diff --git a/warden/api/app.py b/warden/api/app.py index 5b8e485..68906af 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, jobs, qpu, sessions +from warden.api.routes import accessible, acct, jobs, qpu, sessions 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 @@ -23,6 +23,7 @@ def create_app(config: Config): app.include_router(sessions.router, tags=["sessions"]) app.include_router(qpu.router, tags=["qpu"]) app.include_router(accessible.router, tags=["accessible"]) + app.include_router(acct.router, tags=["accounting"]) logger = logging.getLogger(__name__) diff --git a/warden/api/routes/acct.py b/warden/api/routes/acct.py new file mode 100644 index 0000000..1b810c2 --- /dev/null +++ b/warden/api/routes/acct.py @@ -0,0 +1,141 @@ +"""Accounting information route""" + +from logging import getLogger +from typing import Annotated + +from fastapi import APIRouter, Query +from sqlalchemy import and_, func, select + +from warden.api.routes.dependencies.auth import ( + AdminUserDep, +) +from warden.api.routes.dependencies.db import DBSessionDep +from warden.api.schemas.acct import ( + AcctData, + GetAcctRequest, + GetAcctResponse, + JobsSummary, + JobSummaryStats, + PaginationResponse, + SessionsSummary, +) +from warden.lib.models import Job, Session + +logger = getLogger(__name__) +router = APIRouter(prefix="/acct") + + +@router.get("") +async def get_accounting_snapshot( + params: Annotated[GetAcctRequest, Query()], + db_session: DBSessionDep, + _admin: AdminUserDep, +) -> GetAcctResponse: + """High-level accounting data""" + + # Base session filter + session_filters = [Session.revoked_at >= params.start_datetime] + + if params.end_datetime: + session_filters.append(Session.revoked_at < params.end_datetime) + + if params.user_ids: + session_filters.append(Session.user_id.in_(params.user_ids)) + + # Get total count for pagination + total_count_stmt = select(func.count(func.distinct(Session.user_id))).where( + and_(*session_filters) + ) + total_result = await db_session.execute(total_count_stmt) + total_count = total_result.scalar() or 0 + + # Query to get sessions data aggregated by user_id + sessions_stmt = ( + select( + Session.user_id, + func.count(Session.id).label("session_count"), + # Check + func.coalesce( + func.sum( + func.extract("epoch", Session.revoked_at - Session.created_at) + ), + 0, + ).label("total_session_duration"), + ) + .where(and_(*session_filters)) + .group_by(Session.user_id) + .order_by(Session.user_id) + .offset(params.offset) + .limit(params.limit) + ) + + sessions_result = await db_session.execute(sessions_stmt) + sessions_data = sessions_result.fetchall() + + user_sessions_summary: dict[str, SessionsSummary] = {} + + for session_row in sessions_data: + user_id = session_row.user_id + user_sessions_summary[user_id] = SessionsSummary( + count=session_row.session_count, + total_duration=session_row.total_session_duration, + ) + + # Query to get jobs data aggregated by user_id and job status + jobs_stmt = ( + select( + Session.user_id, + Job.status, + func.count(Job.id).label("job_count"), + func.coalesce( + func.sum(func.extract("epoch", Job.ended_at - Job.started_at)), 0 + ).label("total_job_duration"), + ) + .join(Session, Session.id == Job.session_id, isouter=True) + .where(and_(*session_filters)) + .group_by(Session.user_id, Job.status) + .order_by(Session.user_id) + .offset(params.offset) + .limit(params.limit) + ) + + jobs_result = await db_session.execute(jobs_stmt) + jobs_data = jobs_result.fetchall() + + user_jobs_summaries: dict[str, JobsSummary] = {} + + for job_row in jobs_data: + user_id = job_row.user_id + if user_id not in user_jobs_summaries: + user_jobs_summaries[user_id] = JobsSummary(count=0) + user_jobs_summaries[user_id].count += job_row.job_count + user_jobs_summaries[user_id].stats.append( + JobSummaryStats( + status=job_row.status, + count=job_row.job_count, + total_duration=int(job_row.total_job_duration or 0), + ) + ) + + # For each user, get job statistics + acct_data_list = [] + + for user_id in user_sessions_summary.keys(): + user_acct_data = AcctData( + user_id=user_id, + sessions=user_sessions_summary[user_id], + jobs=user_jobs_summaries.get(user_id, JobsSummary(count=0)), + ) + + acct_data_list.append(user_acct_data) + + pagination_resp = PaginationResponse( + total=total_count, + start=params.offset, + end=min(params.offset + len(acct_data_list), total_count), + ) + + return GetAcctResponse( + data=acct_data_list, + pagination=pagination_resp, + ) diff --git a/warden/api/schemas/acct.py b/warden/api/schemas/acct.py new file mode 100644 index 0000000..6a9c889 --- /dev/null +++ b/warden/api/schemas/acct.py @@ -0,0 +1,100 @@ +from datetime import datetime +from typing import Any + +from pydantic import BaseModel, Field + +from warden.api.schemas.common import JobID, SessionID, UserID + + +class PaginationResponse(BaseModel): + total: int + start: int + end: int + + +class SessionsSummary(BaseModel): + count: int + total_duration: int + + +class JobSummaryStats(BaseModel): + status: str + count: int + total_duration: int + + +class JobsSummary(BaseModel): + count: int = Field(0, ge=0) + stats: list[JobSummaryStats] = Field(default_factory=list) + + +class AcctData(BaseModel): + user_id: UserID + sessions: SessionsSummary + jobs: JobsSummary + + +class AcctUserData(BaseModel): + pass + + +class SessionsData(BaseModel): + id: SessionID + user_id: UserID + created_at: datetime + revoked_at: datetime | None + slurm_job_id: str + + +class JobData(BaseModel): + id: JobID + user_id: UserID + session_id: SessionID + execution_time: int + status: str + + +# Common parameters for request and response +class AcctRequest(BaseModel): + # Time filtering + start_datetime: datetime + end_datetime: datetime | None = Field(default=None) + # Pagination + limit: int = Field(100, gt=0, le=100) + offset: int = Field(0, ge=0) + + +class AcctResponse(BaseModel): + data: Any + pagination: PaginationResponse + + +# Routes + + +# GET /acct +class GetAcctRequest(AcctRequest): + user_ids: set[UserID] | None = Field(default=None) + + +class GetAcctResponse(AcctResponse): + data: list[AcctData] + + +# GET /acct/user/{user_id} +class GetAcctUserRequest(AcctRequest): + pass + + +class GetAcctUserResponse(AcctResponse): + pass + + +# GET /acct/jobs +class GetAcctJobsRequest(AcctRequest): + session_id: SessionID | None + user_id: UserID | None + + +class GetAcctJobsResponse(AcctResponse): + data: list[JobData] diff --git a/warden/api/schemas/common.py b/warden/api/schemas/common.py new file mode 100644 index 0000000..b3770f6 --- /dev/null +++ b/warden/api/schemas/common.py @@ -0,0 +1,8 @@ +"""Common definitions""" + +from typing import TypeAlias +from uuid import UUID + +SessionID: TypeAlias = UUID +UserID: TypeAlias = str +JobID: TypeAlias = int diff --git a/warden/api/schemas/jobs.py b/warden/api/schemas/jobs.py index 92928f7..0fa8dae 100644 --- a/warden/api/schemas/jobs.py +++ b/warden/api/schemas/jobs.py @@ -4,6 +4,7 @@ from pydantic import BaseModel, Field, ValidationError +from warden.api.schemas.common import UserID from warden.lib.models.jobs import Job @@ -23,7 +24,7 @@ class JobCreate(BaseModel): class JobResponse(BaseModel): id: int - user_id: str + user_id: UserID created_at: datetime status: str results: str | None diff --git a/warden/api/schemas/sessions.py b/warden/api/schemas/sessions.py index f1a23f2..8413fba 100644 --- a/warden/api/schemas/sessions.py +++ b/warden/api/schemas/sessions.py @@ -1,8 +1,8 @@ from datetime import datetime -from uuid import UUID from pydantic import BaseModel +from warden.api.schemas.common import SessionID, UserID from warden.lib.models.sessions import Session @@ -12,8 +12,8 @@ class CreateSession(BaseModel): class SessionResponse(BaseModel): - id: UUID - user_id: str + id: SessionID + user_id: UserID created_at: datetime revoked_at: datetime | None