feat: Support diracx pilots for jobs

Author: Robin-Van-de-Merghel

diracx-core/src/diracx/core/exceptions.py

@@ -139,6 +139,10 @@ class PilotJobsNotFoundError(DiracFormattedError):
     pattern = "Pilots or Jobs %s not found"
 
 
+class PilotCantAccessJobError(DiracFormattedError):
+    pattern = "Pilot %s can't access jobs"
+
+
 class PilotAlreadyAssociatedWithJobError(DiracFormattedError):
     pattern = "Pilot is already associated with a job %s "
 

diracx-logic/src/diracx/logic/pilots/query.py

@@ -112,6 +112,19 @@ async def get_pilot_jobs_ids_by_pilot_id(
     return [job["JobID"] for job in jobs]
 
 
+async def get_pilot_jobs_ids_by_stamp(
+    pilot_db: PilotAgentsDB, pilot_stamp: str
+) -> list[int]:
+    pilot_ids = await get_pilot_ids_by_stamps(
+        pilot_db=pilot_db,
+        pilot_stamps=[pilot_stamp],
+    )
+
+    return await get_pilot_jobs_ids_by_pilot_id(
+        pilot_db=pilot_db, pilot_id=pilot_ids[0]
+    )
+
+
 async def get_secrets_by_hashed_secrets_bulk(
     pilot_db: PilotAgentsDB, hashed_secrets: list[bytes], parameters: list[str] = []
 ) -> list[dict[Any, Any]]:

diracx-routers/src/diracx/routers/jobs/access_policies.py

@@ -49,7 +49,8 @@ async def policy(
 
         if action == ActionType.PILOT:
             # TODO: For now we map this to MANAGE but it should be changed once
-            # we have pilot credentials
+            # Pilots will eventually be removed from this access policy
+            # See #572
             action = ActionType.MANAGE
 
         if action == ActionType.CREATE:

diracx-routers/src/diracx/routers/jobs/status.py

@@ -74,6 +74,8 @@ async def set_job_statuses(
     check_permissions: CheckWMSPolicyCallable,
     force: bool = False,
 ) -> SetJobStatusReturn:
+    # FIXME: Pilots will eventually be removed from this endpoint
+    # See #572
     await check_permissions(
         action=ActionType.MANAGE, job_db=job_db, job_ids=list(job_update)
     )
@@ -121,6 +123,8 @@ async def add_heartbeat(
 
     The `data` parameter and return value are mappings keyed by job ID.
     """
+    # FIXME: Pilots will eventually be removed from this endpoint
+    # See #572
     await check_permissions(action=ActionType.PILOT, job_db=job_db, job_ids=list(data))
 
     await add_heartbeat_bl(
@@ -171,6 +175,8 @@ async def patch_metadata(
     job_parameters_db: JobParametersDB,
     check_permissions: CheckWMSPolicyCallable,
 ):
+    # FIXME: Pilots will eventually be removed from this endpoint
+    # See #572
     await check_permissions(action=ActionType.MANAGE, job_db=job_db, job_ids=updates)
     try:
         await set_job_parameters_or_attributes_bl(updates, job_db, job_parameters_db)

diracx-routers/src/diracx/routers/pilots/resources.py

@@ -0,0 +1,159 @@
+from __future__ import annotations
+
+from datetime import datetime
+from http import HTTPStatus
+from typing import Annotated, Any
+
+from fastapi import Depends, HTTPException, status
+
+from diracx.core.exceptions import JobNotFoundError, PilotCantAccessJobError
+from diracx.core.models import (
+    HeartbeatData,
+    JobCommand,
+    JobStatusUpdate,
+    SetJobStatusReturn,
+)
+from diracx.logic.jobs.status import add_heartbeat as add_heartbeat_bl
+from diracx.logic.jobs.status import get_job_commands as get_job_commands_bl
+from diracx.logic.jobs.status import (
+    set_job_parameters_or_attributes as set_job_parameters_or_attributes_bl,
+)
+from diracx.logic.jobs.status import set_job_statuses as set_job_statuses_bl
+from diracx.routers.utils.pilots import (
+    AuthorizedPilotInfo,
+    verify_dirac_pilot_access_token,
+    verify_that_pilot_can_access_jobs,
+)
+
+from ..dependencies import (
+    Config,
+    JobDB,
+    JobLoggingDB,
+    JobParametersDB,
+    PilotAgentsDB,
+    TaskQueueDB,
+)
+from ..fastapi_classes import DiracxRouter
+
+router = DiracxRouter()
+
+
+@router.patch("/status")
+async def pilot_set_job_statuses(
+    job_update: dict[int, dict[datetime, JobStatusUpdate]],
+    config: Config,
+    job_db: JobDB,
+    job_logging_db: JobLoggingDB,
+    task_queue_db: TaskQueueDB,
+    job_parameters_db: JobParametersDB,
+    pilot_agents_db: PilotAgentsDB,
+    pilot_info: Annotated[
+        AuthorizedPilotInfo, Depends(verify_dirac_pilot_access_token)
+    ],
+    force: bool = False,
+) -> SetJobStatusReturn:
+    # Endpoint only for DiracX pilots (with a pilot token)
+    try:
+        await verify_that_pilot_can_access_jobs(
+            pilot_db=pilot_agents_db,
+            pilot_stamp=pilot_info.pilot_stamp,
+            job_ids=list(job_update),
+        )
+    except (PilotCantAccessJobError, JobNotFoundError) as e:
+        raise HTTPException(
+            status_code=status.HTTP_403_FORBIDDEN, detail="Pilot can't access this job."
+        ) from e
+
+    try:
+        result = await set_job_statuses_bl(
+            status_changes=job_update,
+            config=config,
+            job_db=job_db,
+            job_logging_db=job_logging_db,
+            task_queue_db=task_queue_db,
+            job_parameters_db=job_parameters_db,
+            force=force,
+        )
+    except ValueError as e:
+        raise HTTPException(
+            status_code=HTTPStatus.BAD_REQUEST,
+            detail=str(e),
+        ) from e
+
+    if not result.success:
+        raise HTTPException(
+            status_code=HTTPStatus.NOT_FOUND,
+            detail=result.model_dump(),
+        )
+
+    return result
+
+
+@router.patch("/heartbeat")
+async def pilot_add_heartbeat(
+    data: dict[int, HeartbeatData],
+    config: Config,
+    job_db: JobDB,
+    job_logging_db: JobLoggingDB,
+    task_queue_db: TaskQueueDB,
+    job_parameters_db: JobParametersDB,
+    pilot_agents_db: PilotAgentsDB,
+    pilot_info: Annotated[
+        AuthorizedPilotInfo, Depends(verify_dirac_pilot_access_token)
+    ],
+) -> list[JobCommand]:
+    """Register a heartbeat from the job.
+
+    This endpoint is used by the JobAgent to send heartbeats to the WMS and to
+    receive job commands from the WMS. It also results in stalled jobs being
+    restored to the RUNNING status.
+
+    The `data` parameter and return value are mappings keyed by job ID.
+    """
+    # Endpoint only for DiracX pilots (with a pilot token)
+    try:
+        await verify_that_pilot_can_access_jobs(
+            pilot_db=pilot_agents_db,
+            pilot_stamp=pilot_info.pilot_stamp,
+            job_ids=list(data),
+        )
+    except (PilotCantAccessJobError, JobNotFoundError) as e:
+        raise HTTPException(
+            status_code=status.HTTP_403_FORBIDDEN, detail="Pilot can't access this job."
+        ) from e
+
+    await add_heartbeat_bl(
+        data, config, job_db, job_logging_db, task_queue_db, job_parameters_db
+    )
+    return await get_job_commands_bl(data, job_db)
+
+
+@router.patch("/metadata", status_code=HTTPStatus.NO_CONTENT)
+async def pilot_patch_metadata(
+    updates: dict[int, dict[str, Any]],
+    job_db: JobDB,
+    job_parameters_db: JobParametersDB,
+    pilot_agents_db: PilotAgentsDB,
+    pilot_info: Annotated[
+        AuthorizedPilotInfo, Depends(verify_dirac_pilot_access_token)
+    ],
+):
+    # Endpoint only for DiracX pilots (with a pilot token)
+    try:
+        await verify_that_pilot_can_access_jobs(
+            pilot_db=pilot_agents_db,
+            pilot_stamp=pilot_info.pilot_stamp,
+            job_ids=list(updates),
+        )
+    except (PilotCantAccessJobError, JobNotFoundError) as e:
+        raise HTTPException(
+            status_code=status.HTTP_403_FORBIDDEN, detail="Pilot can't access this job."
+        ) from e
+
+    try:
+        await set_job_parameters_or_attributes_bl(updates, job_db, job_parameters_db)
+    except ValueError as e:
+        raise HTTPException(
+            status_code=HTTPStatus.BAD_REQUEST,
+            detail=str(e),
+        ) from e

diracx-routers/src/diracx/routers/utils/pilots.py

@@ -8,9 +8,11 @@
 from joserfc.jwt import JWTClaimsRegistry
 from pydantic import BaseModel
 
+from diracx.core.exceptions import PilotCantAccessJobError
 from diracx.core.models import UUID, PilotInfo
 from diracx.logic.auth.utils import read_token
-from diracx.routers.dependencies import AuthSettings
+from diracx.logic.pilots.query import get_pilot_jobs_ids_by_stamp
+from diracx.routers.dependencies import AuthSettings, PilotAgentsDB
 
 
 class AuthInfo(BaseModel):
@@ -71,3 +73,23 @@ async def verify_dirac_pilot_access_token(
             status_code=status.HTTP_401_UNAUTHORIZED,
             detail="Invalid JWT",
         ) from e
+
+
+async def verify_that_pilot_can_access_jobs(
+    pilot_db: PilotAgentsDB, pilot_stamp: str, job_ids: list[int]
+):
+    # Get its jobs
+    pilot_jobs = await get_pilot_jobs_ids_by_stamp(
+        pilot_db=pilot_db, pilot_stamp=pilot_stamp
+    )
+
+    # Equivalent of issubset, but cleaner
+    if set(job_ids) <= set(pilot_jobs):
+        return
+
+    forbidden_jobs_ids = set(job_ids) - set(pilot_jobs)
+
+    if forbidden_jobs_ids:
+        return PilotCantAccessJobError(
+            data={"forbidden_jobs_ids": str(forbidden_jobs_ids)}
+        )