Source code for psij.launchers.script_based_launcher

import hashlib
import logging
import secrets
import threading
from pathlib import Path
from typing import Optional, List, Any

from psij.job_launcher import Launcher
from psij.job_executor_config import JobExecutorConfig
from psij.job import Job

logger = logging.getLogger(__name__)

def _str(obj: Optional[object]) -> str:
    if obj is not None:
        return str(obj)
        return ''

def _path(obj: Optional[object]) -> str:
    if obj is None:
        return '/dev/null'
        return str(obj)

_BLOCK_SZ = 16384

def _copy_and_checksum(src_path: Path, dst_path: Path, s: Any) -> None:
    with open(src_path, 'rb') as src:
        with open(dst_path, 'wb') as dst:
            bytes =

[docs]class ScriptBasedLauncher(Launcher): """ A launcher that uses a script to start the job, possibly by wrapping it in other tools. This launcher is an abstract base class for launchers that wrap the job in a script. The script must be a bash script and is invoked with the first four parameters as: * the job ID * a launcher log file, which is taken from the launcher_log_file configuration setting and defaults to `/dev/null` * the pre- and post- launcher scripts, or empty strings if they are not specified Additional positional arguments to the script can be specified by subclasses by overriding the :meth:`~.get_additional_args` method. The remaining arguments to the script are the job executable and arguments. A simple script library is provided in scripts/ Its use is optional and it is intended to be included at the beginning of a main launcher script using `source $(dirname "$0")/`. It does the following: * sets '-e' mode (exit on error) * sets the variables _PSI_J_JOB_ID, _PSI_J_LOG_FILE, _PSI_J_PRE_LAUNCH, and _PSI_J_POST_LAUNCH from the first arguments, as specified above. * saves the current stdout and stderr in descriptors 3 and 4, respectively * redirects stdout and stderr to the log file, while prepending a timestamp and the job ID to each line * defines the commands "pre_launch" and "post_launch", which can be invoked by the main script. When invoking the job executable (either directly or through a launch command), it is recommended that the stdout and stderr of the job process be redirected to descriptors 3 and 4, respectively, such that they can be captured by the entity invoking the launcher rather than ending up in a the launcher log file. A successful completion of the launcher should be signalled by the launcher by printing the string "_PSI_J_LAUNCHER_DONE" to stdout. The launcher can then exit with the exit code returned by the launched command. This allows executor to distinguish between a non-zero exit code due to application failure or due to a premature launcher failure. The actual launcher scripts, as well as the library, are deployed at run-time into the work directory, where submit scripts are also generated. This directory is meant to be accessible by both the node submitting the job as well as the node launching the job. """ def __init__(self, script_path: Path, config: Optional[JobExecutorConfig] = None) -> None: """ Parameters ---------- script_path A path to a script that is invoked as described above. config An optional configuration. """ super().__init__(config) self._script_path = script_path self._log_file = '' if config and config.launcher_log_file: self._log_file = str(config.launcher_log_file) self._lock = threading.RLock() self._deployed = False def _ensure_launcher_deployed(self) -> None: with self._lock: if self._deployed: return deploy_dir = self._deploy_files(self._files_to_deploy()) self._deployed_script_path = deploy_dir / self._deployed = True def _files_to_deploy(self) -> List[Path]: return [ Path(__file__).parent / 'scripts' / '', self._script_path ] def _deploy_files(self, paths: List[Path]) -> Path: """ Deploys all files in a list to the same subdirectory of `config.work_directory`. The directory is guaranteed to be unique for each distinct list of files. That is, two invocations of this method with byte-by-byte identical files pointed to by `paths` and in the same order, will lead to the same subdirectory; changing the order of files in `paths` or changing the contents of one or more files in the list will lead to a different subdirectory. This allows files deployed together through this method to maintain relative references to each other while ensuring that different processes using different files do not interfere with each other. """ wrk_dir = self.config.work_directory wrk_dir.mkdir(parents=True, exist_ok=True) tmp_prefix = secrets.token_hex() + '_' h = hashlib.sha256() tmp_paths = {} for path in paths: tmp_path = wrk_dir / (tmp_prefix + tmp_paths[path] = tmp_path _copy_and_checksum(path, tmp_path, h) dst_dir = wrk_dir / h.hexdigest() if dst_dir.exists(): for path in paths: tmp_paths[path].unlink() return dst_dir dst_dir.mkdir(exist_ok=True) for path in paths: dst_path = dst_dir / try: # this appears to use os.rename, although I'd wish pathlib docs would state this # explicitly, since the docs for Path.rename mention nothing of the exceptions # thrown tmp_paths[path].rename(dst_path) except FileExistsError: # thrown in Windows if the path already exists, which is fine if the destination is # a file; we were racing another process if dst_path.is_dir(): raise else: continue except IsADirectoryError: # this is thrown in Unix if the destination exists and is a directory; we're not # expecting this raise return dst_dir
[docs] def get_launch_command(self, job: Job, log_file: Optional[str] = None) -> List[str]: """See :func:`~psij.Launcher.get_launch_command`.""" spec = job.spec assert spec is not None self._ensure_launcher_deployed() if log_file is None: log_file = self._log_file args = ['/bin/bash', str(self._deployed_script_path),, _str(log_file), _str(spec.pre_launch), _str(spec.post_launch), _path(spec.stdin_path), _path(spec.stdout_path), _path(spec.stderr_path)] args += self.get_additional_args(job) assert spec.executable is not None args += [spec.executable] if spec.arguments is not None: args += spec.arguments return args
[docs] def get_additional_args(self, job: Job) -> List[str]: """ Returns any additional arguments, after first mandatory four, to be passed to the script. :param job: The job that is being launched. """ return []
[docs] def is_launcher_failure(self, output: str) -> bool: """See :func:`~psij.Launcher.is_launcher_failure`.""" return output.split('\n')[-2] != '_PSI_J_LAUNCHER_DONE'
[docs] def get_launcher_failure_message(self, output: str) -> str: """See :func:`~psij.Launcher.get_launcher_failure_message`.""" return '\n'.join(output.split('\n')[:-2])