"""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)