diff --git a/tests/lib/qpu_client/test_auth.py b/tests/lib/qpu_client/test_auth.py new file mode 100644 index 0000000..28c7677 --- /dev/null +++ b/tests/lib/qpu_client/test_auth.py @@ -0,0 +1,240 @@ +"""Testing lib/qpu_client/auth""" + +import json +import logging + +import pytest +from httpx2 import AsyncClient, HTTPStatusError +from pytest_httpx2 import HTTPXMock + +from warden.lib.config.config import QPUAuthConfig, QPUConfig +from warden.lib.qpu_client.auth import ( + KeycloakClientCredentialsAuth, + TokenRequestError, +) + +TOKEN_URL = "http://keycloak:8080/realms/pasqos/protocol/openid-connect/token" +QPU_URL = "http://qpu:4300/api/v1/system" + + +@pytest.fixture +def auth_conf() -> QPUAuthConfig: + return QPUAuthConfig( + url="http://keycloak:8080", realm="pasqos", id="warden", secret="s3cret" + ) + + +@pytest.mark.asyncio +async def test_token_is_fetched_once_and_reused(httpx_mock: HTTPXMock, auth_conf): + httpx_mock.add_response( + url=TOKEN_URL, + json={"access_token": "tok-1", "expires_in": 300}, + ) + httpx_mock.add_response(url=QPU_URL, json={"data": {}}) + httpx_mock.add_response(url=QPU_URL, json={"data": {}}) + + auth = KeycloakClientCredentialsAuth(auth_conf) + async with AsyncClient(auth=auth) as client: + first = await client.get(QPU_URL) + second = await client.get(QPU_URL) + + assert first.request.headers["Authorization"] == "Bearer tok-1" + assert second.request.headers["Authorization"] == "Bearer tok-1" + token_requests = [r for r in httpx_mock.get_requests() if str(r.url) == TOKEN_URL] + assert len(token_requests) == 1 + + +@pytest.mark.asyncio +async def test_token_request_uses_client_credentials_grant( + httpx_mock: HTTPXMock, auth_conf +): + httpx_mock.add_response( + url=TOKEN_URL, json={"access_token": "tok-1", "expires_in": 300} + ) + httpx_mock.add_response(url=QPU_URL, json={"data": {}}) + + auth = KeycloakClientCredentialsAuth(auth_conf) + async with AsyncClient(auth=auth) as client: + await client.get(QPU_URL) + + token_request = next( + r for r in httpx_mock.get_requests() if str(r.url) == TOKEN_URL + ) + body = token_request.read().decode() + assert "grant_type=client_credentials" in body + assert "client_id=warden" in body + assert "client_secret=s3cret" in body + + +@pytest.mark.asyncio +async def test_expired_token_is_refreshed( + httpx_mock: HTTPXMock, auth_conf, monkeypatch +): + # expires_in 300 with leeway 30 means the token is stale after 270s. + clock = {"now": 1_000.0} + monkeypatch.setattr("warden.lib.qpu_client.auth.monotonic", lambda: clock["now"]) + httpx_mock.add_response( + url=TOKEN_URL, json={"access_token": "tok-1", "expires_in": 300} + ) + httpx_mock.add_response(url=QPU_URL, json={"data": {}}) + httpx_mock.add_response( + url=TOKEN_URL, json={"access_token": "tok-2", "expires_in": 300} + ) + httpx_mock.add_response(url=QPU_URL, json={"data": {}}) + + auth = KeycloakClientCredentialsAuth(auth_conf) + async with AsyncClient(auth=auth) as client: + first = await client.get(QPU_URL) + clock["now"] += 271 + second = await client.get(QPU_URL) + + assert first.request.headers["Authorization"] == "Bearer tok-1" + assert second.request.headers["Authorization"] == "Bearer tok-2" + + +@pytest.mark.asyncio +async def test_401_triggers_one_refresh_and_one_retry(httpx_mock: HTTPXMock, auth_conf): + httpx_mock.add_response( + url=TOKEN_URL, json={"access_token": "stale", "expires_in": 300} + ) + httpx_mock.add_response(url=QPU_URL, status_code=401) + httpx_mock.add_response( + url=TOKEN_URL, json={"access_token": "fresh", "expires_in": 300} + ) + httpx_mock.add_response(url=QPU_URL, json={"data": {}}) + + auth = KeycloakClientCredentialsAuth(auth_conf) + async with AsyncClient(auth=auth) as client: + response = await client.get(QPU_URL) + + assert response.status_code == 200 + assert response.request.headers["Authorization"] == "Bearer fresh" + qpu_requests = [r for r in httpx_mock.get_requests() if str(r.url) == QPU_URL] + assert len(qpu_requests) == 2 + + +@pytest.mark.asyncio +async def test_persistent_401_is_not_retried_forever(httpx_mock: HTTPXMock, auth_conf): + httpx_mock.add_response( + url=TOKEN_URL, json={"access_token": "tok-1", "expires_in": 300} + ) + httpx_mock.add_response(url=QPU_URL, status_code=401) + httpx_mock.add_response( + url=TOKEN_URL, json={"access_token": "tok-2", "expires_in": 300} + ) + httpx_mock.add_response(url=QPU_URL, status_code=401) + + auth = KeycloakClientCredentialsAuth(auth_conf) + async with AsyncClient(auth=auth) as client: + response = await client.get(QPU_URL) + + # The second 401 is surfaced, not retried again. + assert response.status_code == 401 + qpu_requests = [r for r in httpx_mock.get_requests() if str(r.url) == QPU_URL] + assert len(qpu_requests) == 2 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("status_code", [400, 401]) +async def test_bad_credentials_raise_token_request_error( + httpx_mock: HTTPXMock, auth_conf, status_code +): + httpx_mock.add_response( + url=TOKEN_URL, + status_code=status_code, + json={"error": "invalid_client"}, + ) + + auth = KeycloakClientCredentialsAuth(auth_conf) + async with AsyncClient(auth=auth) as client: + with pytest.raises(TokenRequestError, match="invalid_client"): + await client.get(QPU_URL) + + +@pytest.mark.asyncio +async def test_keycloak_5xx_raises_retryable_http_status_error( + httpx_mock: HTTPXMock, auth_conf +): + # 503 must stay an httpx.HTTPStatusError so the existing retry decorator + # recognises it as transient. + httpx_mock.add_response(url=TOKEN_URL, status_code=503) + + auth = KeycloakClientCredentialsAuth(auth_conf) + async with AsyncClient(auth=auth) as client: + with pytest.raises(HTTPStatusError): + await client.get(QPU_URL) + + +@pytest.mark.asyncio +async def test_async_flow_attaches_token(httpx_mock: HTTPXMock, auth_conf): + httpx_mock.add_response( + url=TOKEN_URL, json={"access_token": "tok-async", "expires_in": 300} + ) + httpx_mock.add_response(url=QPU_URL, json={"data": {}}) + + auth = KeycloakClientCredentialsAuth(auth_conf) + async with AsyncClient(auth=auth) as client: + response = await client.get(QPU_URL) + + assert response.request.headers["Authorization"] == "Bearer tok-async" + + +@pytest.mark.asyncio +async def test_short_lived_token_still_caches_with_warning( + httpx_mock: HTTPXMock, auth_conf, caplog +): + # expires_in 30 with the default leeway_s 30 would clamp to 0 without the + # half-lifespan fallback, disabling the cache entirely and forcing a + # Keycloak round-trip on every request. + httpx_mock.add_response( + url=TOKEN_URL, json={"access_token": "tok-1", "expires_in": 30} + ) + httpx_mock.add_response(url=QPU_URL, json={"data": {}}) + httpx_mock.add_response(url=QPU_URL, json={"data": {}}) + + auth = KeycloakClientCredentialsAuth(auth_conf) + with caplog.at_level(logging.WARNING, logger="warden.lib.qpu_client.auth"): + async with AsyncClient(auth=auth) as client: + await client.get(QPU_URL) + await client.get(QPU_URL) + + token_requests = [r for r in httpx_mock.get_requests() if str(r.url) == TOKEN_URL] + assert len(token_requests) == 1 + assert any(record.levelno == logging.WARNING for record in caplog.records) + + +@pytest.mark.asyncio +async def test_client_sends_no_authorization_header_without_auth_config( + httpx_mock: HTTPXMock, +): + httpx_mock.add_response(url=QPU_URL, json={"data": {}}) + + response = await QPUConfig(uri="http://qpu:4300").client.get(QPU_URL) + + assert "Authorization" not in response.request.headers + + +@pytest.mark.asyncio +async def test_401_on_post_retries_with_fresh_token_and_identical_body( + httpx_mock: HTTPXMock, auth_conf +): + httpx_mock.add_response( + url=TOKEN_URL, json={"access_token": "stale", "expires_in": 300} + ) + httpx_mock.add_response(url=QPU_URL, status_code=401) + httpx_mock.add_response( + url=TOKEN_URL, json={"access_token": "fresh", "expires_in": 300} + ) + httpx_mock.add_response(url=QPU_URL, json={"data": {}}) + + body = {"circuit": "bell", "shots": 100} + auth = KeycloakClientCredentialsAuth(auth_conf) + async with AsyncClient(auth=auth) as client: + response = await client.post(QPU_URL, json=body) + + assert response.status_code == 200 + assert response.request.headers["Authorization"] == "Bearer fresh" + qpu_requests = [r for r in httpx_mock.get_requests() if str(r.url) == QPU_URL] + assert len(qpu_requests) == 2 + for request in qpu_requests: + assert json.loads(request.read()) == body diff --git a/tests/lib/qpu_client/test_retry.py b/tests/lib/qpu_client/test_retry.py new file mode 100644 index 0000000..2ebb342 --- /dev/null +++ b/tests/lib/qpu_client/test_retry.py @@ -0,0 +1,32 @@ +"""Testing lib/qpu_client/retry""" + +import pytest + +from warden.lib.qpu_client.auth import TokenRequestError +from warden.lib.qpu_client.retry import UnhandledError, retry + + +@pytest.mark.asyncio +async def test_already_classified_errors_are_not_rewrapped(): + calls = {"n": 0} + + @retry(max=5, sleep_s=0) + async def fails_with_bad_credentials(): + calls["n"] += 1 + raise TokenRequestError("invalid_client") + + with pytest.raises(TokenRequestError): + await fails_with_bad_credentials() + + # Fail fast: a wrong secret will not fix itself. + assert calls["n"] == 1 + + +@pytest.mark.asyncio +async def test_unknown_errors_are_still_wrapped(): + @retry(max=5, sleep_s=0) + async def fails_with_value_error(): + raise ValueError("something unexpected") + + with pytest.raises(UnhandledError): + await fails_with_value_error() diff --git a/tests/scheduler/test_scheduler.py b/tests/scheduler/test_scheduler.py index 2705ccf..8e4d358 100644 --- a/tests/scheduler/test_scheduler.py +++ b/tests/scheduler/test_scheduler.py @@ -18,7 +18,8 @@ from warden.lib.config import Config, SchedulerStrategy from warden.lib.models import Job from warden.scheduler.main import run_scheduler -from warden.scheduler.worker import LocalQPUWorker +from warden.scheduler.types import JobUpdateQueue +from warden.scheduler.worker import TERMINAL_STATUSES, LocalQPUWorker NOW = datetime.now() @@ -35,8 +36,6 @@ SYSTEM_API = API_URI + "/system" PROGRAM_API = API_URI + "/programs" -SUCCESS_CHECK_INTERVAL_S = 0.1 - DUMMY_RESULTS = json.dumps([{"counter": {"0001": 1, "0010": 2, "0100": 3, "1000": 4}}]) @@ -60,7 +59,7 @@ async def test_run_nominal( - To return "RUNNING" and then "DONE" status for each job - Run scheduler until: - All jobs have a "DONE" status is DB - - Test timeout after TEST_TIMEOUT_S + - Test timeout after utils.JOB_WAIT_TIMEOUT_S - Check n (jobs with status "DONE") = N_JOBS - Check "DONE" jobs have the right results and non-empty logs - Check those jobs have `scheduled_at` set @@ -73,7 +72,6 @@ async def test_run_nominal( # Enable warden logging for jobs 'logs' field to be populated caplog.set_level(logging.INFO, logger="warden") - TEST_TIMEOUT_S = 3 N_JOBS = 10 conf: Config = build_conf(strategy, QPU_URI) @@ -147,8 +145,6 @@ async def test_run_nominal( # Populate DB with jobs to run await utils.create_n_jobs(db_session_maker, N_JOBS) - stmt_count = select(func.count(Job.id)).where(Job.status == "DONE") - ################## ### TEST RUN ### ################## @@ -157,13 +153,7 @@ async def test_run_nominal( 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, - ) + await utils.wait_until_jobs_settled(session, main_task, count=N_JOBS) stmt_all = select(Job).where(Job.status == "DONE") jobs_done = (await session.execute(stmt_all)).scalars().all() @@ -200,7 +190,7 @@ async def test_run_resume_job( - 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 + - Test timeout after utils.JOB_WAIT_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 @@ -213,7 +203,6 @@ async def test_run_resume_job( # 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" @@ -330,8 +319,6 @@ async def test_run_resume_job( ], ) - stmt_count = select(func.count(Job.id)).where(Job.status == "DONE") - ################## ### TEST RUN ### ################## @@ -340,13 +327,7 @@ async def test_run_resume_job( 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, - ) + await utils.wait_until_jobs_settled(session, main_task, count=N_JOBS) stmt_all = select(Job).where(Job.status == "DONE") jobs_done = (await session.execute(stmt_all)).scalars().all() @@ -377,7 +358,7 @@ async def test_run_qpu_down( - No need to mock jobs calls - Run scheduler until: - All jobs have an "ERROR" status - - Test timeout after TEST_TIMEOUT_S + - Test timeout after utils.JOB_WAIT_TIMEOUT_S - Check n (jobs with status "ERROR") = N_JOBS - Check those jobs have non-empty logs """ @@ -389,7 +370,6 @@ async def test_run_qpu_down( # Enable warden logging for jobs 'logs' field to be populated caplog.set_level(logging.INFO, logger="warden") - TEST_TIMEOUT_S = 3 N_JOBS = 3 EXPECTED_STATUS = "ERROR" @@ -413,8 +393,6 @@ async def test_run_qpu_down( # Populate DB with jobs to run await utils.create_n_jobs(db_session_maker, N_JOBS) - stmt = select(func.count(Job.id)).where(Job.status == EXPECTED_STATUS) - ################## ### TEST RUN ### ################## @@ -423,13 +401,9 @@ async def test_run_qpu_down( 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, - N_JOBS, - interval=SUCCESS_CHECK_INTERVAL_S, - ) + await utils.wait_until_jobs_settled( + session, main_task, count=N_JOBS, statuses=(EXPECTED_STATUS,) + ) stmt_all = select(Job).where(Job.status == EXPECTED_STATUS) all_jobs = (await session.execute(stmt_all)).scalars().all() @@ -497,7 +471,6 @@ async def test_run_job_timeout( # Enable warden logging for jobs 'logs' field to be populated caplog.set_level(logging.INFO, logger="warden") - TEST_TIMEOUT_S = 3 N_JOBS = 8 N_JOBS_TIMEOUT = 4 @@ -650,13 +623,9 @@ async def test_run_job_timeout( 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_processed, - N_JOBS, - interval=SUCCESS_CHECK_INTERVAL_S, - ) + await utils.wait_until_jobs_settled( + session, main_task, count=N_JOBS, statuses=("DONE", "CANCELED") + ) n_processed = (await session.execute(stmt_processed)).scalar() assert n_processed == N_JOBS @@ -693,7 +662,7 @@ async def test_run_resume_job_timeout( - Accept the job's cancelation request - Run scheduler until: - All jobs have a "DONE" status in DB - - Test timeout after TEST_TIMEOUT_S + - Test timeout after utils.JOB_WAIT_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 @@ -706,7 +675,6 @@ async def test_run_resume_job_timeout( # 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 @@ -776,8 +744,6 @@ async def test_run_resume_job_timeout( backend_ids=[BACKEND_ID], ) - stmt_count = select(func.count(Job.id)).where(Job.status == EXPECTED_JOB_STATUS) - ################## ### TEST RUN ### ################## @@ -786,13 +752,9 @@ async def test_run_resume_job_timeout( 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, - ) + await utils.wait_until_jobs_settled( + session, main_task, count=N_JOBS, statuses=(EXPECTED_JOB_STATUS,) + ) stmt_all = select(Job).where(Job.status == EXPECTED_JOB_STATUS) jobs_done = (await session.execute(stmt_all)).scalars().all() @@ -828,7 +790,7 @@ async def test_run_retry_transient_errors( - To return "RUNNING" and then "DONE" status for each job - Run scheduler until: - All jobs have a "DONE" status in DB - - Test timeout after TEST_TIMEOUT_S + - Test timeout after utils.JOB_WAIT_TIMEOUT_S - Checks: - n (jobs with status "DONE") = N_JOBS - jobs have non-empty logs @@ -841,7 +803,6 @@ async def test_run_retry_transient_errors( # Enable warden logging for jobs 'logs' field to be populated caplog.set_level(logging.INFO, logger="warden") - TEST_TIMEOUT_S = 3 N_JOBS = 1 conf: Config = build_conf(strategy, QPU_URI) @@ -934,7 +895,6 @@ def _add_transient_errors(httpx_mock: HTTPXMock, url: str, method: str): # Populate DB with jobs to run await utils.create_n_jobs(db_session_maker, N_JOBS) - stmt_count = select(func.count(Job.id)).where(Job.status == "DONE") stmt = select(Job).where(Job.status == "DONE") ################## @@ -945,13 +905,7 @@ def _add_transient_errors(httpx_mock: HTTPXMock, url: str, method: str): 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, - ) + await utils.wait_until_jobs_settled(session, main_task, count=N_JOBS) jobs_done = (await session.execute(stmt)).scalars().all() assert len(jobs_done) == N_JOBS @@ -979,7 +933,7 @@ async def test_run_qpu_api_unreachable( - To return QPU status as "Down" - Run scheduler until: - All jobs have a "ERROR" status is DB - - Test timeout after TEST_TIMEOUT_S + - Test timeout after utils.JOB_WAIT_TIMEOUT_S - Checks - n(jobs with status "ERROR") = N_JOBS - All jobs have non-empty logs and an "ERROR" message @@ -992,7 +946,6 @@ async def test_run_qpu_api_unreachable( # Enable warden logging for jobs 'logs' field to be populated caplog.set_level(logging.INFO, logger="warden") - TEST_TIMEOUT_S = 3 N_JOBS = 1 EXPECTED_STATUS = "ERROR" @@ -1008,7 +961,6 @@ async def test_run_qpu_api_unreachable( # Populate DB with jobs to run await utils.create_n_jobs(db_session_maker, N_JOBS) - stmt_count = select(func.count(Job.id)).where(Job.status == EXPECTED_STATUS) stmt = select(Job).where(Job.status == EXPECTED_STATUS) ################## @@ -1019,13 +971,9 @@ async def test_run_qpu_api_unreachable( 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, - ) + await utils.wait_until_jobs_settled( + session, main_task, count=N_JOBS, statuses=(EXPECTED_STATUS,) + ) error_jobs = (await session.execute(stmt)).scalars().all() assert len(error_jobs) == N_JOBS @@ -1055,7 +1003,7 @@ async def test_run_job_creation_client_error( - Return exceptions when attempting to create a job - Run scheduler until: - All jobs have a "ERROR" status is DB - - Test timeout after TEST_TIMEOUT_S + - Test timeout after utils.JOB_WAIT_TIMEOUT_S - Checks - n(jobs with status "ERROR") = N_JOBS - All jobs have non-empty logs and an "ERROR" message @@ -1068,7 +1016,6 @@ async def test_run_job_creation_client_error( # Enable warden logging for jobs 'logs' field to be populated caplog.set_level(logging.INFO, logger="warden") - TEST_TIMEOUT_S = 3 N_JOBS = 3 EXPECTED_STATUS = "ERROR" @@ -1097,7 +1044,6 @@ async def test_run_job_creation_client_error( # Populate DB with jobs to run await utils.create_n_jobs(db_session_maker, N_JOBS) - stmt_count = select(func.count(Job.id)).where(Job.status == EXPECTED_STATUS) stmt = select(Job).where(Job.status == EXPECTED_STATUS) ################## @@ -1108,13 +1054,9 @@ async def test_run_job_creation_client_error( 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, - ) + await utils.wait_until_jobs_settled( + session, main_task, count=N_JOBS, statuses=(EXPECTED_STATUS,) + ) jobs = (await session.execute(stmt)).scalars().all() assert len(jobs) == N_JOBS @@ -1149,7 +1091,7 @@ async def test_run_job_client_error_timeout( (it's the same backend request in QPU) - Run scheduler until: - All jobs have an "ERROR" status is DB - - Test timeout after TEST_TIMEOUT_S + - Test timeout after utils.JOB_WAIT_TIMEOUT_S - Check n (jobs with status "ERROR") = N_JOBS - Check those jobs have `ended_at` set, even though the job was never reported as ended by the QPU: `to_error` must backfill it. @@ -1162,7 +1104,6 @@ async def test_run_job_client_error_timeout( # Enable warden logging for jobs 'logs' field to be populated caplog.set_level(logging.INFO, logger="warden") - TEST_TIMEOUT_S = 3 N_JOBS = 1 EXPECTED_STATUS = "ERROR" @@ -1215,7 +1156,6 @@ async def test_run_job_client_error_timeout( # Populate DB with jobs to run await utils.create_n_jobs(db_session_maker, N_JOBS) - stmt_count = select(func.count(Job.id)).where(Job.status == EXPECTED_STATUS) stmt = select(Job).where(Job.status == EXPECTED_STATUS) ################## @@ -1226,19 +1166,16 @@ async def test_run_job_client_error_timeout( 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, - ) + await utils.wait_until_jobs_settled( + session, main_task, count=N_JOBS, statuses=(EXPECTED_STATUS,) + ) jobs = (await session.execute(stmt)).scalars().all() assert len(jobs) == N_JOBS for job in jobs: assert len(job.logs) > 0 assert "ERROR" in job.logs + assert "Job execution ended with status 'ERROR'" in job.logs assert job.ended_at is not None @@ -1264,7 +1201,7 @@ async def test_run_job_canceled_by_cancellation_worker( - For JOB_ID_CANCELED return "CANCELED" status - Run scheduler until: - All jobs have a "DONE" or "CANCELED" status in DB - - Test timeout after TEST_TIMEOUT_S + - Test timeout after utils.JOB_WAIT_TIMEOUT_S - Check n (jobs with status "DONE") = N_JOBS-1 - Check "DONE" jobs have the right results and non-empty logs - Check n (jobs with status "CANCELED") = 1 @@ -1278,7 +1215,6 @@ async def test_run_job_canceled_by_cancellation_worker( # Enable warden logging for jobs 'logs' field to be populated caplog.set_level(logging.INFO, logger="warden") - TEST_TIMEOUT_S = 3 N_JOBS = 10 JOB_ID_CANCELED = 5 @@ -1384,8 +1320,6 @@ async def test_run_job_canceled_by_cancellation_worker( # Populate DB with jobs to run await utils.create_n_jobs(db_session_maker, N_JOBS) - stmt_count = select(func.count(Job.id)).where(Job.status.in_(("DONE", "CANCELED"))) - ################## ### TEST RUN ### ################## @@ -1394,13 +1328,9 @@ async def test_run_job_canceled_by_cancellation_worker( 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, - ) + await utils.wait_until_jobs_settled( + session, main_task, count=N_JOBS, statuses=("DONE", "CANCELED") + ) stmt_done = select(Job).where(Job.status == "DONE") jobs_done = (await session.execute(stmt_done)).scalars().all() @@ -1555,3 +1485,74 @@ async def crash(*args, **kwargs): job = (await session.execute(select(Job))).scalar_one() assert job.backend_id == "0" assert job.status == "RUNNING" + + +@pytest.mark.asyncio +@pytest.mark.parametrize("final_status", ["DONE", "ERROR"]) +async def test_terminal_status_and_closing_log_are_committed_together( + final_status: str, + httpx_mock: HTTPXMock, + caplog, +): + """A job's terminal status must be queued together with its closing log line. + + One JobUpdate is one transaction, so pushing the status and the closing + "Job execution ended with status '...'" line separately leaves a window + where the DB holds a finished job whose logs are truncated. Anything that + stops polling once the status is terminal - the API, and every test that + waits on a status then asserts on logs - reads incomplete logs. + + This guards the invariant: keep the closing line in the same update as the + terminal status, whether the job reaches it through `update_job` (DONE) + or through `to_error` (ERROR, here triggered by a failing job creation). + """ + + # Enable warden logging for jobs 'logs' field to be populated + caplog.set_level(logging.INFO, logger="warden") + + def job_json(status: str) -> dict: + return { + "data": { + "uid": 0, + "batch_id": SLURM_USER_ID, + "status": status, + "result": DUMMY_RESULTS if status == "DONE" else None, + "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", + url=SYSTEM_OPERATIONAL_API, + json={"data": {"operational_status": "UP"}}, + ) + if final_status == "ERROR": + # Job creation fails outright -> create_job() catches + # QPUClientRequestError and calls to_error() directly. + # is_reusable=True: the QPU client retries on 500s before giving up. + httpx_mock.add_response( + method="POST", status_code=500, url=JOB_API, is_reusable=True + ) + else: + httpx_mock.add_response( + method="POST", status_code=200, url=JOB_API, json=job_json("RUNNING") + ) + httpx_mock.add_response( + method="GET", status_code=200, url=JOB_API + "/0", json=job_json("RUNNING") + ) + httpx_mock.add_response( + method="GET", status_code=200, url=JOB_API + "/0", json=job_json("DONE") + ) + + queue: JobUpdateQueue = JobUpdateQueue() + worker = LocalQPUWorker(conf=build_conf(SchedulerStrategy.FIFO, QPU_URI)) + await worker.execute_job(queue=queue, nb_run=100, sequence="{}") + + updates = [queue.get_nowait() for _ in range(queue.qsize())] + first_terminal = next(u for u in updates if u.status in TERMINAL_STATUSES) + assert ( + f"Job execution ended with status '{final_status}'" in first_terminal.new_logs + ) diff --git a/tests/scheduler/test_scheduler_integration.py b/tests/scheduler/test_scheduler_integration.py index 778b77e..52780c8 100644 --- a/tests/scheduler/test_scheduler_integration.py +++ b/tests/scheduler/test_scheduler_integration.py @@ -80,12 +80,10 @@ async def test_run_scheduler_integration( await utils.create_n_jobs(db_session_maker, N_JOBS) - # The terminal status is committed before the closing "Job execution ended - # with status 'DONE'" log line is flushed, so waiting on the status alone - # races with the log assertions below. Wait for the logs too. - stmt = select(func.count(Job.id)).where( - Job.status == "DONE", Job.logs.contains("DONE") - ) + # Safe to wait on the status alone: the scheduler commits a job's terminal + # status and its complete logs in the same transaction, so the log + # assertions below cannot race it + stmt = select(func.count(Job.id)).where(Job.status == "DONE") ################## ### TEST RUN ### @@ -191,11 +189,8 @@ async def test_run_scheduler_integration_cancellation_worker( JOB_TO_CANCEL_ID = job_to_cancel.id - # Same race as above: wait for the closing log line, not just the status - stmt = select(func.count(Job.id)).where( - Job.status.in_(("CANCELED", "DONE")), - Job.logs.contains("Job execution ended"), - ) + # Status alone is enough here too, see the comment in the test above + stmt = select(func.count(Job.id)).where(Job.status.in_(("CANCELED", "DONE"))) ################## ### TEST RUN ### diff --git a/tests/scheduler/utils.py b/tests/scheduler/utils.py index 255503d..a925c91 100644 --- a/tests/scheduler/utils.py +++ b/tests/scheduler/utils.py @@ -2,15 +2,22 @@ import asyncio from asyncio import Task, timeout +from collections.abc import Sequence from contextlib import asynccontextmanager from typing import Any import pytest +from sqlalchemy import func, select from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker from warden.lib.config import Config, QPUConfig, SchedulerConfig, SchedulerStrategy from warden.lib.models import Job, Session +# Deliberately generous: this budget is only ever spent by a test that is +# already failing, so it costs nothing on the happy path. Tight per-test budgets +# turned a slow CI runner into a flake instead of catching anything. +JOB_WAIT_TIMEOUT_S = 30 + async def wait_until_scalar_equals( session: AsyncSession, @@ -123,6 +130,31 @@ async def scheduler_task_timeout(delay: float, scheduler_task: Task): pass +async def wait_until_jobs_settled( + session: AsyncSession, + scheduler_task: Task, + *, + count: int, + statuses: Sequence[str] = ("DONE",), + timeout_s: float = JOB_WAIT_TIMEOUT_S, + interval: float = 0.1, +) -> None: + """Wait until ``count`` jobs reached one of ``statuses``, then stop the scheduler. + + Use this rather than hand-rolling a wait predicate. The scheduler commits a + job's terminal status and its complete logs in a single transaction (see + ``JobExecutionTracker.update_job``), so waiting on the status alone is enough + to make the assertions that follow - including any on ``logs`` - safe. + + Fails the test on timeout, and always cancels ``scheduler_task`` so it cannot + outlive the test body and interfere with fixture teardown. + """ + + stmt = select(func.count(Job.id)).where(Job.status.in_(tuple(statuses))) + async with scheduler_task_timeout(timeout_s, scheduler_task): + await wait_until_scalar_equals(session, stmt, count, interval=interval) + + def build_conf(strategy: SchedulerStrategy, qpu_uri: str) -> Config: return Config( scheduler=SchedulerConfig( diff --git a/tests/test_config.py b/tests/test_config.py index 6c111da..8275382 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -8,6 +8,8 @@ from warden.lib.config.config import ( APIConfig, Config, + QPUAuthConfig, + QPUConfig, SchedulerConfig, SchedulerStrategy, SqliteConfig, @@ -150,3 +152,77 @@ def test_admin_users_must_not_be_empty(): """ with pytest.raises(ValidationError): APIConfig(admin_users=[]) + + +def test_qpu_auth_absent_by_default(): + assert Config().qpu.auth is None + + +def test_qpu_auth_token_url_is_built_from_base_and_realm(): + auth = QPUAuthConfig( + url="http://keycloak:8080", realm="pasqos", id="warden", secret="s" + ) + + assert ( + auth.token_url + == "http://keycloak:8080/realms/pasqos/protocol/openid-connect/token" + ) + + +def test_qpu_auth_token_url_tolerates_trailing_slash(): + auth = QPUAuthConfig( + url="http://keycloak:8080/", realm="pasqos", id="warden", secret="s" + ) + + assert ( + auth.token_url + == "http://keycloak:8080/realms/pasqos/protocol/openid-connect/token" + ) + + +def test_qpu_auth_rejects_partial_configuration(): + # A half-configured auth section must fail loudly rather than silently + # falling back to unauthenticated requests. + with pytest.raises(ValidationError): + # model_validate, not the constructor: omitting a required field is the + # point of the test, and a static type checker rejects the direct call. + QPUAuthConfig.model_validate({"url": "http://keycloak:8080", "id": "warden"}) + + +def test_qpu_auth_secret_read_from_env(monkeypatch, tmp_path): + monkeypatch.chdir(tmp_path) + monkeypatch.setenv("WARDEN_QPU_AUTH_URL", "http://keycloak:8080") + monkeypatch.setenv("WARDEN_QPU_AUTH_ID", "warden") + monkeypatch.setenv("WARDEN_QPU_AUTH_SECRET", "from-env") + + config = Config() + + assert config.qpu.auth is not None + assert config.qpu.auth.id == "warden" + assert config.qpu.auth.secret == "from-env" + + +def test_auth_flow_is_none_without_auth_config(): + assert Config().qpu.auth_flow is None + + +def test_auth_flow_is_memoized(): + qpu = QPUConfig( + uri="http://qpu:4300", + auth=QPUAuthConfig(url="http://keycloak:8080", id="warden", secret="s"), + ) + + assert qpu.auth_flow is qpu.auth_flow + + +def test_client_is_given_the_auth_flow(): + qpu = QPUConfig( + uri="http://qpu:4300", + auth=QPUAuthConfig(url="http://keycloak:8080", id="warden", secret="s"), + ) + + assert qpu.client.auth is qpu.auth_flow + + +def test_client_has_no_auth_without_auth_config(): + assert QPUConfig(uri="http://qpu:4300").client.auth is None diff --git a/tests/test_generate_config.py b/tests/test_generate_config.py index 161db8c..3f6c387 100644 --- a/tests/test_generate_config.py +++ b/tests/test_generate_config.py @@ -8,6 +8,7 @@ Config, MariadbConfig, PostgresConfig, + QPUAuthConfig, QPUConfig, SchedulerConfig, SqliteConfig, @@ -25,12 +26,17 @@ APIConfig, SchedulerConfig, QPUConfig, + QPUAuthConfig, SqliteConfig, PostgresConfig, MariadbConfig, ) for name in model.model_fields } +# Fields whose type is itself a nested model (e.g. QPUConfig.auth, an Optional +# QPUAuthConfig) are rendered as an uncommented "name:" header, with their own +# fields commented out beneath, rather than a single "# name:" scalar line. +NESTED_MODEL_FIELD_NAMES = {"auth"} def test_generate_config_is_valid_yaml(): @@ -47,7 +53,8 @@ def test_generate_config_documents_every_field(): generated = generate_config() for name in ALL_FIELD_NAMES: - assert f"# {name}:" in generated + prefix = "" if name in NESTED_MODEL_FIELD_NAMES else "# " + assert f"{prefix}{name}:" in generated def test_generate_config_preserves_existing_overrides(): @@ -101,9 +108,40 @@ class Outer(BaseModel): overridden = _render_section_fields(Outer.model_fields, {"inner": {"value": 42}}, 1) assert " value: 42" in overridden - assert yaml.safe_load(f"outer:\n{overridden}") == { - "outer": {"inner": {"value": 42}} - } + + +def test_render_fields_indents_optional_nested_models(): + """Test that an Optional (``Model | None``) nested field is recursed into + just like a plain nested model, falling back to the nested model's own + docstring since it has no field-level description of its own""" + + class Inner(BaseModel): + """Inner docstring.""" + + value: int = Field(default=1, description="An inner value.") + + class Outer(BaseModel): + inner: Inner | None = None + + generated = _render_section_fields(Outer.model_fields, {}, 1) + + assert ( + " # Inner docstring.\n inner:\n # An inner value.\n # value: 1" + in generated + ) + + +def test_generate_config_documents_qpu_auth_section(): + """Test that qpu.auth (an Optional QPUAuthConfig) is recursed into: its + own docstring and every one of its fields must appear, not just an opaque + "# auth: null" line""" + generated = generate_config() + + assert "auth:" in generated + assert "# auth: null" not in generated + assert "Keycloak client_credentials configuration" in generated + for name in QPUAuthConfig.model_fields: + assert f"# {name}:" in generated def test_generate_writes_directly_when_no_previous_file(tmp_path, monkeypatch): diff --git a/warden/lib/config/config.py b/warden/lib/config/config.py index fa8ce2f..c89a26e 100644 --- a/warden/lib/config/config.py +++ b/warden/lib/config/config.py @@ -181,6 +181,39 @@ class SchedulerConfig(WardenSettings): ) +class QPUAuthConfig(WardenSettings): + """Keycloak client_credentials configuration for outbound QPU API calls. + + Presence of this section is what enables authentication. There is + deliberately no separate `enabled` flag: a second switch can drift out of + sync with the credentials it guards. `url`, `id` and `secret` have no + defaults, so a partially configured section is a startup validation error + rather than a silent fallback to unauthenticated requests. + """ + + url: str = Field(description="Keycloak base URL, for example http://keycloak:8080") + + realm: str = Field(default="pasqos") + + id: str = Field(description="OIDC client_id") + + secret: str = Field( + description="OIDC client_secret. Provide via WARDEN_QPU_AUTH_SECRET, never in YAML." + ) + + leeway_s: float = Field( + default=30, + description="Refresh this many seconds before the token actually expires.", + ) + + @property + def token_url(self) -> str: + """Keycloak's OIDC token endpoint for this realm.""" + return ( + f"{self.url.rstrip('/')}/realms/{self.realm}/protocol/openid-connect/token" + ) + + class QPUConfig(WardenSettings): """QPU backend connection configuration.""" @@ -188,6 +221,8 @@ class QPUConfig(WardenSettings): default="http://localhost:8000", description="Local Pasqal QPU API URI." ) + auth: QPUAuthConfig | None = None + retry_max: int = Field( default=10, description=( @@ -211,6 +246,7 @@ class QPUConfig(WardenSettings): ) _client: httpx2.AsyncClient | None = PrivateAttr(default=None) + _auth_flow: httpx2.Auth | None = PrivateAttr(default=None) @property def verify(self) -> bool | ssl.SSLContext: @@ -219,10 +255,26 @@ def verify(self) -> bool | ssl.SSLContext: return ssl.create_default_context(cafile=self.tls_verify) return self.tls_verify + @property + def auth_flow(self) -> httpx2.Auth | None: + """Memoized Keycloak auth flow, or None when auth is not configured. + + Memoized so that every client built from this config shares a single + cached token. Imported lazily because ``qpu_client.auth`` imports this + module. + """ + if self.auth is None: + return None + if self._auth_flow is None: + from warden.lib.qpu_client.auth import KeycloakClientCredentialsAuth + + self._auth_flow = KeycloakClientCredentialsAuth(self.auth) + return self._auth_flow + @property def client(self) -> httpx2.AsyncClient: if self._client is None: - self._client = httpx2.AsyncClient(verify=self.verify) + self._client = httpx2.AsyncClient(verify=self.verify, auth=self.auth_flow) self._client.base_url = self.uri + API_PREFIX return self._client diff --git a/warden/lib/config/generate_config.py b/warden/lib/config/generate_config.py index 5631a4a..959e750 100644 --- a/warden/lib/config/generate_config.py +++ b/warden/lib/config/generate_config.py @@ -5,6 +5,8 @@ import re import sys import textwrap +import types +import typing from enum import Enum from pathlib import Path @@ -85,10 +87,14 @@ def _wrap_indented_text(text: str, indent: str) -> list[str]: def _get_nested_model(field: FieldInfo) -> type[BaseModel] | None: - """Returns if the field type is indeed a BaseModel subclass""" - field_annotation = field.annotation - if isinstance(field_annotation, type) and issubclass(field_annotation, BaseModel): - return field_annotation + """Returns the field's BaseModel type, unwrapping an Optional (``Model | None``).""" + annotation = field.annotation + if typing.get_origin(annotation) in (typing.Union, types.UnionType): + non_none = [a for a in typing.get_args(annotation) if a is not type(None)] + if len(non_none) == 1: + annotation = non_none[0] + if isinstance(annotation, type) and issubclass(annotation, BaseModel): + return annotation return None @@ -126,10 +132,19 @@ def _render_field( indent = INDENT_UNIT * depth + # Check if contains a nested model, to recursively render it below and, + # absent a field-level description, fall back to its docstring's summary + # line (the rest of a multi-paragraph docstring is skipped). + nested_model = _get_nested_model(field_info) + nested_doc = nested_model.__doc__ if nested_model else None + description = field_info.description or ( + nested_doc.strip().splitlines()[0] if nested_doc else None + ) + # Add comment lines lines = [ wrapped - for paragraph in (field_info.description or "").split("\n") + for paragraph in (description or "").split("\n") if paragraph for sentence in SENTENCE_RE.split(paragraph) if sentence @@ -139,8 +154,6 @@ def _render_field( # Get matching previously set data that we need to migrate to new config data_to_migrate = _existing_value(previous_section_data, name, field_info) - # Check if contains a nested model and recursively renders it - nested_model = _get_nested_model(field_info) if nested_model is not None: nested_existing = data_to_migrate if isinstance(data_to_migrate, dict) else {} lines.append(f"{indent}{name}:") diff --git a/warden/lib/qpu_client/auth.py b/warden/lib/qpu_client/auth.py new file mode 100644 index 0000000..fad4323 --- /dev/null +++ b/warden/lib/qpu_client/auth.py @@ -0,0 +1,114 @@ +"""Keycloak client_credentials authentication for outbound QPU API calls.""" + +import logging +from time import monotonic +from typing import Generator + +import httpx2 + +from warden.lib.config.config import QPUAuthConfig +from warden.lib.qpu_client.retry import QPUClientRequestError + +logger = logging.getLogger(__name__) + +# Token-endpoint statuses that will never succeed on retry: the credentials or +# the grant itself are wrong. Anything else (transport errors, 5xx) is left to +# propagate so the existing retry decorator can treat it as transient. +FATAL_TOKEN_STATUSES = (400, 401, 403) + + +class TokenRequestError(QPUClientRequestError): + """Keycloak refused to issue a token and retrying cannot help.""" + + +class KeycloakClientCredentialsAuth(httpx2.Auth): + """Attach a Keycloak service-account bearer token to each request. + + Implemented as an ``httpx.Auth`` so it runs inside the transport, below + Warden's ``retry`` decorator. That matters because 401 is not in + ``RETRY_HTTP_EXIT_CODES``: a token expiring mid-job would otherwise surface + as an immediate, non-retryable ``NotRetriedHTTPStatus``. Here it is just a + refresh. + + ``requires_response_body`` tells httpx to read the token response before + handing it back, since ``_store`` needs the JSON body. + + Args: + conf: Keycloak credentials and endpoint. + """ + + requires_response_body = True + + def __init__(self, conf: QPUAuthConfig) -> None: + self.conf = conf + self._token: str | None = None + # monotonic() deadline after which the cached token is considered stale. + self._expires_at: float = 0.0 + + # Note: unlocked check-then-fetch. Two concurrent requests can both miss + # and both fetch a token; one wins and the loser wasted a request. Add a + # lock only if token-endpoint traffic ever becomes a problem. + def auth_flow( + self, request: httpx2.Request + ) -> Generator[httpx2.Request, httpx2.Response, None]: + if not self._is_fresh(): + token_response = yield self._token_request() + self._store(token_response) + assert self._token is not None + request.headers["Authorization"] = f"Bearer {self._token}" + response = yield request + if response.status_code == httpx2.codes.UNAUTHORIZED: + logger.info("QPU API returned 401, refreshing token and retrying once") + token_response = yield self._token_request() + self._store(token_response) + request.headers["Authorization"] = f"Bearer {self._token}" + yield request + + def _is_fresh(self) -> bool: + return self._token is not None and monotonic() < self._expires_at + + def _token_request(self) -> httpx2.Request: + return httpx2.Request( + "POST", + self.conf.token_url, + data={ + "grant_type": "client_credentials", + "client_id": self.conf.id, + "client_secret": self.conf.secret, + }, + ) + + def _store(self, response: httpx2.Response) -> None: + """Validate a token response and cache the token.""" + if response.status_code in FATAL_TOKEN_STATUSES: + # Never log the response body of a token request: it may echo + # credentials. The error field alone is the useful part. + try: + error = response.json().get("error", "unknown_error") + except (ValueError, AttributeError): + error = "unknown_error" + raise TokenRequestError( + f"Keycloak refused to issue a token for client " + f"'{self.conf.id}' at {self.conf.token_url}: " + f"{response.status_code} {error}" + ) + # Transport errors and 5xx stay as httpx exceptions so the existing + # retry decorator sees them as transient. + response.raise_for_status() + + payload = response.json() + self._token = payload["access_token"] + expires_in = float(payload.get("expires_in", 0)) + configured_ttl = expires_in - self.conf.leeway_s + ttl = max(configured_ttl, expires_in / 2) + if ttl > configured_ttl: + logger.warning( + f"Configured leeway {self.conf.leeway_s}s leaves less than " + f"half of the {expires_in}s token lifespan; caching for " + f"{ttl}s (half the lifespan) instead." + ) + self._expires_at = monotonic() + ttl + logger.debug( + f"Obtained QPU API token for client '{self.conf.id}', " + f"expires in {expires_in}s" + ) diff --git a/warden/lib/qpu_client/retry.py b/warden/lib/qpu_client/retry.py index d14648b..f3bd945 100644 --- a/warden/lib/qpu_client/retry.py +++ b/warden/lib/qpu_client/retry.py @@ -50,6 +50,8 @@ def retry(max: int, sleep_s: float, no_retry: bool = False) -> Callable: UnhandledError: If decorator encounters an unnexpected exception. NotRetriedHTTPStatus: If the HTTP request returns with a non-retryable error code. MaxRetryError: If the maximum number of retries without success has been reached. + QPUClientRequestError: If `no_retry=True` or any subclass already classified as + non-retryable by the wrapped function (e.g. TokenRequestError) propagates unchanged. """ def decorator(func: Callable): @@ -60,6 +62,10 @@ def _handle_exception(e: Exception): elif isinstance(e, HTTPStatusError): if e.response.status_code not in RETRY_HTTP_EXIT_CODES: raise NotRetriedHTTPStatus(e) from e + elif isinstance(e, QPUClientRequestError): + # Already classified as non-retryable by the raiser (e.g. bad + # Keycloak credentials). Do not rewrap it as UnhandledError. + raise else: raise UnhandledError(e) from e diff --git a/warden/scheduler/worker.py b/warden/scheduler/worker.py index d6ce6fa..89e41aa 100644 --- a/warden/scheduler/worker.py +++ b/warden/scheduler/worker.py @@ -19,6 +19,8 @@ logger = logging.getLogger(__name__) +TERMINAL_STATUSES: tuple[JobStatus, ...] = ("ERROR", "DONE", "CANCELED") + class JobExecutionTracker: """Handles current job status and sends updates to db""" @@ -44,10 +46,6 @@ def job(self) -> QPUJobInfo: def is_error(self) -> bool: return self.status == "ERROR" - @property - def is_in_terminal_state(self) -> bool: - return self.status in ("DONE", "CANCELED", "ERROR") - @property def created_datetime(self) -> UTCDatetime: return self.job.created_datetime @@ -55,23 +53,41 @@ def created_datetime(self) -> UTCDatetime: async def update_job( self, qpu_job_info: QPUJobInfo, enforce_end_datetime: bool = False ): + was_terminal = self._status in TERMINAL_STATUSES self._qpu_job_info = qpu_job_info self._status = qpu_job_info.status or "ERROR" - if enforce_end_datetime and self._qpu_job_info.end_datetime is None: - self._qpu_job_info.end_datetime = datetime.now(timezone.utc) - await self.push_update() + await self.push_update(was_terminal, enforce_end_datetime) async def to_error(self): + was_terminal = self._status in TERMINAL_STATUSES self._status = "ERROR" - if self._qpu_job_info and self._qpu_job_info.end_datetime is None: - self._qpu_job_info.end_datetime = datetime.now(timezone.utc) - await self.push_update() + await self.push_update(was_terminal, enforce_end_datetime=True) def log(self, msg: str) -> None: self._log_buffer.append(msg + "\n") - async def push_update(self): - """Push update of job execution to db commit task through queue""" + async def push_update( + self, was_terminal: bool | None = None, enforce_end_datetime: bool = False + ): + """Push update of job execution to db commit task through queue + + `was_terminal` and `enforce_end_datetime` are set by `update_job`/ + `to_error` on a status transition, so the closing log line and the + backfilled `end_datetime` land in the same `JobUpdate` - hence the + same transaction - as the terminal status itself. Flushed separately, + the DB would briefly hold a finished job whose logs/dates are + incomplete, and anything that stops polling once the status is + terminal would read that incomplete state. + """ + if ( + enforce_end_datetime + and self._qpu_job_info + and self._qpu_job_info.end_datetime is None + ): + self._qpu_job_info.end_datetime = datetime.now(timezone.utc) + if was_terminal is False and self._status in TERMINAL_STATUSES: + logger.info("Job execution ended with status '%s'", self._status) + new_logs = "".join(self._log_buffer) self._log_buffer = [] @@ -172,7 +188,6 @@ async def execute_job( return await self.await_job_execution(job_tracker) - logger.info("Job execution ended with status '%s'", job_tracker.status) # Flush potential last updates before return await job_tracker.push_update() @@ -247,7 +262,7 @@ async def await_job_execution(self, job_tracker: JobExecutionTracker) -> None: polling_start = job_tracker.created_datetime await self._get_job_poll(job_tracker) - while not job_tracker.is_in_terminal_state: + while job_tracker.status not in TERMINAL_STATUSES: if self.is_timed_out(self.conf_sched.job_polling_timeout_s, polling_start): logger.warning( f"Job timed out (max {self.conf_sched.job_polling_timeout_s} s). " @@ -256,6 +271,10 @@ async def await_job_execution(self, job_tracker: JobExecutionTracker) -> None: ) try: qpu_job_info = await self.qpu_client.cancel_job(job_tracker.job.uid) + # Logged before the update so it is buffered into the same + # JobUpdate, and stays ahead of the closing line that a + # terminal status appends + logger.info("Job cancellation done") await job_tracker.update_job( qpu_job_info, enforce_end_datetime=True ) @@ -263,7 +282,6 @@ async def await_job_execution(self, job_tracker: JobExecutionTracker) -> None: logger.error(f"Failed cancelling job: {e}") await job_tracker.to_error() continue - logger.info("Job cancellation done") continue await asyncio.sleep(self.conf_sched.job_polling_interval_s) await self._get_job_poll(job_tracker)