schedulers/PanDAiDDS/runner.py

0001
0002 """PanDA / iDDS-backed scheduler runner.
0003
0004 This runner integrates with the iDDS workflow API when available to submit
0005 function-based work to PanDA. For simpler command-based jobs it falls back
0006 to local execution (previously done with joblib) so the scheduler remains
0007 usable without iDDS installed.
0008
0009 The implementation favors clarity and conservative behavior:
0010 - If a job dict contains a "function" key, we try to submit it via iDDS.
0011 - Otherwise we execute the provided "command" locally and collect outputs.
0012
0013 This file intentionally keeps iDDS imports inside functions so the module
0014 can be imported even when iDDS is not installed.
0015 """
0016
0017 import json
0018 import logging
0019 import os
0020 import subprocess
0021 import datetime
0022 from typing import Dict, Any, List, Optional, Tuple
0023
0024 import threading
0025 import uuid
0026 from time import time
0027 import time as _time
0028
0029 from aid2e.schedulers.base import BaseScheduler, JobStatus, StageExecutionResult
0030 from aid2e.schedulers.PanDAiDDS.config import PanDAiDDSRunnerConfig
0031
0032
0033 class PanDAiDDSScheduler(BaseScheduler):
0034     """Scheduler that prefers PanDA/iDDS for function-style work.
0035
0036     Notes:
0037     - This is a pragmatic adapter. To enable full iDDS capabilities you must
0038       install the `idds` package and ensure any function objects submitted are
0039       importable in the remote environment used by PanDA.
0040     """
0041
0042     def __init__(self, config: Optional[PanDAiDDSRunnerConfig] = None) -> None:
0043         super().__init__(config)
0044         self.config = config or PanDAiDDSRunnerConfig()
0045         self.logger = logging.getLogger("PanDAiDDSScheduler")
0046
0047         # in-memory bookkeeping organized by stage_name
0048         self.running_jobs: Dict[str, Dict[str, Any]] = {}
0049         self.running_stages: Dict[str, Dict[str, Any]] = {}
0050         # running_funcs: stage_name -> job_id -> func_name -> {work, tf_id, status, results}
0051         self.running_funcs: Dict[str, Dict[str, Any]] = {}
0052         # jobs: stage_name -> job_id -> tf_id
0053         self.jobs: Dict[str, Dict[str, Any]] = {}
0054         self.num_checks: int = 0
0055         self.workflow = None
0056         self.workflow_id = None
0057         self.lock = threading.Lock()
0058         # Cache workflows per stage_name to ensure one workflow per stage
0059         self.stage_workflows: Dict[str, Any] = {}
0060
0061
0062     # --- Run stage (synchronous convenience wrapper) -----------------------
0063     def run_stage(
0064         self,
0065         stage_name: str,
0066         job_definitions: List[Dict[str, Any]],
0067         parallelism_policy: Optional[Dict[str, Any]] = None,
0068         working_dir: Optional[str] = None,
0069     ) -> StageExecutionResult:
0070         """Execute a stage: prefer iDDS submissions for function-jobs, fallback to local commands.
0071
0072         This method runs synchronously: it will submit all jobs first, then poll
0073         until all jobs complete, and finally return a StageExecutionResult.
0074         """
0075         policy = parallelism_policy or {}
0076         poll_interval = policy.get("poll_interval", 5)
0077
0078         self.logger.info("Running stage '%s' with %d jobs", stage_name, len(job_definitions))
0079
0080         job_statuses: List[JobStatus] = []
0081         all_artifacts: Dict[str, Any] = {}
0082         all_success = True
0083
0084         # Map job_id to job_def for easy lookup during polling
0085         job_id_to_job_def: Dict[str, Dict[str, Any]] = {}
0086
0087         # Create a simple stage id for bookkeeping
0088         stage_id = uuid.uuid4().hex
0089         self.running_stages[stage_id] = {"status": "running", "jobs": [], "result": None}
0090
0091         # Phase 1: Submit all jobs
0092         submitted_job_ids = []
0093         local_job_results = {}
0094
0095         for index, job_def in enumerate(job_definitions):
0096             job_name = job_def.get("name", f"job_{index}")
0097             job_id = f"{stage_name}_{job_name}_{index}"
0098             job_id_to_job_def[job_id] = job_def
0099
0100             # If the job provides a function object, try to submit to iDDS.
0101             if job_def.get("function") is not None:
0102                 try:
0103                     self.logger.info("Submitting function job %s to iDDS/PanDA", job_id)
0104                     # normalize job to the shape expected by submit_job
0105                     job_dict = job_def.copy()
0106                     job_dict.setdefault("job_id", job_id)
0107                     self.submit_job(stage_name, job_dict, working_dir)
0108                     submitted_job_ids.append(job_id)
0109                 except RuntimeError as exc:
0110                     self.logger.exception("Failed to submit function job %s: %s", job_id, exc)
0111                     # mark as failed immediately
0112                     local_job_results[job_id] = {
0113                         "status": "failed",
0114                         "return_code": -1,
0115                         "stdout": "",
0116                         "stderr": str(exc),
0117                         "outputs": {},
0118                     }
0119                     all_success = False
0120             else:
0121                 # Fallback: execute the command locally
0122                 self.logger.info("Executing local job %s", job_id)
0123                 result = self._execute_job(job_def, working_dir)
0124                 local_job_results[job_id] = result
0125
0126         # Phase 2: Poll all submitted iDDS jobs until they finish
0127         poll_count = 0
0128         while submitted_job_ids:
0129             _time.sleep(poll_interval)
0130             poll_count += 1
0131             # Only log polling message every 10th iteration to reduce noise
0132             if poll_count % 60 == 1:
0133                 self.logger.info("Polling %d remaining jobs (poll #%d)", len(submitted_job_ids), poll_count)
0134
0135             # check status of all remaining jobs
0136             for job_id in list(submitted_job_ids):
0137                 try:
0138                     # Find job_def for this job_id to extract job_context
0139                     job_def = job_id_to_job_def.get(job_id, None)
0140                     job_context = job_def.get("job_context") if job_def else None
0141                     self.check_single_job_status({"job_id": job_id, "stage_name": stage_name}, job_context)
0142                 except Exception as exc:
0143                     # if job still running, check_single_job_status may raise until finished
0144                     self.logger.debug("Job %s not finished yet: %s", job_id, exc)
0145
0146                 # if running_funcs no longer contains job_id, it's done or failed
0147                 stage_funcs = self.running_funcs.get(stage_name, {})
0148                 if job_id not in stage_funcs:
0149                     self.logger.info("Job %s finished", job_id)
0150                     submitted_job_ids.remove(job_id)
0151                     # store a placeholder result; detailed results are in self.jobs mapping
0152                     local_job_results[job_id] = {
0153                         "status": "completed",
0154                         "return_code": 0,
0155                         "stdout": "",
0156                         "stderr": "",
0157                         "outputs": {},
0158                     }
0159
0160         # Phase 3: Consolidate results
0161         for index, job_def in enumerate(job_definitions):
0162             job_name = job_def.get("name", f"job_{index}")
0163             job_id = f"{stage_name}_{job_name}_{index}"
0164             result = local_job_results.get(job_id, {})
0165
0166             return_code = result.get("return_code", -1)
0167             success = return_code == 0
0168             status = result.get("status", "completed" if success else "failed")
0169
0170             job_statuses.append(
0171                 JobStatus(
0172                     job_id=job_id,
0173                     status=status,
0174                     return_code=return_code,
0175                     stdout=result.get("stdout", ""),
0176                     stderr=result.get("stderr", ""),
0177                 outputs=result.get("outputs"),
0178                 )
0179             )
0180
0181             if result.get("outputs"):
0182                 all_artifacts.update(result["outputs"])
0183
0184             if not success:
0185                 all_success = False
0186
0187         self.running_stages[stage_id]["status"] = "completed" if all_success else "failed"
0188         result = StageExecutionResult(
0189             stage_name=stage_name,
0190             job_statuses=job_statuses,
0191             artifacts=all_artifacts,
0192             success=all_success,
0193             error_message=None if all_success else f"Some jobs failed in stage '{stage_name}'",
0194         )
0195
0196         return result
0197
0198     def get_stage_results(self, stage_id: str) -> StageExecutionResult:
0199         """Return stored StageExecutionResult for a stage.
0200
0201         Raises KeyError if unknown, RuntimeError if not finished.
0202         """
0203         if stage_id not in self.running_stages:
0204             raise KeyError(f"Unknown stage_id: {stage_id}")
0205         state = self.running_stages[stage_id]
0206         if state.get("status") in ("queued", "running"):
0207             raise RuntimeError(f"Stage {stage_id} is not yet finished (status={state.get('status')})")
0208         result = state.get("result")
0209         if result is None:
0210             raise RuntimeError(f"Stage {stage_id} completed but no result is available")
0211         return result
0212
0213     def check_status(self, job_id: str) -> JobStatus:
0214         """Check the status of a previously submitted job.
0215
0216         Args:
0217             job_id: Unique job identifier.
0218
0219         Returns:
0220             JobStatus with current state and metrics.
0221         """
0222         # Extract stage_name from job_id (format: stage_name_job_name_index)
0223         parts = job_id.split("_")
0224         stage_name = parts[0] if parts else None
0225
0226         if not stage_name or stage_name not in self.running_funcs:
0227             # Job not found in running funcs, check if it's completed
0228             return JobStatus(
0229                 job_id=job_id,
0230                 status="unknown",
0231                 return_code=None,
0232             )
0233
0234         stage_jobs = self.running_funcs.get(stage_name, {})
0235         if job_id not in stage_jobs:
0236             return JobStatus(
0237                 job_id=job_id,
0238                 status="completed",
0239                 return_code=0,
0240             )
0241
0242         # Job is still running, return running status
0243         return JobStatus(
0244             job_id=job_id,
0245             status="running",
0246             return_code=None,
0247         )
0248
0249     def cancel_job(self, job_id: str) -> bool:
0250         """Cancel a job if it is still running.
0251
0252         Args:
0253             job_id: Unique job identifier.
0254
0255         Returns:
0256             True if the job was cancelled, False otherwise.
0257         """
0258         # Extract stage_name from job_id
0259         parts = job_id.split("_")
0260         stage_name = parts[0] if parts else None
0261
0262         if not stage_name or stage_name not in self.running_funcs:
0263             self.logger.warning("Cannot cancel job %s: stage not found", job_id)
0264             return False
0265
0266         entry = self.running_funcs[stage_name].get(job_id, {})
0267         cancelled = False
0268         for func_name, g in entry.get("funcs", {}).items():
0269             work = g.get("work")
0270             try:
0271                 if work and not work.is_terminated():
0272                     work.cancel()
0273                     cancelled = True
0274             except Exception:
0275                 self.logger.exception("Failed to cancel work for job %s", job_id)
0276
0277         return cancelled
0278
0279     def submit_stage(
0280         self,
0281         stage_name: str,
0282         job_definitions: List[Dict[str, Any]],
0283         parallelism_policy: Optional[Dict[str, Any]] = None,
0284         working_dir: Optional[str] = None,
0285     ) -> str:
0286         """Submit a stage for asynchronous execution and return a stage_id.
0287
0288         This schedules the full `run_stage` call in a background thread and
0289         stores the `StageExecutionResult` in memory for later retrieval.
0290         """
0291         stage_id = uuid.uuid4().hex
0292         self.logger.info("Submitting stage '%s' as %s", stage_name, stage_id)
0293
0294         # state holder
0295         state: Dict[str, Any] = {
0296             "thread": None,
0297             "status": "queued",
0298             "result": None,
0299             "submitted_at": time(),
0300             "job_count": len(job_definitions),
0301         }
0302
0303         def _target() -> None:
0304             try:
0305                 state["status"] = "running"
0306                 result = self.run_stage(stage_name, job_definitions, parallelism_policy, working_dir)
0307                 state["result"] = result
0308                 state["status"] = "completed" if result.success else "failed"
0309             except Exception as exc:  # pragma: no cover - background safety
0310                 self.logger.exception("Asynchronous stage execution failed: %s", exc)
0311                 state["result"] = StageExecutionResult(
0312                     stage_name=stage_name, job_statuses=[], artifacts={}, success=False, error_message=str(exc)
0313                 )
0314                 state["status"] = "failed"
0315
0316         thread = threading.Thread(target=_target, name=f"PanDAiDDS-stage-{stage_id}", daemon=True)
0317         state["thread"] = thread
0318         self.running_stages[stage_id] = state
0319         thread.start()
0320
0321         return stage_id
0322
0323     def check_stage_status(self, stage_id: str):
0324         """Return a StageStatus summarizing progress for a submitted stage.
0325
0326         If the stage_id is unknown, raises KeyError.
0327         """
0328         from aid2e.schedulers.base import StageStatus
0329
0330         if stage_id not in self.running_stages:
0331             raise KeyError(f"Unknown stage_id: {stage_id}")
0332
0333         state = self.running_stages[stage_id]
0334         status = state.get("status", "unknown")
0335         result: Optional[StageExecutionResult] = state.get("result")
0336         total = state.get("job_count")
0337         completed_jobs = 0
0338         job_statuses = None
0339
0340         if result is not None:
0341             job_statuses = result.job_statuses
0342             completed_jobs = len(job_statuses)
0343
0344         progress = None
0345         if total and total > 0:
0346             progress = float(completed_jobs) / float(total) if total else None
0347
0348         return StageStatus(
0349             stage_id=stage_id,
0350             status=status,
0351             completed_jobs=completed_jobs,
0352             total_jobs=total,
0353             progress=progress,
0354             job_statuses=job_statuses,
0355         )
0356
0357     # --- IDDS / PanDA integration helpers (best-effort, optional) ---------
0358     def submit_idds_workflow(self, stage_name: str):
0359         """Define and submit an iDDS workflow for a stage. Returns the workflow object.
0360
0361         This method is idempotent per stage_name: if a workflow for the given
0362         stage has already been submitted, it returns the cached workflow.
0363         Otherwise, it creates, submits, and caches a new workflow.
0364
0365         Raises RuntimeError if iDDS is not available.
0366         """
0367         # Check cache first
0368         if stage_name in self.stage_workflows:
0369             self.logger.debug("Returning cached workflow for stage '%s'", stage_name)
0370             return self.stage_workflows[stage_name]
0371
0372         try:
0373             from idds.iworkflow.workflow import workflow as workflow_def  # type: ignore
0374         except Exception as exc:  # pragma: no cover - optional dependency
0375             raise RuntimeError("idds.iworkflow is not available; install idds to use PanDA runner") from exc
0376
0377         workflow_name = f"{self.config.name or 'aid2e'}.{stage_name}.{datetime.datetime.now().strftime('%Y%m%d_%H_%M_%S')}"
0378         self.logger.info("Defining workflow for experiment %s", workflow_name)
0379
0380         wf_builder = workflow_def(
0381             func=lambda: None,
0382             name=workflow_name,
0383             service="panda",
0384             cloud=self.config.cloud,
0385             queue=self.config.queue,
0386             init_env=self.config.init_env,
0387             source_dir=self.config.source_dir,
0388             source_dir_parent_level=self.config.source_dir_parent_level,
0389             exclude_source_files=self.config.exclude_source_files,
0390             max_walltime=self.config.max_walltime,
0391             core_count=self.config.core_count,
0392             total_memory=self.config.total_memory,
0393             enable_separate_log=self.config.enable_separate_log,
0394             local=True,
0395             return_workflow=True,
0396             post_script=self.config.post_script,
0397         )
0398
0399         workflow = wf_builder()
0400         workflow.pre_run()
0401         workflow.prepare()
0402         req_id = workflow.submit()
0403         self.logger.info("Workflow id for experiment %s: %s", workflow_name, req_id)
0404         if not req_id:
0405             raise RuntimeError(f"Failed to submit workflow for experiment {workflow_name} to PanDA")
0406
0407         # store for potential future use
0408         self.workflow = workflow
0409         self.workflow_id = req_id
0410         # Cache the workflow per stage_name
0411         self.stage_workflows[stage_name] = workflow
0412         return workflow
0413
0414     def submit_job(self, stage_name: str, job_definition: Dict[str, Any], working_dir: Optional[str] = None) -> None:
0415         """Submit a single function-based job to iDDS/PanDA.
0416
0417         The method expects job_definition to contain at least a 'function' key.
0418         This mirrors the upstream runner but intentionally keeps the interface
0419         loose: any missing integration points raise a RuntimeError describing
0420         the problem.
0421         """
0422         try:
0423             from idds.iworkflow.work import work as work_def  # type: ignore
0424         except Exception as exc:  # pragma: no cover - optional dependency
0425             raise RuntimeError("idds.iworkflow.work is not available; install idds to use PanDA runner") from exc
0426
0427         # ensure workflow exists
0428         workflow = self.submit_idds_workflow(stage_name)
0429
0430         job = job_definition
0431         job_id = job.get("job_id") or uuid.uuid4().hex
0432
0433         func = job.get("function")
0434         if func is None:
0435             raise ValueError("Job dict must contain 'function' to submit to PanDA/iDDS")
0436
0437         # Check that the function is not from __main__
0438         func_mod = getattr(func, "__module__", None)
0439         if func_mod == "__main__":
0440             raise RuntimeError(
0441                 f"Function '{getattr(func, '__name__', func)}' comes from module __main__. "
0442                 "Remote execution requires the function to be defined in a real importable module, not in __main__. "
0443                 "Otherwise, the remote environment will import the module __main__ which will execute the top-level code "
0444                 "and not find the function definition. "
0445                 "Please move the function to a proper module and reference it by its full module path."
0446             )
0447
0448         func_name = getattr(func, "__name__", str(func))
0449         work_name = f"{self.config.name or 'aid2e'}.{stage_name}.{job_id}.{func_name}"
0450         self.logger.info("Defining work %s", work_name)
0451
0452         # Initialize stage-level tracking if needed
0453         if stage_name not in self.running_funcs:
0454             self.running_funcs[stage_name] = {}
0455         if stage_name not in self.jobs:
0456             self.jobs[stage_name] = {}
0457
0458         # simple bookkeeping entry for this job
0459         self.running_funcs[stage_name][job_id] = {"funcs": {}}
0460
0461         # create a work object depending on job content
0462         params = job.get("params", {})
0463         work_builder = work_def(
0464             func=func,
0465             workflow=workflow,
0466             return_work=True,
0467             map_results=True,
0468             name=work_name,
0469             job_key=work_name,
0470             log_dataset_name=f"{work_name}.$WORKFLOWID.log/",
0471         )
0472         work = work_builder(**params)
0473
0474         # apply resource hints
0475         try:
0476             work.core_count = int(getattr(self.config, "core_count", 1))
0477         except Exception:
0478             pass
0479
0480         tf_id = work.submit()
0481         self.logger.info("Submitted work %s to PanDA/iDDS, transform id %s", work_name, tf_id)
0482         if not tf_id:
0483             raise RuntimeError(f"Failed to submit {work_name} to PanDA")
0484
0485         # store mapping under stage_name -> job_id -> funcs -> func_name
0486         self.running_funcs[stage_name][job_id]["funcs"][func_name] = {"work": work, "tf_id": tf_id, "status": "New", "results": None}
0487         self.jobs[stage_name][job_id] = tf_id
0488
0489     def check_single_job_status(self, job: Dict[str, Any], job_context: Any = None) -> None:
0490         """Check status of a single submitted job and update running_funcs state.
0491
0492         Expects job dict with 'job_id' and optionally 'stage_name'.
0493         If stage_name is not provided, searches all stages.
0494         Optionally accepts job_context for passing execution context.
0495         Raises on irrecoverable errors.
0496         """
0497         job_id = job.get("job_id")
0498         if not job_id:
0499             raise ValueError("job must contain 'job_id'")
0500
0501         # Extract stage_name from job_id if not provided (format: stage_name_job_name_index)
0502         stage_name = job.get("stage_name")
0503         if not stage_name:
0504             # Try to extract from job_id
0505             parts = job_id.split("_")
0506             if len(parts) >= 1:
0507                 # Assume first part is stage_name
0508                 stage_name = parts[0]
0509             else:
0510                 # Fallback: search all stages
0511                 for sname, stage_jobs in self.running_funcs.items():
0512                     if job_id in stage_jobs:
0513                         stage_name = sname
0514                         break
0515
0516         if not stage_name or stage_name not in self.running_funcs:
0517             raise RuntimeError(f"No running entry for job {job_id} (stage not found)")
0518
0519         # find the job entry in the stage
0520         entry = self.running_funcs[stage_name].get(job_id)
0521         if not entry:
0522             raise RuntimeError(f"No running entry for job {job_id} in stage {stage_name}")
0523
0524         # entry may have multiple func names; pick the first
0525         func_name, info = next(iter(entry.get("funcs", {}).items()), (None, None))
0526         if func_name is None:
0527             # older-style storage format
0528             func_name = next(iter(entry.keys()))
0529             info = entry[func_name]
0530
0531         work = info.get("work")
0532         tf_id = info.get("tf_id")
0533         if not work or not tf_id:
0534             raise RuntimeError(f"Job {job_id} has no work or no transform id")
0535
0536         # ensure async result initialized if available
0537         try:
0538             work.init_async_result()
0539         except Exception:
0540             pass
0541
0542         status = work.get_status()
0543         if work.is_finished(status):
0544             self.logger.info("Job %s finished (transform %s)", job_id, tf_id)
0545             try:
0546                 ret = work.get_results()
0547                 # try to extract mapped results
0548                 results = None
0549                 try:
0550                     results, _details = ret.get_result(name=work.name, key=info.get("job_key", work.name), verbose=True, with_details=True)
0551                     self.logger.debug(f"Extracted results for job {job_id}: {results}, details: {_details}")
0552
0553                     if job:
0554                         self.logger.debug("Job %s has context: %s", job_id, job_context)
0555                         # Optionally, you could store or use this context information as needed for your application
0556                         job_context.xcom_push("objectives", results)
0557                         job_context.xcom_push({"results_details": _details})
0558                 except Exception:
0559                     results = ret
0560                 info["results"] = results
0561             except Exception:
0562                 self.logger.exception("Failed to fetch results for job %s", job_id)
0563             info["status"] = "finished"
0564             # cleanup bookkeeping from stage
0565             self.running_funcs[stage_name].pop(job_id, None)
0566         elif work.is_failed(status):
0567             self.logger.info("Job %s failed (transform %s)", job_id, tf_id)
0568             info["status"] = "failed"
0569             self.running_funcs[stage_name].pop(job_id, None)
0570
0571     def check_job_status(self, job: Dict[str, Any]) -> None:
0572         """Wrapper around check_single_job_status that throttles verbose logs."""
0573         if self.num_checks % 60 == 0:
0574             self.logger.info("Check job %s status", job.get("job_id"))
0575         self.check_single_job_status(job, job.get("job_context"))
0576         self.num_checks += 1
0577
0578     # convenience alias for upstream name
0579     submit_workflow = submit_idds_workflow
0580