diff --git a/tests/scheduler/test_scheduler.py b/tests/scheduler/test_scheduler.py index 2ad8e70..97f03d8 100644 --- a/tests/scheduler/test_scheduler.py +++ b/tests/scheduler/test_scheduler.py @@ -172,6 +172,189 @@ async def test_run_nominal( assert len(job.logs) > 0 +@pytest.mark.asyncio +@pytest.mark.parametrize("strategy", list(SchedulerStrategy)) +async def test_run_resume_job( + strategy: SchedulerStrategy, + db_engine: AsyncEngine, + db_session_maker: async_sessionmaker, + httpx_mock: HTTPXMock, + caplog, +): + """Test that the scheduler is able to resume + execution of a job when Warden is restarted and a job is already running + on the QPU + + Test rationale: + - Create 3 dummy jobs with a backend_id already set: + - One with a valid backend_id (BACKEND_ID) + - One with a non-existing backend_id (NON_EXISTING_BACKEND_ID) + - One with a valid backend_id whose execution is already "DONE" on the QPU + - QPU API is mocked: + - To return QPU status as "UP" + - To return 400 for the non-existing backend_id job status request + - To accept job creation request for the re-created job (NEW_BACKEND_ID) + - To return "RUNNING" and then "DONE" status for BACKEND_ID and NEW_BACKEND_ID + - To return "DONE" status for ALREADY_DONE_BACKEND_ID + - Run scheduler until: + - All jobs have a "DONE" status in DB + - Test timeout after TEST_TIMEOUT_S + - Check n (jobs with status "DONE") = 3 + - Check "DONE" jobs have the right results and non-empty logs + - Check those jobs have `scheduled_at` set + """ + + ################## + ### TEST CONF ### + ################## + + # Enable warden logging for jobs 'logs' field to be populated + caplog.set_level(logging.INFO, logger="warden") + + TEST_TIMEOUT_S = 5 + NORMAL_BACKEND_ID = "1" + NON_EXISTING_BACKEND_ID = "9999" + NEW_BACKEND_ID = "2" + ALREADY_DONE_BACKEND_ID = "3" + N_JOBS = 3 + + conf: Config = build_conf(strategy, QPU_URI) + + ################## + ### TEST SETUP ### + ################## + + # QPU status + httpx_mock.add_response( + method="GET", + url=SYSTEM_OPERATIONAL_API, + json={"data": {"operational_status": "UP"}}, + is_reusable=True, + ) + # Non-existing backend id + httpx_mock.add_response( + method="GET", + url=JOB_API + f"/{NON_EXISTING_BACKEND_ID}", + status_code=400, + json={"message": "Bad request."}, + ) + # Creating a new job for non-existing backend + return_create_json = { + "data": { + "uid": NEW_BACKEND_ID, + "batch_id": SLURM_USER_ID, + "status": "PENDING", + "result": None, + "program_id": QPU_PROGRAM_UID, + "created_datetime": NOW.isoformat(), + "start_datetime": None, + "end_datetime": None, + } + } + httpx_mock.add_response( + method="POST", status_code=200, url=JOB_API, json=return_create_json + ) + # Nominal running scenario + for backend_id in [NORMAL_BACKEND_ID, NEW_BACKEND_ID]: + return_running_json = { + "data": { + "uid": backend_id, + "batch_id": SLURM_USER_ID, + "status": "RUNNING", + "result": None, + "program_id": QPU_PROGRAM_UID, + "created_datetime": NOW.isoformat(), + "start_datetime": (NOW + timedelta(seconds=1)).isoformat(), + "end_datetime": None, + } + } + + return_done_json = { + "data": { + "uid": backend_id, + "batch_id": SLURM_USER_ID, + "status": "DONE", + "result": DUMMY_RESULTS, + "program_id": QPU_PROGRAM_UID, + "created_datetime": NOW.isoformat(), + "start_datetime": (NOW + timedelta(seconds=1)).isoformat(), + "end_datetime": (NOW + timedelta(seconds=2)).isoformat(), + } + } + # Job running + for _ in range(2): + httpx_mock.add_response( + method="GET", + status_code=200, + url=JOB_API + f"/{backend_id}", + json=return_running_json, + ) + # Job done + httpx_mock.add_response( + method="GET", + status_code=200, + url=JOB_API + f"/{backend_id}", + json=return_done_json, + ) + # Job that was already done + return_done_json = { + "data": { + "uid": ALREADY_DONE_BACKEND_ID, + "batch_id": SLURM_USER_ID, + "status": "DONE", + "result": DUMMY_RESULTS, + "program_id": QPU_PROGRAM_UID, + "created_datetime": NOW.isoformat(), + "start_datetime": (NOW + timedelta(seconds=1)).isoformat(), + "end_datetime": (NOW + timedelta(seconds=2)).isoformat(), + } + } + httpx_mock.add_response( + method="GET", + status_code=200, + url=JOB_API + f"/{ALREADY_DONE_BACKEND_ID}", + json=return_done_json, + is_reusable=True, + ) + + # Populate DB with jobs to run + await utils.create_n_jobs( + db_session_maker, + 3, + backend_ids=[ + NORMAL_BACKEND_ID, + NON_EXISTING_BACKEND_ID, + ALREADY_DONE_BACKEND_ID, + ], + ) + + stmt_count = select(func.count(Job.id)).where(Job.status == "DONE") + + ################## + ### TEST RUN ### + ################## + + # RUN SCHEDULER + main_task = asyncio.create_task(run_scheduler(db_engine, conf)) + + async with db_session_maker() as session: + async with utils.scheduler_task_timeout(TEST_TIMEOUT_S, main_task): + await utils.wait_until_scalar_equals( + session, + stmt_count, + N_JOBS, + interval=SUCCESS_CHECK_INTERVAL_S, + ) + + stmt_all = select(Job).where(Job.status == "DONE") + jobs_done = (await session.execute(stmt_all)).scalars().all() + assert len(jobs_done) == N_JOBS + for job in jobs_done: + assert job.scheduled_at is not None + assert job.results == DUMMY_RESULTS + assert len(job.logs) > 0 + + @pytest.mark.asyncio @pytest.mark.parametrize("strategy", list(SchedulerStrategy)) async def test_run_qpu_down( @@ -484,6 +667,136 @@ async def test_run_job_timeout( assert len(job.logs) > 0 +@pytest.mark.asyncio +@pytest.mark.parametrize("strategy", list(SchedulerStrategy)) +async def test_run_resume_job_timeout( + strategy: SchedulerStrategy, + db_engine: AsyncEngine, + db_session_maker: async_sessionmaker, + httpx_mock: HTTPXMock, + caplog, +): + """Test that the scheduler is able to resume + execution of a job that is supposed to be timedout + + Test rationale: + - Create 1 dummy jobs with a backend_id and a created_at timestamp 120s ago + - QPU API is mocked: + - To return the status of the job as "RUNNING" + - Accept the job's cancelation request + - Run scheduler until: + - All jobs have a "DONE" status in DB + - Test timeout after TEST_TIMEOUT_S + - Check n (jobs with status "CANCELED") == 1 + - Check "CANCELED" jobs have the right results and non-empty logs + - Check those jobs have `scheduled_at` set + """ + + ################## + ### TEST CONF ### + ################## + + # Enable warden logging for jobs 'logs' field to be populated + caplog.set_level(logging.INFO, logger="warden") + + TEST_TIMEOUT_S = 5 + N_JOBS = 1 + BACKEND_ID = "1" + # Setting the job's created_at at a time that is already timedout + JOB_CREATED_AT = NOW - timedelta(seconds=120) + EXPECTED_JOB_STATUS = "CANCELED" + + conf: Config = build_conf(strategy, QPU_URI) + # Setting the scheduler's config to a timeout shorter than the job's age + conf.scheduler.job_polling_timeout_s = 60 + + ################## + ### TEST SETUP ### + ################## + + # QPU status + httpx_mock.add_response( + method="GET", + url=SYSTEM_OPERATIONAL_API, + json={"data": {"operational_status": "UP"}}, + is_reusable=True, + ) + # Initial job status probe + return_running_json = { + "data": { + "uid": BACKEND_ID, + "batch_id": SLURM_USER_ID, + "status": "RUNNING", + "result": None, + "program_id": QPU_PROGRAM_UID, + "created_datetime": JOB_CREATED_AT.isoformat(), + "start_datetime": (JOB_CREATED_AT + timedelta(seconds=1)).isoformat(), + "end_datetime": None, + } + } + for _ in range(2): + httpx_mock.add_response( + method="GET", + status_code=200, + url=JOB_API + f"/{BACKEND_ID}", + json=return_running_json, + ) + # Job cancellation requests + return_cancelled_job_status = { + "data": { + "uid": BACKEND_ID, + "batch_id": SLURM_USER_ID, + "status": "CANCELED", + "result": None, + "program_id": QPU_PROGRAM_UID, + # TODO: CHECK RETURN FIELDS HERE + "created_datetime": JOB_CREATED_AT.isoformat(), + "start_datetime": (JOB_CREATED_AT + timedelta(seconds=1)).isoformat(), + "end_datetime": None, + } + } + httpx_mock.add_response( + method="PUT", + status_code=200, + url=JOB_API + f"/{BACKEND_ID}/cancel", + json=return_cancelled_job_status, + ) + + # Populate DB with jobs to run + await utils.create_n_jobs( + db_session_maker, + n_jobs=1, + backend_ids=[BACKEND_ID], + ) + + stmt_count = select(func.count(Job.id)).where(Job.status == EXPECTED_JOB_STATUS) + + ################## + ### TEST RUN ### + ################## + + # RUN SCHEDULER + main_task = asyncio.create_task(run_scheduler(db_engine, conf)) + + async with db_session_maker() as session: + async with utils.scheduler_task_timeout(TEST_TIMEOUT_S, main_task): + await utils.wait_until_scalar_equals( + session, + stmt_count, + N_JOBS, + interval=SUCCESS_CHECK_INTERVAL_S, + ) + + stmt_all = select(Job).where(Job.status == EXPECTED_JOB_STATUS) + jobs_done = (await session.execute(stmt_all)).scalars().all() + assert len(jobs_done) == N_JOBS + for job in jobs_done: + assert job.scheduled_at is not None + assert job.results is None + assert len(job.logs) > 0 + assert "CANCELED" in job.logs + + @pytest.mark.asyncio @pytest.mark.parametrize("strategy", list(SchedulerStrategy)) async def test_run_retry_transient_errors( diff --git a/tests/scheduler/utils.py b/tests/scheduler/utils.py index 783cfd0..255503d 100644 --- a/tests/scheduler/utils.py +++ b/tests/scheduler/utils.py @@ -5,6 +5,7 @@ from contextlib import asynccontextmanager from typing import Any +import pytest from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker from warden.lib.config import Config, QPUConfig, SchedulerConfig, SchedulerStrategy @@ -32,19 +33,26 @@ async def wait_until_scalar_equals( async def create_n_jobs( - db_session_maker: async_sessionmaker, n_jobs: int, shots: int = 100 + db_session_maker: async_sessionmaker, + n_jobs: int, + shots: int = 100, + backend_ids: list[str] | None = None, ) -> list[Job]: """Creates n_jobs mock jobs to run in the warden db""" SLURM_USER_ID = "1234" + if backend_ids and len(backend_ids) != n_jobs: + raise Exception("Length of 'backend_ids' should match 'n_jobs' parameter") + jobs_to_run = [ Job( sequence="{}", status="PENDING", shots=shots, session=Session(slurm_job_id="1", user_id=SLURM_USER_ID), + backend_id=backend_ids[i] if backend_ids else None, ) - for _ in range(n_jobs) + for i in range(n_jobs) ] async with db_session_maker() as session: @@ -79,14 +87,40 @@ def raise_task_exception(async_task: Task) -> None: @asynccontextmanager async def scheduler_task_timeout(delay: float, scheduler_task: Task): - """Manage test timeout and cleanup main scheduler task on test timeout""" + """Manage test timeout and cleanup main scheduler task on test end. + + On timeout, fails the test with a descriptive pytest message. + If the scheduler task itself raised an exception, that is reported + as the cause so the real error is immediately visible. + """ try: async with timeout(delay): yield except TimeoutError: - # cleanup - raise_task_exception(scheduler_task) + # Check whether the scheduler crashed + scheduler_exc: BaseException | None = None + if scheduler_task.done() and not scheduler_task.cancelled(): + scheduler_exc = scheduler_task.exception() + + if scheduler_exc is not None: + pytest.fail( + f"Test timed out after {delay}s because the scheduler task " + f"raised an exception: {type(scheduler_exc).__name__}: {scheduler_exc}" + ) + else: + pytest.fail( + f"Test timed out after {delay}s waiting for the scheduler " + "to reach the expected state." + ) + finally: + # Always cancel the scheduler task so it does not keep running after + # the test body finishes and fixture teardown begins + scheduler_task.cancel() + try: + await scheduler_task + except (asyncio.CancelledError, Exception): + pass def build_conf(strategy: SchedulerStrategy, qpu_uri: str) -> Config: diff --git a/tests/test_qpu_client_types.py b/tests/test_qpu_client_types.py new file mode 100644 index 0000000..4a5dafb --- /dev/null +++ b/tests/test_qpu_client_types.py @@ -0,0 +1,86 @@ +"""Testing lib/qpu_client/types — UTCDatetime Pydantic validator""" + +from datetime import datetime, timedelta, timezone + +from warden.lib.qpu_client.types import QPUJobInfo + + +def _api_data(**overrides) -> dict: + """Return a minimal valid QPUJobInfo payload.""" + data = dict( + uid=1, + batch_id=None, + status=None, + result=None, + program_id=None, + created_datetime=datetime(2024, 1, 15, 12, 0, 0), + start_datetime=None, + end_datetime=None, + ) + data.update(overrides) + return data + + +def test_naive_datetime_is_coerced_to_utc(): + """A naive created_datetime is assumed UTC and given UTC tzinfo.""" + naive = datetime.now() + assert naive.tzinfo is None + + job = QPUJobInfo(**_api_data(created_datetime=naive)) + + assert job.created_datetime.tzinfo == timezone.utc + assert job.created_datetime.replace(tzinfo=None) == naive + + +def test_utc_aware_datetime_is_unchanged(): + """A UTC-aware created_datetime passes through without modification.""" + utc_dt = datetime.now(tz=timezone.utc) + + job = QPUJobInfo(**_api_data(created_datetime=utc_dt)) + + assert job.created_datetime == utc_dt + assert job.created_datetime.tzinfo == timezone.utc + + +def test_non_utc_aware_datetime_is_converted_to_utc(): + """A non-UTC timezone-aware created_datetime is converted to UTC.""" + minus_five = timezone(timedelta(hours=-5)) + minus_five_dt = datetime(2024, 1, 15, 7, 0, 0, tzinfo=minus_five) + expected_utc = datetime(2024, 1, 15, 12, 0, 0, tzinfo=timezone.utc) + + job = QPUJobInfo(**_api_data(created_datetime=minus_five_dt)) + + assert job.created_datetime == expected_utc + assert job.created_datetime.tzinfo == timezone.utc + + +def test_optional_utc_datetime_field_none_is_accepted(): + """Optional UTCDatetime fields (start/end) accept None without error.""" + job = QPUJobInfo(**_api_data(start_datetime=None, end_datetime=None)) + + assert job.start_datetime is None + assert job.end_datetime is None + + +def test_optional_utc_datetime_field_naive_is_coerced(): + """Optional UTCDatetime fields coerce naive datetimes to UTC.""" + naive = datetime.now() + + job = QPUJobInfo(**_api_data(start_datetime=naive, end_datetime=naive)) + + assert job.start_datetime is not None + assert job.start_datetime.tzinfo == timezone.utc + assert job.end_datetime is not None + assert job.end_datetime.tzinfo == timezone.utc + assert job.start_datetime.replace(tzinfo=None) == naive + + +def test_optional_utc_datetime_field_non_utc_is_converted(): + """Optional UTCDatetime fields convert non-UTC aware datetimes to UTC.""" + plus_two = timezone(timedelta(hours=2)) + local_dt = datetime(2024, 6, 1, 10, 0, 0, tzinfo=plus_two) + expected_utc = datetime(2024, 6, 1, 8, 0, 0, tzinfo=timezone.utc) + + job = QPUJobInfo(**_api_data(end_datetime=local_dt)) + + assert job.end_datetime == expected_utc diff --git a/warden/lib/qpu_client/__init__.py b/warden/lib/qpu_client/__init__.py index 11bf280..e8d3427 100644 --- a/warden/lib/qpu_client/__init__.py +++ b/warden/lib/qpu_client/__init__.py @@ -5,6 +5,7 @@ QPUInfo, QPUJobInfo, QPUOperationalStatus, + UTCDatetime, ) __all__ = [ @@ -16,4 +17,5 @@ "QPUClientRequestError", "JobCancelationError", "JobStatus", + "UTCDatetime", ] diff --git a/warden/lib/qpu_client/types.py b/warden/lib/qpu_client/types.py index 8c3abe8..881d41b 100644 --- a/warden/lib/qpu_client/types.py +++ b/warden/lib/qpu_client/types.py @@ -1,10 +1,23 @@ """Useful data types for QPU API parsing""" -from datetime import datetime -from typing import Literal, TypeAlias +from datetime import datetime, timezone +from typing import Annotated, Literal, TypeAlias +from pydantic import AfterValidator from pydantic.dataclasses import dataclass + +def _as_utc(value: datetime | None) -> datetime | None: + """Ensure a datetime is UTC-aware. Naive datetimes from the API are assumed UTC.""" + if value is None: + return None + if value.tzinfo is None: + return value.replace(tzinfo=timezone.utc) + return value.astimezone(timezone.utc) + + +UTCDatetime = Annotated[datetime, AfterValidator(_as_utc)] + JobStatus: TypeAlias = Literal["PENDING", "RUNNING", "ERROR", "CANCELED", "DONE"] QPUStatus: TypeAlias = Literal["UP", "DOWN"] @@ -16,9 +29,9 @@ class QPUJobInfo: status: JobStatus | None result: str | None program_id: int | None - created_datetime: datetime - start_datetime: datetime | None - end_datetime: datetime | None + created_datetime: UTCDatetime + start_datetime: UTCDatetime | None + end_datetime: UTCDatetime | None @dataclass diff --git a/warden/scheduler/main.py b/warden/scheduler/main.py index 1295a25..90079a4 100644 --- a/warden/scheduler/main.py +++ b/warden/scheduler/main.py @@ -75,6 +75,7 @@ async def run_scheduler(engine: AsyncEngine, conf: Config): queue=queue, nb_run=job.shots, sequence=job.sequence, + backend_id=job.backend_id, batch_id=job.session.slurm_job_id, ), name=f"Job {job.id} execution worker", diff --git a/warden/scheduler/worker.py b/warden/scheduler/worker.py index ba55ecf..a29a2a1 100644 --- a/warden/scheduler/worker.py +++ b/warden/scheduler/worker.py @@ -3,7 +3,7 @@ import asyncio import logging from contextlib import contextmanager -from datetime import datetime +from datetime import datetime, timezone from warden.lib.config import Config from warden.lib.qpu_client import ( @@ -12,6 +12,7 @@ QPUClient, QPUClientRequestError, QPUJobInfo, + UTCDatetime, ) from warden.scheduler.errors import QPUDownError from warden.scheduler.types import JobUpdate, JobUpdateQueue @@ -43,6 +44,10 @@ def job(self) -> QPUJobInfo: def is_error(self) -> bool: return self.status == "ERROR" + @property + def created_datetime(self) -> UTCDatetime: + return self.job.created_datetime + async def update_job(self, qpu_job_info: QPUJobInfo): self._qpu_job_info = qpu_job_info self._status = qpu_job_info.status or "ERROR" @@ -128,16 +133,17 @@ def is_operational(self): return self.operational_status == "UP" @staticmethod - def is_timed_out(timeout_s: int | float, start: datetime) -> bool: + def is_timed_out(timeout_s: int | float, start: UTCDatetime) -> bool: if timeout_s < 0: return False - return (datetime.now() - start).total_seconds() > timeout_s + return (datetime.now(tz=timezone.utc) - start).total_seconds() > timeout_s async def execute_job( self, queue: JobUpdateQueue, nb_run: int, sequence: str, + backend_id: str | None = None, batch_id: str | None = None, ) -> None: """Execute job on the QPU""" @@ -153,11 +159,11 @@ async def execute_job( job_tracker=job_tracker, nb_run=nb_run, sequence=sequence, + backend_id=backend_id, batch_id=batch_id, ) if job_tracker.is_error: return - logger.info("Job created on QPU") await self.await_job_execution(job_tracker) logger.info("Job execution ended with status '%s'", job_tracker.status) @@ -168,7 +174,7 @@ async def execute_job( async def poll_qpu(self, job_tracker: JobExecutionTracker) -> None: """Check the QPU status""" try: - polling_start = datetime.now() + polling_start = datetime.now(tz=timezone.utc) while not self.is_operational: if self.is_timed_out( self.conf_sched.qpu_polling_timeout_s, polling_start @@ -195,9 +201,28 @@ async def create_job( job_tracker: JobExecutionTracker, nb_run: int, sequence: str, + backend_id: str | None, batch_id: str | None, ) -> None: """Create the job on the QPU""" + + # Try to re-fetch the job status on the QPU + if backend_id and backend_id.isdigit(): + logger.info( + "Job already has a backend_id '%s'. Trying to fetch status from QPU.", + backend_id, + ) + try: + qpu_job_info = self.qpu_client.get_job(job_uid=int(backend_id)) + await job_tracker.update_job(qpu_job_info) + logger.info("Job status fetched from QPU: %s", job_tracker.status) + return + except QPUClientRequestError as e: + logger.warning( + "Failed fetching job status: %s. Re-starting corresponding job on the QPU", + e, + ) + try: qpu_job_info = self.qpu_client.create_job( nb_run=nb_run, @@ -206,13 +231,15 @@ async def create_job( ) await job_tracker.update_job(qpu_job_info) + logger.info("Job created on QPU") except QPUClientRequestError as e: logger.error(f"Failed creating job: {e}") await job_tracker.to_error() async def await_job_execution(self, job_tracker: JobExecutionTracker) -> None: """Polling the job status until completion, error, or cancellation""" - polling_start = datetime.now() + + polling_start = job_tracker.created_datetime await self._get_job_poll(job_tracker) while job_tracker.status not in ("ERROR", "DONE", "CANCELED"): if self.is_timed_out(self.conf_sched.job_polling_timeout_s, polling_start):