mockslurm.mock_sbatch

Implement a mock of the sbatch command of slurm.

The implementation mimics the sbatch API, while using subprocess to start the processes on the host machine without slurm daemons.

Double fork is used to detach the subprocess, allowing the sbatch mock to return while the created process is still running.

Classes:

Name	Description
DependencyActions	Scheduler behaviour with respect to job dependencies.

Functions:

Name	Description
check_dependency	Check if a job's dependency are satisfied.
launch_job	Run cmd as a detached subprocess.
parse_dependency	Parse a job's dependency specification string in sbatch format.

DependencyActions

Bases: Enum

              flowchart TD
              mockslurm.mock_sbatch.DependencyActions[DependencyActions]

              

              click mockslurm.mock_sbatch.DependencyActions href "" "mockslurm.mock_sbatch.DependencyActions"
            

Scheduler behaviour with respect to job dependencies.

Source code in src/mockslurm/mock_sbatch.py

class DependencyActions(Enum):
    """Scheduler behaviour with respect to job dependencies."""

    OK = 1
    WAIT = 2
    NEVER = 3

check_dependency

check_dependency(job_dependency)

Check if a job's dependency are satisfied.

Parameters:

Name	Type	Description	Default
job_dependency	Tuple[bool, List[Dict[str, List[int]]]]	Job dependencies. Each tuple in the list corresponds to a dependency specification. See mockslurm.mock_sbatch.parse_dependency	required

Warning

This function only supports the "afterok" dependency specification, anything else causes a NotImplementedError to be raised.

Raise

NotImplementedError If the dependency type is not "afterok"

Returns:

Type	Description
DependencyActions	The action to follow based on the dependency evaluation, for instance to wait, to start job, etc.

Source code in src/mockslurm/mock_sbatch.py

def check_dependency(job_dependency: list[tuple[bool, dict[str, list[int]]]]) -> str:
    """Check if a job's dependency are satisfied.

    Parameters
    ----------
    job_dependency : Tuple[bool, List[Dict[str, List[int]]]]
        Job dependencies. Each tuple in the list corresponds to a dependency specification.
        See `mockslurm.mock_sbatch.parse_dependency`

    Warning
    -------
    This function only supports the "afterok" dependency specification, anything else causes
    a `NotImplementedError` to be raised.

    Raise
    -----
    NotImplementedError
        If the dependency type is not "afterok"

    Returns
    -------
    DependencyActions
        The action to follow based on the dependency evaluation, for instance to wait, to start job, etc.
    """
    if not job_dependency[1]:
        return DependencyActions.OK

    deps_value = []
    with get_db_file_handle(find_db_file()) as db_file:
        db = get_db(db_file)
        for dep in job_dependency[1]:
            internal_deps_ok = True
            for dep_type, job_idx in dep.items():
                if dep_type != "afterok":
                    raise NotImplementedError(f"Dependency type {dep_type} is not implemented.")
                dep_state = db[job_idx if isinstance(job_idx, list) else [job_idx]]["STATE"]
                if any(dep_state == JobState.FAILED) or any(dep_state == JobState.CANCELLED):
                    return DependencyActions.NEVER
                internal_deps_ok &= all(dep_state == JobState.COMPLETED)

            deps_value.append(internal_deps_ok)

    combined_dep = all(deps_value) if job_dependency[0] else any(deps_value)
    return DependencyActions.OK if combined_dep else DependencyActions.WAIT

launch_job

launch_job(cmd, stdout, stderr, job_idx, job_dependency)

Run cmd as a detached subprocess.

The return code of cmd is retrieved by the detached subprocess and updated in the DB.

Parameters:

Name	Type	Description	Default
cmd	str	Command to run, similar to the content of the --wrap= argument of sbatch	required
stdout	str	Path to a file where the output of cmd will be piped	required
stderr	str	Path to a file where the errors of cmd will be piped	required
job_idx	int	Index of the job in the DB	required
job_dependency	List[Tuple[bool, Dict[str, List[int]]]]	Job dependencies. Each tuple in the list corresponds to a dependency specification. See mockslurm.mock_sbatch.parse_dependency	required

