schedulers/panda-idds/check_panda_jobs.py

0001 #!/usr/bin/env python3
0002 """Inspect AID2E PanDA/iDDS requests, transforms, tasks, and jobs.
0003
0004 Source ``setup_aid2e.sh`` and ``setup_panda.sh`` before using this script. The
0005 script performs read-only queries through ``panda-client`` and never reads or
0006 prints authentication tokens.
0007
0008 Examples::
0009
0010     python check_panda_jobs.py --task-id 3665
0011     python check_panda_jobs.py --request-id 3665 --transform-id 32469
0012     python check_panda_jobs.py --jedi-task-id 40758
0013     python check_panda_jobs.py --job-id 4878999 --job-id 4879000
0014 """
0015
0016 from __future__ import annotations
0017
0018 import argparse
0019 import json
0020 import sys
0021 from typing import Any, Dict, Iterable, List, Optional, Sequence
0022
0023 REQUEST_FIELDS = (
0024     "request_id",
0025     "name",
0026     "scope",
0027     "status",
0028     "substatus",
0029     "request_type",
0030     "username",
0031     "site",
0032     "created_at",
0033     "updated_at",
0034     "total_transforms",
0035     "finished_transforms",
0036     "subfinished_transforms",
0037     "failed_transforms",
0038     "errors",
0039 )
0040
0041 TRANSFORM_FIELDS = (
0042     "transform_id",
0043     "transform_name",
0044     "transform_type",
0045     "transform_tag",
0046     "status",
0047     "substatus",
0048     "workload_id",
0049     "internal_id",
0050     "parent_transform_id",
0051     "parent_internal_id",
0052     "created_at",
0053     "updated_at",
0054     "finished_at",
0055     "errors",
0056 )
0057
0058 TASK_FIELDS = (
0059     "jediTaskID",
0060     "status",
0061     "taskName",
0062     "parent_tid",
0063     "errordialog",
0064     "processingType",
0065     "taskType",
0066     "site",
0067     "nFiles",
0068     "nFilesFinished",
0069     "nFilesFailed",
0070 )
0071
0072 JOB_FIELDS = (
0073     "PandaID",
0074     "jediTaskID",
0075     "jobStatus",
0076     "jobSubStatus",
0077     "jobName",
0078     "attemptNr",
0079     "computingSite",
0080     "creationTime",
0081     "startTime",
0082     "endTime",
0083     "transExitCode",
0084     "pilotErrorCode",
0085     "pilotErrorDiag",
0086     "ddmErrorCode",
0087     "ddmErrorDiag",
0088     "exeErrorCode",
0089     "exeErrorDiag",
0090     "taskBufferErrorCode",
0091     "taskBufferErrorDiag",
0092 )
0093
0094
0095 class QueryError(RuntimeError):
0096     """A read-only PanDA/iDDS query failed."""
0097
0098
0099 def _display_value(value: Any) -> Any:
0100     """Convert PanDA/iDDS values into compact JSON-compatible values."""
0101     if value is None or isinstance(value, (str, int, float, bool)):
0102         return value
0103     if isinstance(value, dict):
0104         # iDDS values passed through panda-client can be encoded enum records.
0105         attributes = value.get("attributes")
0106         if isinstance(attributes, dict) and attributes.get("_name_"):
0107             return attributes["_name_"]
0108         return {str(key): _display_value(item) for key, item in value.items()}
0109     if isinstance(value, (list, tuple, set)):
0110         return [_display_value(item) for item in value]
0111     enum_name = getattr(value, "name", None)
0112     if enum_name:
0113         return str(enum_name)
0114     return str(value)
0115
0116
0117 def _select_mapping_fields(
0118     mapping: Dict[str, Any], fields: Iterable[str]
0119 ) -> Dict[str, Any]:
0120     return {
0121         field: _display_value(mapping[field])
0122         for field in fields
0123         if field in mapping and mapping[field] is not None
0124     }
0125
0126
0127 def _select_object_fields(obj: Any, fields: Iterable[str]) -> Dict[str, Any]:
0128     result: Dict[str, Any] = {}
0129     for field in fields:
0130         value = getattr(obj, field, None)
0131         if value is not None:
0132             result[field] = _display_value(value)
0133     return result
0134
0135
0136 def _unwrap_idds_response(response: Any, operation: str) -> Any:
0137     """Unwrap panda-client's ``(status, (success, data))`` response."""
0138     if not isinstance(response, (list, tuple)) or len(response) < 2:
0139         raise QueryError(f"{operation} returned an unexpected response")
0140     status, payload = response[0], response[1]
0141     if status != 0:
0142         raise QueryError(f"{operation} failed with panda-client status {status}")
0143     if not isinstance(payload, (list, tuple)) or len(payload) < 2:
0144         raise QueryError(f"{operation} returned an unexpected payload")
0145     success, data = payload[0], payload[1]
0146     if not success:
0147         raise QueryError(f"{operation} failed: {_display_value(data)}")
0148
0149     # Some manager calls, including get_transforms in released client versions,
0150     # add a second ``(success, data)`` envelope and leave the data JSON encoded.
0151     while (
0152         isinstance(data, (list, tuple)) and len(data) == 2 and isinstance(data[0], bool)
0153     ):
0154         nested_success, data = data
0155         if not nested_success:
0156             raise QueryError(f"{operation} failed: {_display_value(data)}")
0157     if isinstance(data, str) and data.lstrip().startswith(("[", "{")):
0158         try:
0159             data = json.loads(data)
0160         except json.JSONDecodeError as exc:
0161             raise QueryError(f"{operation} returned invalid JSON data") from exc
0162     return data
0163
0164
0165 def _get_idds_api(verbose: bool = False) -> Any:
0166     try:
0167         import idds.common.utils as idds_utils
0168         import pandaclient.idds_api as idds_api
0169     except ImportError as exc:
0170         raise QueryError(
0171             "PanDA/iDDS client packages are unavailable; source setup_aid2e.sh first"
0172         ) from exc
0173
0174     return idds_api.get_api(
0175         idds_utils.json_dumps,
0176         compress=True,
0177         manager=True,
0178         verbose=verbose,
0179     )
0180
0181
0182 def _query_request_rows(
0183     request_id: int, *, verbose: bool = False
0184 ) -> List[Dict[str, Any]]:
0185     """Return the detailed rows used to describe an iDDS workflow."""
0186     api = _get_idds_api(verbose=verbose)
0187     rows = (
0188         _unwrap_idds_response(
0189             api.get_requests(
0190                 request_id=request_id,
0191                 with_detail=True,
0192                 with_metadata=False,
0193             ),
0194             "get_requests",
0195         )
0196         or []
0197     )
0198     return [row for row in rows if isinstance(row, dict)]
0199
0200
0201 def query_workflow(request_id: int, *, verbose: bool = False) -> Dict[str, Any]:
0202     """Return workflow-level information for an iDDS request/task ID."""
0203     rows = _query_request_rows(request_id, verbose=verbose)
0204     if not rows:
0205         raise QueryError(f"iDDS workflow/task {request_id} was not found")
0206     workflow = _select_mapping_fields(rows[0], REQUEST_FIELDS)
0207     workflow.setdefault("request_id", int(request_id))
0208     workflow["task_id"] = int(request_id)
0209     return workflow
0210
0211
0212 def query_transforms(
0213     request_id: int,
0214     transform_id: Optional[int] = None,
0215     *,
0216     verbose: bool = False,
0217 ) -> List[Dict[str, Any]]:
0218     """Return compact transform records for an iDDS request."""
0219     api = _get_idds_api(verbose=verbose)
0220     if transform_id is not None:
0221         raw = _unwrap_idds_response(
0222             api.get_transform(request_id=request_id, transform_id=transform_id),
0223             "get_transform",
0224         )
0225         raw_transforms = [raw] if raw else []
0226     else:
0227         # The bulk get_transforms call is incompatible with some released
0228         # iDDS manager/client combinations. Detailed request rows reliably
0229         # expose transform IDs, which can then be queried individually.
0230         request_rows = _query_request_rows(request_id, verbose=verbose)
0231         transform_ids = sorted(
0232             {
0233                 int(row["transform_id"])
0234                 for row in request_rows
0235                 if isinstance(row, dict) and row.get("transform_id") is not None
0236             }
0237         )
0238         raw_transforms = [
0239             _unwrap_idds_response(
0240                 api.get_transform(request_id=request_id, transform_id=item),
0241                 f"get_transform({item})",
0242             )
0243             for item in transform_ids
0244         ]
0245
0246     return [
0247         _select_mapping_fields(transform, TRANSFORM_FIELDS)
0248         for transform in raw_transforms
0249         if isinstance(transform, dict)
0250     ]
0251
0252
0253 def query_task(task_id: int, *, verbose: bool = False) -> Dict[str, Any]:
0254     """Return task details and all discoverable PanDA IDs."""
0255     try:
0256         from pandaclient import Client
0257     except ImportError as exc:
0258         raise QueryError(
0259             "panda-client is unavailable; source setup_aid2e.sh first"
0260         ) from exc
0261
0262     status, task = Client.getJediTaskDetails(
0263         {"jediTaskID": int(task_id)},
0264         False,
0265         True,
0266         verbose=verbose,
0267     )
0268     if status != 0 or not isinstance(task, dict):
0269         raise QueryError(f"getJediTaskDetails({task_id}) failed with status {status}")
0270
0271     job_status, panda_ids = Client.getPandaIDsWithTaskID(int(task_id), verbose=verbose)
0272     if job_status != 0:
0273         raise QueryError(
0274             f"getPandaIDsWithTaskID({task_id}) failed with status {job_status}"
0275         )
0276
0277     result = _select_mapping_fields(task, TASK_FIELDS)
0278     result.setdefault("jediTaskID", int(task_id))
0279     result["panda_ids"] = [int(job_id) for job_id in (panda_ids or [])]
0280     return result
0281
0282
0283 def query_jobs(
0284     job_ids: Sequence[int], *, verbose: bool = False
0285 ) -> List[Dict[str, Any]]:
0286     """Return compact job records, including archived jobs when available."""
0287     if not job_ids:
0288         return []
0289     try:
0290         from pandaclient import Client
0291     except ImportError as exc:
0292         raise QueryError(
0293             "panda-client is unavailable; source setup_aid2e.sh first"
0294         ) from exc
0295
0296     status, jobs = Client.getFullJobStatus(
0297         [int(job_id) for job_id in job_ids], verbose=verbose
0298     )
0299     if status != 0 or jobs is None:
0300         raise QueryError(f"getFullJobStatus failed with status {status}")
0301     return [_select_object_fields(job, JOB_FIELDS) for job in jobs if job is not None]
0302
0303
0304 def inspect_identifiers(
0305     *,
0306     request_id: Optional[int],
0307     transform_id: Optional[int],
0308     task_ids: Sequence[int],
0309     job_ids: Sequence[int],
0310     max_jobs: int,
0311     verbose: bool,
0312 ) -> Dict[str, Any]:
0313     """Resolve identifiers down the request → transform → task → job chain."""
0314     report: Dict[str, Any] = {}
0315     discovered_task_ids = {int(task_id) for task_id in task_ids}
0316     transforms: List[Dict[str, Any]] = []
0317
0318     if request_id is not None:
0319         transforms = query_transforms(request_id, transform_id, verbose=verbose)
0320         discovered_task_ids.update(
0321             int(transform["workload_id"])
0322             for transform in transforms
0323             if transform.get("workload_id") not in (None, "")
0324         )
0325
0326     tasks = [
0327         query_task(task_id, verbose=verbose) for task_id in sorted(discovered_task_ids)
0328     ]
0329     discovered_job_ids = {int(job_id) for job_id in job_ids}
0330     for task in tasks:
0331         discovered_job_ids.update(task.get("panda_ids", []))
0332
0333     selected_job_ids = sorted(discovered_job_ids)
0334     if len(selected_job_ids) > max_jobs:
0335         report["job_limit"] = {
0336             "discovered": len(selected_job_ids),
0337             "queried": max_jobs,
0338         }
0339         selected_job_ids = selected_job_ids[:max_jobs]
0340     if selected_job_ids:
0341         jobs = query_jobs(selected_job_ids, verbose=verbose)
0342     else:
0343         jobs = []
0344
0345     tasks_by_id = {int(task["jediTaskID"]): task for task in tasks}
0346     jobs_by_task: Dict[int, List[Dict[str, Any]]] = {}
0347     for job in jobs:
0348         jedi_task_id = job.get("jediTaskID")
0349         if jedi_task_id is not None:
0350             jobs_by_task.setdefault(int(jedi_task_id), []).append(job)
0351
0352     if request_id is not None:
0353         workflow = query_workflow(request_id, verbose=verbose)
0354         workflow_transforms: List[Dict[str, Any]] = []
0355         for transform in transforms:
0356             transform_report = dict(transform)
0357             workload_id = transform.get("workload_id")
0358             if workload_id not in (None, ""):
0359                 task_id = int(workload_id)
0360                 task = tasks_by_id.get(task_id)
0361                 if task is not None:
0362                     transform_report["jedi_task"] = task
0363                 transform_report["jobs"] = jobs_by_task.get(task_id, [])
0364             else:
0365                 transform_report["jobs"] = []
0366             workflow_transforms.append(transform_report)
0367
0368         status_counts: Dict[str, int] = {}
0369         for job in jobs:
0370             status = str(job.get("jobStatus", "unknown"))
0371             status_counts[status] = status_counts.get(status, 0) + 1
0372         workflow["summary"] = {
0373             "transforms": len(transforms),
0374             "jedi_tasks": len(tasks),
0375             "panda_jobs": len(jobs),
0376             "job_statuses": status_counts,
0377         }
0378         workflow["transforms"] = workflow_transforms
0379         report["workflow"] = workflow
0380     else:
0381         if tasks:
0382             report["jedi_tasks"] = [
0383                 {
0384                     "task": task,
0385                     "jobs": jobs_by_task.get(int(task["jediTaskID"]), []),
0386                 }
0387                 for task in tasks
0388             ]
0389         if job_ids:
0390             report["jobs"] = jobs
0391
0392     return report
0393
0394
0395 def build_parser() -> argparse.ArgumentParser:
0396     parser = argparse.ArgumentParser(
0397         description=(
0398             "Inspect PanDA/iDDS execution by request ID, transform ID, JEDI task ID, "
0399             "or PanDA job ID. Queries are read-only."
0400         )
0401     )
0402     parser.add_argument(
0403         "--request-id",
0404         "--task-id",
0405         "--runtaskid",
0406         dest="request_id",
0407         type=int,
0408         help="AID2E/iDDS workflow task ID (iDDS request ID)",
0409     )
0410     parser.add_argument(
0411         "--transform-id",
0412         type=int,
0413         help="iDDS transform ID; requires --request-id",
0414     )
0415     parser.add_argument(
0416         "--jedi-task-id",
0417         dest="task_ids",
0418         action="append",
0419         type=int,
0420         default=[],
0421         help="individual JEDI task ID; may be supplied more than once",
0422     )
0423     parser.add_argument(
0424         "--job-id",
0425         "--panda-id",
0426         dest="job_ids",
0427         action="append",
0428         type=int,
0429         default=[],
0430         help="PanDA job ID; may be supplied more than once",
0431     )
0432     parser.add_argument(
0433         "--max-jobs",
0434         type=int,
0435         default=100,
0436         help="maximum number of discovered jobs to query (default: 100)",
0437     )
0438     parser.add_argument(
0439         "--verbose",
0440         action="store_true",
0441         help="enable panda-client request diagnostics",
0442     )
0443     return parser
0444
0445
0446 def main(argv: Optional[Sequence[str]] = None) -> int:
0447     parser = build_parser()
0448     args = parser.parse_args(argv)
0449
0450     if args.transform_id is not None and args.request_id is None:
0451         parser.error("--transform-id requires --request-id")
0452     if args.max_jobs < 1:
0453         parser.error("--max-jobs must be at least 1")
0454     if args.request_id is None and not args.task_ids and not args.job_ids:
0455         parser.error("provide --task-id/--request-id, --jedi-task-id, or --job-id")
0456
0457     try:
0458         report = inspect_identifiers(
0459             request_id=args.request_id,
0460             transform_id=args.transform_id,
0461             task_ids=args.task_ids,
0462             job_ids=args.job_ids,
0463             max_jobs=args.max_jobs,
0464             verbose=args.verbose,
0465         )
0466     except QueryError as exc:
0467         print(f"error: {exc}", file=sys.stderr)
0468         return 1
0469
0470     print(json.dumps(report, indent=2, sort_keys=True, default=str))
0471     return 0
0472
0473
0474 if __name__ == "__main__":
0475     raise SystemExit(main())