Source code for job_execution.execution

"""Public execution facade used by long-running applications."""

import json
import logging
import subprocess
import time
from pathlib import Path

import yaml

from simtools.context import RunContext
from simtools.io.file_type import product_filename
from simtools.job_execution.backends.registry import get_backend
from simtools.job_execution.job import (
    CommandOptions,
    ExecutionOptions,
    JobResult,
    JobSpec,
    SubmissionHandle,
)
from simtools.job_execution.log_inspector import get_log_excerpt
from simtools.job_execution.worker import execute_command

logger = logging.getLogger(__name__)


[docs] class JobExecutionError(RuntimeError): """Raised when a command job fails."""
[docs] def run_command( command, *, stdin=None, stdout=None, stderr=None, env=None, runtime_environment=None, configuration=None, application_log=None, check=True, capture_output=True, ): """ Run one command synchronously through the execution worker. Parameters ---------- command : str or list[str] or Path Command (a string or script path runs through the shell) or argument list. stdin : file object or str or Path, optional Standard input: an open file, or a file path that is opened for reading. stdout, stderr : str or Path, optional Files receiving standard output and standard error. env : dict, optional Environment variables added to the current environment. runtime_environment : str or list[str], optional Command prefix running the command in a container. configuration : dict, optional Application configuration appended as command-line arguments. application_log : str or Path, optional Application log; an excerpt is added to the error message on failure. check : bool, optional Raise ``JobExecutionError`` on a non-zero return code. capture_output : bool, optional Capture standard output and error that are not written to files; otherwise they are inherited. Returns ------- JobResult Return code, wall-clock runtime and the captured output (``output``) and error output (``error_output``). Raises ------ JobExecutionError If the command fails and ``check`` is True; the message contains the return code, the captured standard error and excerpts from the log files. """ job = JobSpec( job_id="command", index=0, command=tuple(command) if isinstance(command, list) else command, options=CommandOptions( stdin=stdin, stdout=Path(stdout) if stdout is not None else None, stderr=Path(stderr) if stderr is not None else None, env=dict(env or {}), runtime_environment=runtime_environment, configuration=configuration, application_log=Path(application_log) if application_log else None, check=check, capture_output=capture_output, ), ) try: return_code, runtime, output, error_output = execute_command(job) except subprocess.CalledProcessError as exc: _raise_command_error(exc, job) return JobResult( job_id=job.job_id, index=job.index, return_code=return_code, runtime=runtime, output=output, error_output=error_output, )
def _raise_command_error(exc, job): """Raise a command error with the return code, captured stderr and log excerpts.""" details = [f"Command failed with return code {exc.returncode}."] if exc.stderr: details.append(f"Standard error:\n{exc.stderr}") for label, path in (("Output", job.options.stdout), ("Error", job.options.stderr)): if path is not None: details.append(f"{label} log excerpt from {path}:\n{get_log_excerpt(path)}") log = job.options.application_log if log and log.exists() and (time.time() - log.stat().st_mtime) / 60 < 5: details.append(f"Application log excerpt from {log}:\n{get_log_excerpt(log)}") raise JobExecutionError("\n".join(details)) from exc
[docs] def execute_jobs(job_specs, options=None): """Execute ordered jobs through the selected backend.""" options = options or ExecutionOptions() jobs = list(job_specs) _validate_jobs(jobs) if not jobs: return [] backend = get_backend(options.backend) logger.info("Executing %d job(s) with backend %s", len(jobs), options.backend) submission = _submit_validated(jobs, options, backend) results = wait_for_submission(submission, options=options, backend=backend) logger.info("Completed %d job(s) with backend %s", len(results), options.backend) return sorted(results, key=lambda result: result.index)
[docs] def submit_jobs(job_specs, options=None): """Submit independent jobs, detaching when the backend supports it. Scheduler-backed implementations return without waiting. Backends that do not advertise submit-only support complete synchronously. """ options = options or ExecutionOptions() jobs = list(job_specs) _validate_jobs(jobs) if not jobs: return SubmissionHandle( backend=options.backend, work_dir=Path(options.work_dir or Path.cwd()), job_ids=(), ) backend = get_backend(options.backend) submission = _submit_validated(jobs, options, backend) supports_submit_only = getattr(backend, "supports_submit_only", False) if not supports_submit_only: wait_for_submission(submission, options=options, backend=backend) if jobs: if supports_submit_only: logger.info( "Submitted %d job(s) with backend %s; manifest: %s", len(jobs), options.backend, submission.work_dir / product_filename("submission_manifest"), ) else: logger.info( "Completed %d job(s) synchronously with backend %s", len(jobs), options.backend, ) return submission
[docs] def wait_for_submission(submission, *, options=None, backend=None): """Wait for a previously submitted job set and validate its outputs.""" backend = backend or get_backend(submission.backend) try: results = backend.wait(submission) _validate_manifest_outputs(submission) except KeyboardInterrupt: _mark_submission(submission, "interrupted") _cancel_if_requested(backend, submission, options) raise except Exception: _mark_submission(submission, "failed") raise _mark_submission(submission, "completed") return sorted(results, key=lambda result: result.index)
[docs] def collect_submission(submission, *, backend=None): """Collect finished scheduler jobs without waiting for active jobs. Parameters ---------- submission : SubmissionHandle Persisted scheduler submission to inspect. backend : object or None Execution backend; constructed from the submission when omitted. Returns ------- list[JobResult] or None Completed results, or None while jobs remain in the scheduler queue. """ backend = backend or get_backend(submission.backend) if not backend.is_finished(submission): return None try: results = backend.collect(submission) _validate_manifest_outputs(submission) except Exception: _mark_submission(submission, "failed") raise _mark_submission(submission, "completed") return results
[docs] def load_submission(path): """Load a submission handle from a JSON manifest.""" manifest = Path(path) try: payload = json.loads(manifest.read_text(encoding="utf-8")) except OSError as exc: raise FileNotFoundError(f"Submission manifest not found: {manifest}") from exc except json.JSONDecodeError as exc: raise ValueError(f"Invalid submission manifest: {manifest}") from exc return SubmissionHandle.from_dict(payload)
def _submit_validated(jobs, options, backend): """Submit validated jobs and persist a remote manifest.""" submission = backend.submit(jobs, options) submission.metadata["expected_outputs"] = { job.job_id: [str(Path(path).expanduser().resolve()) for path in job.output_paths] for job in jobs if job.output_paths } _mark_submission(submission, "submitted") return submission def _validate_jobs(jobs): """Validate stable identifiers and expected outputs before submission.""" if len({job.job_id for job in jobs}) != len(jobs): raise ValueError("Execution jobs must have unique job IDs.") if len({job.index for job in jobs}) != len(jobs): raise ValueError("Execution jobs must have unique input-order indices.") expected_outputs = [ Path(path).expanduser().resolve() for job in jobs for path in job.output_paths ] if len(expected_outputs) != len(set(expected_outputs)): raise ValueError("Execution jobs contain duplicate expected output paths.") def _mark_submission(submission, state): """Update the remote submission manifest when applicable.""" if submission.backend == "local": return submission.metadata["state"] = state _write_manifest(submission) def _cancel_if_requested(backend, submission, options): """Cancel a remote submission when interruption policy requests it.""" configured = submission.metadata.get("cancel_on_interrupt", False) if options is not None: configured = configured or (options.backend_config or {}).get("cancel_on_interrupt", False) if configured: backend.cancel(submission) def _validate_manifest_outputs(submission): """Validate output paths recorded in a durable submission manifest.""" expected = submission.metadata.get("expected_outputs", {}) if isinstance(expected, list): expected = {"unknown": expected} missing_outputs = [ path for paths in expected.values() for path in paths if not Path(path).exists() ] if missing_outputs: raise FileNotFoundError( "Execution completed without expected output(s): " + ", ".join(missing_outputs) )
[docs] def map_ordered( function, items, *, backend="local", max_workers=None, backend_config=None, initializer=None, initargs=(), mp_start_method="fork", work_dir=None, context=None, ): """Apply ``function`` to items and return values in input order.""" context = context or RunContext.from_environment() job_specs = [ JobSpec( job_id=f"job-{index:06d}", index=index, function=function, item=item, initializer=initializer if backend != "local" else None, initargs=tuple(initargs) if backend != "local" else (), context_spec=context.to_spec(), ) for index, item in enumerate(items) ] execution_config = _load_backend_config(backend_config) if backend == "local": execution_config.setdefault("mp_start_method", mp_start_method) options = ExecutionOptions( backend=backend, max_workers=max_workers, work_dir=Path(work_dir) if work_dir else None, backend_config=execution_config, initializer=initializer, initargs=tuple(initargs), ) return [result.value for result in execute_jobs(job_specs, options)]
[docs] def options_from_args(args, *, max_workers=None, work_dir=None): """Build execution options from an application argument dictionary.""" config_value = args.get("backend_config") backend_config = _load_backend_config(config_value) return ExecutionOptions( backend=args.get("backend", "local"), max_workers=max_workers, work_dir=Path(work_dir) if work_dir else None, backend_config=backend_config, )
def _load_backend_config(config_value): """Load an inline or file-based backend configuration.""" if isinstance(config_value, dict): return dict(config_value) if config_value: config_path = Path(config_value) if not config_path.is_file(): raise FileNotFoundError(f"Backend configuration file not found: {config_path}") with config_path.open(encoding="utf-8") as handle: backend_config = yaml.safe_load(handle) or {} if not isinstance(backend_config, dict): raise ValueError("Backend configuration must contain a mapping.") else: return {} return backend_config def _write_manifest(submission): """Write a durable submission manifest.""" submission.work_dir.mkdir(parents=True, exist_ok=True) manifest = submission.work_dir / product_filename("submission_manifest") temporary = manifest.with_suffix(".tmp") temporary.write_text( json.dumps(submission.as_dict(), indent=2, default=str) + "\n", encoding="utf-8", ) temporary.replace(manifest)