schedulers/JobLib/runner.py

0001 """JobLib-based scheduler for local parallel job execution."""
0002
0003 import json
0004 import logging
0005 import os
0006 import pickle
0007 import subprocess
0008 from typing import Dict, Any, List, Optional
0009
0010 import joblib
0011
0012 from aid2e.schedulers.base import BaseScheduler, JobStatus, StageExecutionResult
0013 from aid2e.schedulers.JobLib.config import JobLibRunnerConfig
0014
0015
0016 class JobLibScheduler(BaseScheduler):
0017     """Execute workflow stage jobs in parallel using joblib."""
0018
0019     def __init__(self, config: Optional[JobLibRunnerConfig] = None) -> None:
0020         """Initialize JobLib scheduler with the provided configuration."""
0021
0022         super().__init__(config)
0023         self.config = config or JobLibRunnerConfig()
0024         self.logger = logging.getLogger("JobLibScheduler")
0025         self.running_jobs: Dict[str, Dict[str, Any]] = {}
0026
0027     def _execute_job(self, job_def: Dict[str, Any], working_dir: Optional[str] = None) -> Dict[str, Any]:
0028         """Execute a single job command and collect outputs."""
0029
0030         job_name = job_def.get("name", "unknown_job")
0031         command = job_def.get("command", "")
0032         payload = job_def.get("payload", {})
0033         output_specs = job_def.get("outputs", [])
0034
0035         # Check if this is a Python callable job
0036         python_callable = job_def.get("function")
0037         if python_callable and callable(python_callable):
0038             return self._execute_python_callable(job_def, python_callable, working_dir)
0039
0040         try:
0041             env = os.environ.copy()
0042
0043             # Try to serialize payload with pickle first (handles more types including functions)
0044             try:
0045                 import base64
0046                 pickled_payload = pickle.dumps(payload)
0047                 env["JOB_PAYLOAD_PICKLE"] = base64.b64encode(pickled_payload).decode('ascii')
0048                 env["JOB_PAYLOAD_TYPE"] = "pickle"
0049             except Exception as pickle_err:
0050                 self.logger.debug("Cannot pickle payload, falling back to JSON: %s", pickle_err)
0051                 # Fall back to JSON for simple payloads
0052                 serializable_payload = {}
0053                 for key, value in payload.items():
0054                     if not callable(value):
0055                         try:
0056                             # Test if it's JSON serializable
0057                             json.dumps(value)
0058                             serializable_payload[key] = value
0059                         except (TypeError, ValueError):
0060                             # Skip non-serializable values
0061                             self.logger.debug("Skipping non-serializable payload key: %s", key)
0062
0063                 env["JOB_PAYLOAD"] = json.dumps(serializable_payload)
0064                 env["JOB_PAYLOAD_TYPE"] = "json"
0065
0066             cwd = working_dir or os.getcwd()
0067             self.logger.info("Executing job '%s': %s", job_name, command)
0068
0069             timeout_sec = self.config.timeout if self.config.timeout else None
0070             result = subprocess.run(
0071                 command,
0072                 shell=True,
0073                 capture_output=True,
0074                 text=True,
0075                 cwd=cwd,
0076                 env=env,
0077                 timeout=timeout_sec,
0078             )
0079
0080             artifacts: Dict[str, Any] = {}
0081             for spec in output_specs:
0082                 output_path = spec.get("path", "")
0083                 if not output_path:
0084                     continue
0085                 full_path = os.path.join(cwd, output_path)
0086                 if os.path.exists(full_path):
0087                     with open(full_path, "r", encoding="utf-8") as handle:
0088                         artifacts[output_path] = handle.read()
0089                 else:
0090                     self.logger.warning("Expected output file not found: %s", full_path)
0091
0092             return {
0093                 "stdout": result.stdout,
0094                 "stderr": result.stderr,
0095                 "return_code": result.returncode,
0096                 "outputs": artifacts,
0097             }
0098
0099         except subprocess.TimeoutExpired:
0100             self.logger.error("Job '%s' timed out after %ss", job_name, self.config.timeout)
0101             return {
0102                 "stdout": "",
0103                 "stderr": f"Job timed out after {self.config.timeout}s",
0104                 "return_code": -1,
0105                 "outputs": {},
0106             }
0107         except Exception as exc:  # pragma: no cover - catch-all safety
0108             self.logger.error("Job '%s' raised exception: %s", job_name, exc)
0109             return {
0110                 "stdout": "",
0111                 "stderr": str(exc),
0112                 "return_code": -1,
0113                 "outputs": {},
0114             }
0115
0116     def _execute_python_callable(
0117         self, job_def: Dict[str, Any], python_callable, working_dir: Optional[str] = None
0118     ) -> Dict[str, Any]:
0119         """Execute a Python callable directly (for function-based jobs).
0120
0121         Args:
0122             job_def: Job definition dict.
0123             python_callable: The Python function to call.
0124             working_dir: Working directory (unused for Python callables).
0125
0126         Returns:
0127             Dict with execution results.
0128         """
0129         job_name = job_def.get("name", "unknown_job")
0130         params = job_def.get("params", {})
0131
0132         try:
0133             self.logger.info("Executing Python callable for job '%s'", job_name)
0134
0135             # Extract the context if provided
0136             context = params.get("context")
0137             if context:
0138                 # Call with context as first argument
0139                 result = python_callable(context, **{k: v for k, v in params.items() if k != "context"})
0140             else:
0141                 # Call with just kwargs
0142                 result = python_callable(**params)
0143
0144             # Convert result to string for stdout
0145             result_str = str(result) if result is not None else ""
0146
0147             # Extract xcom data from context (if context was provided)
0148             outputs = {"result": result}
0149             if context and hasattr(context, 'xcom'):
0150                 # Extract xcom entries pushed by the callable
0151                 # XCom keys are in format "job_id:key", extract just the key part
0152                 for xcom_key, xcom_value in context.xcom.items():
0153                     # Extract the key after the job_id prefix
0154                     if ':' in xcom_key:
0155                         key = xcom_key.split(':', 1)[1]
0156                         outputs[key] = xcom_value
0157
0158             return {
0159                 "stdout": result_str,
0160                 "stderr": "",
0161                 "return_code": 0,
0162                 "outputs": outputs,
0163             }
0164
0165         except Exception as exc:
0166             self.logger.exception("Python callable '%s' raised exception", job_name)
0167             return {
0168                 "stdout": "",
0169                 "stderr": str(exc),
0170                 "return_code": -1,
0171                 "outputs": {},
0172             }
0173
0174     def run_stage(
0175         self,
0176         stage_name: str,
0177         job_definitions: List[Dict[str, Any]],
0178         parallelism_policy: Optional[Dict[str, Any]] = None,
0179         working_dir: Optional[str] = None,
0180     ) -> StageExecutionResult:
0181         """Execute all jobs in a stage using joblib.Parallel."""
0182
0183         policy = parallelism_policy or {}
0184         max_concurrent = policy.get("max_concurrent", self.config.n_jobs)
0185         retry_max = policy.get("retry_max", 2)
0186
0187         self.logger.info("Running stage '%s' with %d jobs", stage_name, len(job_definitions))
0188         self.logger.info("  Max concurrent: %s, Max retries: %s", max_concurrent, retry_max)
0189
0190         n_jobs = max_concurrent if max_concurrent and max_concurrent > 0 else self.config.n_jobs
0191
0192         try:
0193             parallel = joblib.Parallel(
0194                 n_jobs=n_jobs,
0195                 backend=self.config.backend,
0196                 verbose=self.config.verbose,
0197             )
0198
0199             job_results = parallel(
0200                 joblib.delayed(self._execute_job)(job_def, working_dir)
0201                 for job_def in job_definitions
0202             )
0203
0204             job_statuses: List[JobStatus] = []
0205             all_artifacts: Dict[str, Any] = {}
0206             all_success = True
0207
0208             for index, (job_def, result) in enumerate(zip(job_definitions, job_results)):
0209                 job_name = job_def.get("name", f"job_{index}")
0210                 job_id = job_def.get("job_id") or f"{stage_name}_{job_name}_{index}"
0211
0212                 return_code = result.get("return_code", -1)
0213                 success = return_code == 0
0214                 status = "completed" if success else "failed"
0215
0216                 job_statuses.append(
0217                     JobStatus(
0218                         job_id=job_id,
0219                         status=status,
0220                         return_code=return_code,
0221                         stdout=result.get("stdout", ""),
0222                         stderr=result.get("stderr", ""),
0223                         outputs=result.get("outputs"),
0224                     )
0225                 )
0226
0227                 if result.get("outputs"):
0228                     all_artifacts.update(result["outputs"])
0229
0230                 if not success:
0231                     all_success = False
0232                     self.logger.warning("Job '%s' failed with code %s", job_id, return_code)
0233
0234             return StageExecutionResult(
0235                 stage_name=stage_name,
0236                 job_statuses=job_statuses,
0237                 artifacts=all_artifacts,
0238                 success=all_success,
0239                 error_message=None if all_success else f"Some jobs failed in stage '{stage_name}'",
0240             )
0241
0242         except Exception as exc:  # pragma: no cover - safety net
0243             self.logger.error("Stage '%s' execution failed: %s", stage_name, exc)
0244             return StageExecutionResult(
0245                 stage_name=stage_name,
0246                 job_statuses=[],
0247                 artifacts={},
0248                 success=False,
0249                 error_message=str(exc),
0250             )
0251
0252     def check_status(self, job_id: str) -> JobStatus:
0253         """Return cached status (JobLib is synchronous, so jobs finish in run_stage)."""
0254
0255         if job_id in self.running_jobs:
0256             cached = self.running_jobs[job_id]
0257             return JobStatus(
0258                 job_id=job_id,
0259                 status=cached.get("status", "unknown"),
0260                 return_code=cached.get("return_code"),
0261             )
0262
0263         return JobStatus(job_id=job_id, status="unknown", return_code=None)
0264
0265     def cancel_job(self, job_id: str) -> bool:
0266         """Indicate that cancellation is not supported for synchronous JobLib jobs."""
0267
0268         self.logger.warning("Cannot cancel job '%s' (JobLib execution is synchronous)", job_id)
0269         return False
0270
0271     def submit_stage(
0272         self,
0273         stage_name: str,
0274         job_definitions: List[Dict[str, Any]],
0275         parallelism_policy: Optional[Dict[str, Any]] = None,
0276         working_dir: Optional[str] = None,
0277     ) -> str:
0278         """Submit a stage for execution.
0279
0280         For JobLib (synchronous execution), this immediately runs the stage
0281         and returns a stage_id for the completed stage.
0282         """
0283         import uuid
0284
0285         stage_id = uuid.uuid4().hex
0286         self.logger.info("Submitting stage '%s' as %s (synchronous execution)", stage_name, stage_id)
0287
0288         # Execute the stage immediately (JobLib is synchronous)
0289         result = self.run_stage(stage_name, job_definitions, parallelism_policy, working_dir)
0290
0291         # Cache the result
0292         if not hasattr(self, '_stage_results'):
0293             self._stage_results = {}
0294         self._stage_results[stage_id] = {
0295             'stage_name': stage_name,
0296             'result': result,
0297             'status': 'completed' if result.success else 'failed',
0298         }
0299
0300         return stage_id
0301
0302     def check_stage_status(self, stage_id: str):
0303         """Return status for a submitted stage.
0304
0305         Since JobLib is synchronous, stages are always completed by the time this is called.
0306         """
0307         from aid2e.schedulers.base import StageStatus
0308
0309         if not hasattr(self, '_stage_results'):
0310             self._stage_results = {}
0311
0312         if stage_id not in self._stage_results:
0313             raise KeyError(f"Unknown stage_id: {stage_id}")
0314
0315         stage_data = self._stage_results[stage_id]
0316         result = stage_data['result']
0317
0318         return StageStatus(
0319             stage_id=stage_id,
0320             status=stage_data['status'],
0321             completed_jobs=len(result.job_statuses),
0322             total_jobs=len(result.job_statuses),
0323             progress=1.0,  # Always complete for synchronous execution
0324             job_statuses=result.job_statuses,
0325         )
0326
0327     def get_stage_results(self, stage_id: str) -> StageExecutionResult:
0328         """Return results for a completed stage.
0329
0330         Since JobLib is synchronous, results are available immediately after submit_stage.
0331         """
0332         if not hasattr(self, '_stage_results'):
0333             self._stage_results = {}
0334
0335         if stage_id not in self._stage_results:
0336             raise KeyError(f"Unknown stage_id: {stage_id}")
0337
0338         stage_data = self._stage_results[stage_id]
0339         return stage_data['result']
0340
0341     def shutdown(self) -> None:
0342         """No-op shutdown hook for JobLib scheduler."""
0343
0344         self.logger.debug("JobLibScheduler shutdown complete")
0345         return None