Source code in src/mockslurm/mock_sbatch.py

def launch_job(
    cmd: str, stdout: str, stderr: str, job_idx: int, job_dependency: list[tuple[bool, dict[str, list[int]]]]
):
    """Run `cmd` as a detached subprocess.

    The return code of `cmd` is retrieved by the detached subprocess and updated in the DB.

    Parameters
    ----------
    cmd : str
        Command to run, similar to the content of the --wrap= argument of sbatch
    stdout : str
        Path to a file where the output of `cmd` will be piped
    stderr : str
        Path to a file where the errors of `cmd` will be piped
    job_idx : int
        Index of the job in the DB
    job_dependency : List[Tuple[bool, Dict[str, List[int]]]]
        Job dependencies. Each tuple in the list corresponds to a dependency specification.
        See `mockslurm.mock_sbatch.parse_dependency`
    """
    logging.debug(
        f"lauch_job with args cmd: {cmd}, stdout: {stdout}, stderr: {stderr}, job_idx: {job_idx}, job_dependency: {job_dependency}"
    )
    # Wait for dependencies to be ready
    dependency_check = check_dependency(job_dependency)
    if dependency_check == DependencyActions.WAIT:
        logging.debug(f"Job {job_idx} dependency {job_dependency} 1st check is WAIT, updating db Reason DEPENDENCY")
        with get_db_file_handle(find_db_file()) as db_file:
            db = get_db(db_file)
            update_db_value(db_file, job_idx, key="REASON", value=JobReason.Dependency)
    while dependency_check == DependencyActions.WAIT:
        time.sleep(0.25)
        logging.debug(f"Job {job_idx} dependency {job_dependency} check remains WAIT, sleeping")
        dependency_check = check_dependency(job_dependency)

    # If not ok: do not start job and mark its state as FAILED
    if dependency_check != DependencyActions.OK:
        logging.debug(
            f"Job {job_idx} dependency check is {dependency_check}, updating DB with REASON DependencyNeverSatisfied and STATE FAILED"
        )
        with get_db_file_handle(find_db_file()) as db_file:
            db = get_db(db_file)
            update_db_value(db_file, job_idx, key="REASON", value=JobReason.DependencyNeverSatisfied)
            update_db_value(db_file, job_idx, key="STATE", value=JobState.FAILED)
    else:
        logging.debug(f"Job {job_idx} dependency check is OK, checking job STATE")
        # lock the DB here, so STATE can not change before we start the process
        with get_db_file_handle(find_db_file()) as db_file:
            db = get_db(db_file)
            if db[job_idx]["STATE"] == JobState.PENDING:
                logging.debug(f"Job {job_idx} STATE is still pending, starting job")
                try:
                    stdout_f = open(stdout, "a")
                    stderr_f = sp.STDOUT if stderr == stdout else open(stderr, "a")
                    # Can not use shlex here to split command, otherwise we would split bash commands at each
                    # space so we use a single string with shell=True
                    p = sp.Popen(cmd, stdout=stdout_f, stderr=stderr_f, start_new_session=True, shell=True)
                    logging.debug(
                        f"Job {job_idx} started with PID {p.pid}, cmd {cmd}, stdout {stdout_f}, stderr {stderr_f}"
                    )
                    update_db_value(db_file, job_idx, key="PID", value=p.pid)
                    update_db_value(db_file, job_idx, key="STATE", value=JobState.RUNNING)
                    update_db_value(db_file, job_idx, key="REASON", value=JobReason.NOREASON)
                    logging.debug(f"Job {job_idx} updated db with PID {p.pid}, STATE RUNNING, REASON NOREASON")
                except:
                    logging.debug("Job {} failed to start. Updating DB with STATE FAILED and REASON JobLaunchFailure")
                    update_db_value(db_file, job_idx, key="STATE", value=JobState.FAILED)
                    update_db_value(db_file, job_idx, key="REASON", value=JobReason.JobLaunchFailure)
                    logging.debug("Job {} starting error", exc_info=True)
                    raise
            else:
                logging.debug(
                    "Job {} STATE is {}. Do not start start job and exit".format(job_idx, db[job_idx]["STATE"])
                )
                # Process is not pending anymore, it must have been killed already, so we won't start it
                return

        logging.debug(f"Job {job_idx} is running, waiting its completion")
        # wait for process to be done
        exit_code = p.wait()
        logging.debug(f"Job {job_idx} completed with exit code {exit_code}")
        # closing stdout and stderr file of job
        stdout_f.close()
        if stderr != stdout:
            stderr_f.close()
        logging.debug(f"Job {job_idx} closed stdout {stdout} and stderr {stderr}")
        # Update job state in db
        with get_db_file_handle(find_db_file()) as db_file:
            job_state = JobState.COMPLETED if exit_code == 0 else JobState.FAILED
            job_reason = JobReason.NOREASON if exit_code == 0 else JobReason.NonZeroExitCode
            update_db_value(db_file, job_idx, key="EXIT_CODE", value=exit_code)
            update_db_value(db_file, job_idx, key="STATE", value=job_state)
            update_db_value(db_file, job_idx, key="REASON", value=job_reason)
            logging.debug(
                f"Job {job_idx} Updated db with EXIT_CODE {exit_code}, STATE {JobState(job_state).name}, REASON {JobReason(job_reason).name}"
            )

