diff --git a/Makefile b/Makefile index 5b8b528..ebd9812 100644 --- a/Makefile +++ b/Makefile @@ -49,7 +49,7 @@ dev-s3df: deps # --- Docker / GHCR targets --- GHCR_USERNAME ?= slaclab GHCR_IMAGE ?= ghcr.io/$(GHCR_USERNAME)/iri-s3df -IMAGE_TAG ?= dev +IMAGE_TAG ?= prod-07132026 docker-build: docker build --platform linux/amd64 -t $(GHCR_IMAGE):$(IMAGE_TAG) . diff --git a/app/s3df/compute_adapter.py b/app/s3df/compute_adapter.py index d40d30d..1746e10 100644 --- a/app/s3df/compute_adapter.py +++ b/app/s3df/compute_adapter.py @@ -24,6 +24,7 @@ import jwt # PyJWT from slurmrestd_client.api_client import ApiClient from slurmrestd_client.api.slurm_api import SlurmApi +from slurmrestd_client.api.slurmdb_api import SlurmdbApi from slurmrestd_client.configuration import Configuration from slurmrestd_client.exceptions import ApiException from slurmrestd_client.models.slurm_v0041_post_job_submit_request import ( @@ -134,14 +135,23 @@ def _mint_slurm_jwt(unix_username: str) -> str: # slurmrestd_client helpers # --------------------------------------------------------------------------- -def _build_slurm_api(token: str) -> SlurmApi: +def _build_api_client(token: str) -> ApiClient: url = os.environ.get("SLURM_REST_URL") if not url: raise RuntimeError("SLURM_REST_URL environment variable is not set") cfg = Configuration(host=url) cfg.verify_ssl = False #os.environ.get("SLURM_VERIFY_SSL", "true").lower() == "true" # api_key on the config is not used for header auth — we pass headers manually - return SlurmApi(ApiClient(cfg)) + return ApiClient(cfg) + + +def _build_slurm_api(token: str) -> SlurmApi: + return SlurmApi(_build_api_client(token)) + + +def _build_slurmdb_api(token: str) -> SlurmdbApi: + """Client for the Slurm accounting DB (slurmdbd) endpoints.""" + return SlurmdbApi(_build_api_client(token)) def _auth_headers(unix_username: str, token: str) -> dict: @@ -209,6 +219,59 @@ def _job_from_slurm_info(job_info, include_spec: bool = False) -> dict: } return job_dict +def _job_from_slurmdb_info(job_record, include_spec: bool = False) -> dict: + """ + Convert a slurmdbd accounting record (sacct-backed) to a plain dict that + IRI can deserialise into its Job model. + + The accounting record has a different shape than the live-scheduler job + object handled by `_job_from_slurm_info`: + - the final state lives under `state.current` (a list of strings) + - the numeric id is `job_id`, the wall-clock limit is under `time.limit` + The returned dict shape is kept identical to `_job_from_slurm_info` so the + router's `models.Job` deserialisation is unchanged. + """ + state_obj = getattr(job_record, "state", None) + current = getattr(state_obj, "current", None) if state_obj is not None else None + state_str = _map_state(current) + + status: dict = {"state": state_str} + exit_code = getattr(job_record, "exit_code", None) + return_code = getattr(exit_code, "return_code", None) if exit_code is not None else None + if return_code is not None and getattr(return_code, "set", False): + status["exit_code"] = getattr(return_code, "number", None) + + job_dict = { + "id": str(job_record.job_id), + "status": status, + } + + if include_spec: + time_obj = getattr(job_record, "time", None) + tl = getattr(time_obj, "limit", None) if time_obj is not None else None + if tl is not None and getattr(tl, "set", False): + duration_secs = int(getattr(tl, "number", 0) or 0) * 60 + else: + duration_secs = 0 + + # NB: the model field is `job_spec` (compute_models.Job). The live-job + # helper `_job_from_slurm_info` uses the key "spec", which does not match + # and is silently dropped by IRIBaseModel's serializer — a pre-existing + # bug in that helper's include_spec path. + job_dict["job_spec"] = { + "name": getattr(job_record, "name", None), + "resources": { + "node_count": getattr(job_record, "allocation_nodes", None), + }, + "attributes": { + "queue_name": getattr(job_record, "partition", None), + "account": getattr(job_record, "account", None), + "duration": duration_secs, + }, + } + return job_dict + + def _build_batch_script(job_spec: compute_models.JobSpec) -> str: """ Build a Slurm batch-script body from an IRI JobSpec. @@ -287,6 +350,18 @@ def _get_slurm_context(self, user): api = _build_slurm_api(token) headers = _auth_headers(unix_user, token) return api, headers + + def _get_slurmdb_context(self, user): + """Return (slurmdb_api, headers, unix_user) for the authenticated user. + + The per-user JWT (`sun` claim) scopes accounting results to this user via + Slurm accounting permissions, so no separate username filter is required. + """ + unix_user = getattr(user, "unix_username", user.id) + token = _mint_slurm_jwt(unix_user) + api = _build_slurmdb_api(token) + headers = _auth_headers(unix_user, token) + return api, headers, unix_user # -- submit_job --------------------------------------------------------- @@ -500,7 +575,31 @@ async def get_job( raise HTTPException(status_code=500, detail="Slurm get_job failed") from exc if historical: - raise HTTPException(status_code=501, detail="Historical job lookup is not implemented yet") + # The job is no longer live — fall back to slurmdbd (accounting DB) + # for its final state. Query by job_id alone + dbapi, db_headers, unix_user = self._get_slurmdb_context(user) + try: + db_resp = dbapi.slurmdb_v0041_get_job(str(job_id), _headers=db_headers) + except ApiException as exc: + if exc.status == 404: + raise HTTPException(status_code=404, detail=f"Job {job_id} not found") from exc + logger.exception("slurmdbd get_job failed for job %s", job_id) + raise HTTPException(status_code=500, detail="Slurm accounting lookup failed") from exc + + records = (db_resp.jobs if db_resp else None) or [] + # A job_id can map to multiple accounting records (job arrays / + # duplicate submissions); the last is the most recent. + record = records[-1] if records else None + if record is None: + raise HTTPException(status_code=404, detail=f"Job {job_id} not found") + + # Defense in depth: the JWT already scopes results to this user, but + # never surface another user's job even if that ever changes. + record_user = getattr(record, "user", None) + if record_user and record_user != unix_user: + raise HTTPException(status_code=404, detail=f"Job {job_id} not found") + + return _job_from_slurmdb_info(record, include_spec) raise HTTPException(status_code=404, detail=f"Job {job_id} not found") @@ -520,6 +619,10 @@ async def get_jobs( api, headers = self._get_slurm_context(user) if historical: + # Historical *listing* would query slurmdbd's /jobs endpoint + # (e.g. by users=), which scans a user's entire accounting history + # and can be expensive. Only single-job historical lookup (get_job) + # is supported for now; revisit listing once load is characterised. raise HTTPException(status_code=501, detail="Historical job listing is not implemented yet") try: diff --git a/app/s3df/tests/test_compute_adapter.py b/app/s3df/tests/test_compute_adapter.py new file mode 100644 index 0000000..3d581c1 --- /dev/null +++ b/app/s3df/tests/test_compute_adapter.py @@ -0,0 +1,160 @@ +"""Tests for the S3DF compute adapter's historical (slurmdbd) job lookup.""" + +from types import SimpleNamespace +from unittest.mock import MagicMock + +import pytest +from fastapi import HTTPException + +from app.routers.compute.models import JobState +from app.s3df.compute_adapter import SLACComputeAdapter +from slurmrestd_client.exceptions import ApiException + + +TEST_USER = SimpleNamespace(id="amithm", unix_username="amithm") + + +def _db_record(*, job_id=12345, state="COMPLETED", user="amithm", return_code=0): + """Build a minimal slurmdbd accounting record (duck-typed).""" + return SimpleNamespace( + job_id=job_id, + name="my-job", + state=SimpleNamespace(current=[state]) if state is not None else None, + user=user, + partition="milano", + account="proj123", + allocation_nodes=2, + exit_code=SimpleNamespace( + return_code=SimpleNamespace(set=True, number=return_code) + ), + time=SimpleNamespace(limit=SimpleNamespace(set=True, number=60)), + ) + + +def _make_adapter(monkeypatch, *, live_side_effect, db_records): + """Wire an adapter whose live lookup and slurmdbd lookup are mocked.""" + adapter = SLACComputeAdapter() + + live_api = MagicMock() + live_api.slurm_v0041_get_job = MagicMock(side_effect=live_side_effect) + monkeypatch.setattr(adapter, "_get_slurm_context", lambda user: (live_api, {})) + + db_api = MagicMock() + if isinstance(db_records, Exception): + db_api.slurmdb_v0041_get_job = MagicMock(side_effect=db_records) + else: + db_api.slurmdb_v0041_get_job = MagicMock( + return_value=SimpleNamespace(jobs=db_records) + ) + monkeypatch.setattr( + adapter, "_get_slurmdb_context", lambda user: (db_api, {}, user.unix_username) + ) + return adapter, live_api, db_api + + +def _not_live(*_a, **_k): + """Simulate the live scheduler no longer knowing about the job.""" + raise ApiException(status=404) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "slurm_state,expected", + [ + ("COMPLETED", JobState.COMPLETED), + ("FAILED", JobState.FAILED), + ("CANCELLED", JobState.CANCELED), + ("TIMEOUT", JobState.FAILED), + ], +) +async def test_historical_returns_final_state(monkeypatch, slurm_state, expected): + adapter, _, db_api = _make_adapter( + monkeypatch, + live_side_effect=_not_live, + db_records=[_db_record(state=slurm_state)], + ) + job = await adapter.get_job(resource=None, user=TEST_USER, job_id="12345", historical=True) + assert job["id"] == "12345" + assert job["status"]["state"] == expected + # Queried by job_id alone — the lightest slurmdbd query. + db_api.slurmdb_v0041_get_job.assert_called_once() + assert db_api.slurmdb_v0041_get_job.call_args.args[0] == "12345" + + +@pytest.mark.asyncio +async def test_historical_picks_most_recent_record(monkeypatch): + adapter, _, _ = _make_adapter( + monkeypatch, + live_side_effect=_not_live, + db_records=[_db_record(state="FAILED"), _db_record(state="COMPLETED")], + ) + job = await adapter.get_job(resource=None, user=TEST_USER, job_id="12345", historical=True) + assert job["status"]["state"] == JobState.COMPLETED + + +@pytest.mark.asyncio +async def test_historical_include_spec(monkeypatch): + adapter, _, _ = _make_adapter( + monkeypatch, + live_side_effect=_not_live, + db_records=[_db_record()], + ) + job = await adapter.get_job( + resource=None, user=TEST_USER, job_id="12345", historical=True, include_spec=True + ) + assert job["job_spec"]["attributes"]["queue_name"] == "milano" + assert job["job_spec"]["attributes"]["duration"] == 3600 # 60 min -> secs + assert job["job_spec"]["resources"]["node_count"] == 2 + + +@pytest.mark.asyncio +async def test_historical_empty_records_is_404(monkeypatch): + adapter, _, _ = _make_adapter( + monkeypatch, live_side_effect=_not_live, db_records=[] + ) + with pytest.raises(HTTPException) as exc: + await adapter.get_job(resource=None, user=TEST_USER, job_id="12345", historical=True) + assert exc.value.status_code == 404 + + +@pytest.mark.asyncio +async def test_historical_other_users_job_is_404(monkeypatch): + adapter, _, _ = _make_adapter( + monkeypatch, + live_side_effect=_not_live, + db_records=[_db_record(user="someone_else")], + ) + with pytest.raises(HTTPException) as exc: + await adapter.get_job(resource=None, user=TEST_USER, job_id="12345", historical=True) + assert exc.value.status_code == 404 + + +@pytest.mark.asyncio +async def test_historical_slurmdbd_404_is_404(monkeypatch): + adapter, _, _ = _make_adapter( + monkeypatch, live_side_effect=_not_live, db_records=ApiException(status=404) + ) + with pytest.raises(HTTPException) as exc: + await adapter.get_job(resource=None, user=TEST_USER, job_id="12345", historical=True) + assert exc.value.status_code == 404 + + +@pytest.mark.asyncio +async def test_historical_slurmdbd_error_is_500(monkeypatch): + adapter, _, _ = _make_adapter( + monkeypatch, live_side_effect=_not_live, db_records=ApiException(status=503) + ) + with pytest.raises(HTTPException) as exc: + await adapter.get_job(resource=None, user=TEST_USER, job_id="12345", historical=True) + assert exc.value.status_code == 500 + + +@pytest.mark.asyncio +async def test_non_historical_missing_job_does_not_hit_slurmdbd(monkeypatch): + adapter, _, db_api = _make_adapter( + monkeypatch, live_side_effect=_not_live, db_records=[_db_record()] + ) + with pytest.raises(HTTPException) as exc: + await adapter.get_job(resource=None, user=TEST_USER, job_id="12345", historical=False) + assert exc.value.status_code == 404 + db_api.slurmdb_v0041_get_job.assert_not_called()