parse_dependency

parse_dependency(job_dependency)

Parse a job's dependency specification string in sbatch format.

Parameters:

Name	Type	Description	Default
job_dependency	str	sbatch job's dependency string. Examples: afterok:20:21,afterany:23	required

Returns:

Type	Description
Tuple[bool, List[Dict[str, List[int]]]]	Tuple 1st element is whether all dependencies must be satisfied for a job to start (True, "," in slurm str), or if a single dependency is enough (False, "?" in slurm str). When a single dependency is present: True. The list values are dictionaries with a single key been the dependency type ("afterok", "afterany", etc.) and the values been the job IDs for these dependency

Example

parse_dependency("afterok:20:21?afterany:23") (False, [{'afterok': [20, 21]}, {'afterany': [23]}])

Source code in src/mockslurm/mock_sbatch.py

def parse_dependency(job_dependency: str) -> tuple[bool, list[dict[str, list[int]]]]:
    """Parse a job's dependency specification string in sbatch format.

    Parameters
    ----------
    job_dependency : str
        sbatch job's dependency string. Examples: afterok:20:21,afterany:23

    Returns
    -------
    Tuple[bool, List[Dict[str, List[int]]]]
        Tuple 1st element is whether all dependencies must be satisfied for a job to start (True, "," in slurm str),
        or if a single dependency is enough (False, "?" in slurm str). When a single dependency is present: True.
        The list values are dictionaries with a single key been the dependency type ("afterok", "afterany", etc.)
        and the values been the job IDs for these dependency
    Example
    -------
    >>> parse_dependency("afterok:20:21?afterany:23")
    (False, [{'afterok': [20, 21]}, {'afterany': [23]}])
    """
    if "," in job_dependency and "?" in job_dependency:
        raise ValueError('Multiple separator "?", ","in job dependency not supported.')
    separator = "?" if "?" in job_dependency else ","  # default is "," if only 1 dep
    dependencies = []
    while job_dependency:
        if job_dependency[0] == separator:
            job_dependency = job_dependency[1:]

        dep_str = job_dependency.split(separator, maxsplit=1)[0]
        job_dependency = job_dependency[len(dep_str) :]
        split = dep_str.split(":")
        dependencies.append({split[0]: [int(v) for v in split[1:]]})

    return (separator == ",", dependencies)