nshtrainer 0.1.0__py3-none-any.whl → 0.2.0__py3-none-any.whl

nshtrainer/_submit/session/slurm.py

@@ -1,573 +0,0 @@
-import copy
-import logging
-import math
-import os
-import signal
-from collections.abc import Callable, Mapping, Sequence
-from datetime import timedelta
-from pathlib import Path
-from typing import Any, Literal, cast
-
-from deepmerge import always_merger
-from typing_extensions import TypeAlias, TypedDict, TypeVarTuple, Unpack
-
-from ._output import SubmitOutput
-from ._script import helper_script_to_command, write_helper_script
-
-log = logging.getLogger(__name__)
-
-TArgs = TypeVarTuple("TArgs")
-
-_Path: TypeAlias = str | Path | os.PathLike
-MailType: TypeAlias = Literal[
-    "NONE",
-    "BEGIN",
-    "END",
-    "FAIL",
-    "REQUEUE",
-    "ALL",
-    "INVALID_DEPEND",
-    "STAGE_OUT",
-    "TIME_LIMIT",
-    "TIME_LIMIT_90",
-    "TIME_LIMIT_80",
-    "TIME_LIMIT_50",
-    "ARRAY_TASKS",
-]
-
-
-class SlurmJobKwargs(TypedDict, total=False):
-    name: str
-    """
-    The name of the job.
-
-    This corresponds to the "-J" option in sbatch.
-    """
-
-    account: str
-    """
-    The account to charge resources used by this job to.
-
-    This corresponds to the "-A" option in sbatch.
-    """
-
-    partition: str | Sequence[str]
-    """
-    The partition to submit the job to.
-
-    This corresponds to the "-p" option in sbatch. If not specified, the default partition will be used.
-    Multiple partitions can be specified, and they will be combined using logical OR.
-    """
-
-    qos: str
-    """
-    The quality of service to submit the job to.
-
-    This corresponds to the "--qos" option in sbatch.
-    """
-
-    output_file: _Path
-    """
-    The file to write the job output to.
-
-    This corresponds to the "-o" option in sbatch. If not specified, the output will be written to the default output file.
-    """
-
-    error_file: _Path
-    """
-    The file to write the job errors to.
-
-    This corresponds to the "-e" option in sbatch. If not specified, the errors will be written to the default error file.
-    """
-
-    time: timedelta | Literal[0]
-    """
-    The maximum time for the job.
-
-    This corresponds to the "-t" option in sbatch. A value of 0 means no time limit.
-    """
-
-    memory_mb: int
-    """
-    The maximum memory for the job in MB.
-
-    This corresponds to the "--mem" option in sbatch. If not specified, the default memory limit will be used.
-    """
-
-    memory_per_cpu_mb: int
-    """
-    The minimum memory required per usable allocated CPU.
-
-    This corresponds to the "--mem-per-cpu" option in sbatch. If not specified, the default memory limit will be used.
-    """
-
-    memory_per_gpu_mb: int
-    """
-    The minimum memory required per allocated GPU.
-
-    This corresponds to the "--mem-per-gpu" option in sbatch. If not specified, the default memory limit will be used.
-    """
-
-    cpus_per_task: int
-    """
-    Advise the Slurm controller that ensuing job steps will require _ncpus_ number of processors per task.
-
-    This corresponds to the "-c" option in sbatch.
-    """
-
-    nodes: int
-    """
-    The number of nodes to use for the job.
-
-    This corresponds to the "-N" option in sbatch. The default is 1 node.
-    """
-
-    ntasks: int
-    """
-    The number of tasks to use for the job.
-
-    This corresponds to the "-n" option in sbatch. The default is 1 task.
-    """
-
-    ntasks_per_node: int
-    """
-    The number of tasks for each node.
-
-    This corresponds to the "--ntasks-per-node" option in sbatch.
-    """
-
-    constraint: str | Sequence[str]
-    """
-    Nodes can have features assigned to them by the Slurm administrator. Users can specify which of these features are required by their job using the constraint option.
-
-    This corresponds to the "-C" option in sbatch.
-    """
-
-    gres: str | Sequence[str]
-    """
-    Specifies a comma-delimited list of generic consumable resources.
-
-    This corresponds to the "--gres" option in sbatch.
-    """
-
-    gpus: int | str
-    """
-    Specify the total number of GPUs required for the job. An optional GPU type specification can be supplied.
-
-    This corresponds to the "-G" option in sbatch.
-    """
-
-    gpus_per_node: int | str
-    """
-    Specify the number of GPUs required for the job on each node included in the job's resource allocation. An optional GPU type specification can be supplied.
-
-    This corresponds to the "--gpus-per-node" option in sbatch.
-    """
-
-    gpus_per_task: int
-    """
-    Specify the number of GPUs required for the job on each task to be spawned in the job's resource allocation. An optional GPU type specification can be supplied.
-
-    This corresponds to the "--gpus-per-task" option in sbatch.
-    """
-
-    mail_user: str
-    """
-    User to receive email notification of state changes as defined by mail_type.
-
-    This corresponds to the "--mail-user" option in sbatch.
-    """
-
-    mail_type: MailType | Sequence[MailType]
-    """
-    Notify user by email when certain event types occur.
-
-    This corresponds to the "--mail-type" option in sbatch.
-    """
-
-    dependency: str
-    """
-    Defer the start of this job until the specified dependencies have been satisfied.
-
-    This corresponds to the "-d" option in sbatch.
-    """
-
-    exclusive: bool
-    """
-    The job allocation can not share nodes with other running jobs.
-
-    This corresponds to the "--exclusive" option in sbatch.
-    """
-
-    signal: signal.Signals
-    """
-    The signal to send to the job when the job is being terminated.
-
-    This corresponds to the "--signal" option in sbatch.
-    """
-
-    signal_delay: timedelta
-    """
-    The delay before sending the signal to the job.
-
-    This corresponds to the "--signal ...@[delay]" option in sbatch.
-    """
-
-    open_mode: str
-    """
-    The open mode for the output and error files.
-
-    This corresponds to the "--open-mode" option in sbatch.
-
-    Valid values are "append" and "truncate".
-    """
-
-    requeue: bool
-    """
-    Requeues the job if it's pre-empted.
-
-    This corresponds to the "--requeue" option in sbatch.
-    """
-
-    setup_commands: Sequence[str]
-    """
-    The setup commands to run before the job.
-
-    These commands will be executed prior to everything else in the job script.
-    """
-
-    environment: Mapping[str, str]
-    """
-    The environment variables to set for the job.
-
-    These variables will be set prior to executing any commands in the job script.
-    """
-
-    command_prefix: str
-    """
-    A command to prefix the job command with.
-
-    This is used to add commands like `srun` to the job command.
-    """
-
-    command_template: str
-    """
-    The template for the command to execute the helper script.
-
-    Default: `bash {/path/to/helper.sh}`.
-    """
-
-    srun_flags: str | Sequence[str]
-    """
-    The flags to pass to the `srun` command.
-    """
-
-
-DEFAULT_KWARGS: SlurmJobKwargs = {
-    "name": "ll",
-    # "nodes": 1,
-    # "time": timedelta(hours=2),
-    "signal": signal.SIGURG,
-    "signal_delay": timedelta(seconds=90),
-    "open_mode": "append",
-    # "requeue": True,
-}
-
-
-def _determine_gres(kwargs: SlurmJobKwargs) -> Sequence[str] | None:
-    """
-    There are many different ways to specify GPU resources, but some are buggy.
-
-    This function normalizes all other ways to specify GPU resources to the `gres` option.
-    """
-
-    # If `--gres` is set, just return it
-    if (gres := kwargs.get("gres")) is not None:
-        if isinstance(gres, str):
-            gres = [gres]
-        return gres
-
-    # We will only support `--gpus` if `--nodes` is set to 1
-    if (gpus := kwargs.get("gpus")) is not None:
-        if kwargs.get("nodes") != 1:
-            raise ValueError("Cannot specify `gpus` without `nodes` set to 1.")
-        if isinstance(gpus, int):
-            gpus = [f"gpu:{gpus}"]
-        return gpus
-
-    # `--gpus-per-task` is only supported if `--ntasks-per-node` is set (or can be inferred).
-    if (gpus_per_task := kwargs.get("gpus_per_task")) is not None:
-        if (ntasks_per_node := _determine_ntasks_per_node(kwargs)) is None:
-            raise ValueError(
-                "Cannot specify `gpus_per_task` without `ntasks_per_node`."
-            )
-
-        gpus_per_node = ntasks_per_node * gpus_per_task
-        return [f"gpu:{gpus_per_node}"]
-
-    # `--gpus-per-node` has no restrictions
-    if (gpus_per_node := kwargs.get("gpus_per_node")) is not None:
-        if isinstance(gpus_per_node, int):
-            gpus_per_node = [f"gpu:{gpus_per_node}"]
-        return gpus_per_node
-
-    return None
-
-
-def _determine_ntasks_per_node(kwargs: SlurmJobKwargs) -> int | None:
-    # If `--ntasks-per-node` is set, just return it
-    if (ntasks_per_node := kwargs.get("ntasks_per_node")) is not None:
-        return ntasks_per_node
-
-    # If `--ntasks` is set, we can infer `--ntasks-per-node`
-    if (ntasks := kwargs.get("ntasks")) is not None:
-        if (nodes := kwargs.get("nodes")) is None:
-            raise ValueError("Cannot infer `ntasks_per_node` without `nodes`.")
-
-        # If nnodes is not divisible by ntasks, raise an error
-        if nodes % ntasks != 0:
-            raise ValueError(
-                "The number of nodes must be divisible by the number of tasks."
-            )
-
-        return ntasks // nodes
-
-    return None
-
-
-def _write_batch_script_to_file(
-    path: Path,
-    kwargs: SlurmJobKwargs,
-    command: str,
-    job_array_n_jobs: int | None = None,
-):
-    with path.open("w") as f:
-        f.write("#!/bin/bash\n")
-
-        if kwargs.get("requeue"):
-            f.write("#SBATCH --requeue\n")
-
-        if job_array_n_jobs is not None:
-            f.write(f"#SBATCH --array=1-{job_array_n_jobs}\n")
-
-        if (name := kwargs.get("name")) is not None:
-            f.write(f"#SBATCH -J {name}\n")
-
-        if (account := kwargs.get("account")) is not None:
-            f.write(f"#SBATCH --account={account}\n")
-
-        if (time := kwargs.get("time")) is not None:
-            # A time limit of zero requests that no time limit be imposed. Acceptable time formats include "minutes", "minutes:seconds", "hours:minutes:seconds", "days-hours", "days-hours:minutes" and "days-hours:minutes:seconds".
-            if time == 0:
-                time_str = "0"
-            else:
-                total_seconds = time.total_seconds()
-                hours, remainder = divmod(total_seconds, 3600)
-                minutes, seconds = divmod(remainder, 60)
-                if hours > 24:
-                    days, hours = divmod(hours, 24)
-                    time_str = f"{int(days)}-{int(hours):02d}:{int(minutes):02d}"
-                else:
-                    time_str = f"{int(hours):02d}:{int(minutes):02d}:{int(seconds):02d}"
-            f.write(f"#SBATCH --time={time_str}\n")
-
-        if (nodes := kwargs.get("nodes")) is not None:
-            f.write(f"#SBATCH --nodes={nodes}\n")
-
-        if (ntasks := kwargs.get("ntasks")) is not None:
-            f.write(f"#SBATCH --ntasks={ntasks}\n")
-
-        if (ntasks_per_node := kwargs.get("ntasks_per_node")) is not None:
-            f.write(f"#SBATCH --ntasks-per-node={ntasks_per_node}\n")
-
-        if (output_file := kwargs.get("output_file")) is not None:
-            output_file = str(Path(output_file).absolute())
-            f.write(f"#SBATCH --output={output_file}\n")
-
-        if (error_file := kwargs.get("error_file")) is not None:
-            error_file = str(Path(error_file).absolute())
-            f.write(f"#SBATCH --error={error_file}\n")
-
-        if (partition := kwargs.get("partition")) is not None:
-            if isinstance(partition, str):
-                partition = [partition]
-            f.write(f"#SBATCH --partition={','.join(partition)}\n")
-
-        if (qos := kwargs.get("qos")) is not None:
-            f.write(f"#SBATCH --qos={qos}\n")
-
-        if (memory_mb := kwargs.get("memory_mb")) is not None:
-            f.write(f"#SBATCH --mem={memory_mb}\n")
-
-        if (memory_per_cpu_mb := kwargs.get("memory_per_cpu_mb")) is not None:
-            f.write(f"#SBATCH --mem-per-cpu={memory_per_cpu_mb}\n")
-
-        if (memory_per_gpu_mb := kwargs.get("memory_per_gpu_mb")) is not None:
-            f.write(f"#SBATCH --mem-per-gpu={memory_per_gpu_mb}\n")
-
-        if (cpus_per_task := kwargs.get("cpus_per_task")) is not None:
-            f.write(f"#SBATCH --cpus-per-task={cpus_per_task}\n")
-
-        if gres := _determine_gres(kwargs):
-            f.write(f"#SBATCH --gres={','.join(gres)}\n")
-
-        if (mail_user := kwargs.get("mail_user")) is not None:
-            f.write(f"#SBATCH --mail-user={mail_user}\n")
-
-        if (mail_type := kwargs.get("mail_type")) is not None:
-            if isinstance(mail_type, str):
-                mail_type = [mail_type]
-            f.write(f"#SBATCH --mail-type={','.join(mail_type)}\n")
-
-        if (dependency := kwargs.get("dependency")) is not None:
-            f.write(f"#SBATCH --dependency={dependency}\n")
-
-        if kwargs.get("exclusive"):
-            f.write("#SBATCH --exclusive\n")
-
-        if (open_mode := kwargs.get("open_mode")) is not None:
-            f.write(f"#SBATCH --open-mode={open_mode}\n")
-
-        if (constraint := kwargs.get("constraint")) is not None:
-            if isinstance(constraint, str):
-                constraint = [constraint]
-            f.write(f"#SBATCH --constraint={','.join(constraint)}\n")
-
-        if (signal := kwargs.get("signal")) is not None:
-            signal_str = signal.name
-            if (signal_delay := kwargs.get("signal_delay")) is not None:
-                signal_str += f"@{math.ceil(signal_delay.total_seconds())}"
-            f.write(f"#SBATCH --signal={signal_str}\n")
-
-        f.write("\n")
-
-        if (command_prefix := kwargs.get("command_prefix")) is not None:
-            command = " ".join(
-                x_stripped
-                for x in (command_prefix, command)
-                if (x_stripped := x.strip())
-            )
-        f.write(f"{command}\n")
-
-    return path
-
-
-def _update_kwargs(kwargs_in: SlurmJobKwargs, base_path: Path):
-    # Update the kwargs with the default values
-    kwargs = copy.deepcopy(DEFAULT_KWARGS)
-
-    # Merge the kwargs
-    kwargs = cast(SlurmJobKwargs, always_merger.merge(kwargs, kwargs_in))
-    del kwargs_in
-
-    # If out/err files are not specified, set them
-    logs_base = base_path.parent / "logs"
-    logs_base.mkdir(exist_ok=True)
-
-    if kwargs.get("output_file") is None:
-        kwargs["output_file"] = logs_base / "output_%j_%a.out"
-
-    if kwargs.get("error_file") is None:
-        kwargs["error_file"] = logs_base / "error_%j_%a.err"
-
-    # Update the command_prefix to add srun:
-    command_parts: list[str] = ["srun"]
-    if (srun_flags := kwargs.get("srun_flags")) is not None:
-        if isinstance(srun_flags, str):
-            srun_flags = [srun_flags]
-        command_parts.extend(srun_flags)
-
-    # Add ntasks/cpus/gpus
-    if (ntasks := kwargs.get("ntasks")) is not None:
-        command_parts.append(f"--ntasks={ntasks}")
-
-    if (ntasks_per_node := kwargs.get("ntasks_per_node")) is not None:
-        command_parts.append(f"--ntasks-per-node={ntasks_per_node}")
-
-    if (cpus_per_task := kwargs.get("cpus_per_task")) is not None:
-        command_parts.append(f"--cpus-per-task={cpus_per_task}")
-
-    if gres := _determine_gres(kwargs):
-        command_parts.append(f"--gres={','.join(gres)}")
-
-    command_parts.append("--unbuffered")
-
-    # Add the task id to the output filenames
-    if (f := kwargs.get("output_file")) is not None:
-        f = Path(f).absolute()
-        command_parts.extend(
-            [
-                "--output",
-                str(f.with_name(f"{f.stem}-%t{f.suffix}").absolute()),
-            ]
-        )
-    if (f := kwargs.get("error_file")) is not None:
-        f = Path(f).absolute()
-        command_parts.extend(
-            [
-                "--error",
-                str(f.with_name(f"{f.stem}-%t{f.suffix}").absolute()),
-            ]
-        )
-
-    # If there is already a command prefix, combine them.
-    if (existing_command_prefix := kwargs.get("command_prefix")) is not None:
-        command_parts.extend(existing_command_prefix.split())
-    # Add the command prefix to the kwargs.
-    kwargs["command_prefix"] = " ".join(command_parts)
-
-    return kwargs
-
-
-def to_array_batch_script(
-    dest: Path,
-    callable: Callable[[Unpack[TArgs]], Any],
-    args_list: Sequence[tuple[Unpack[TArgs]]],
-    /,
-    job_index_variable: str = "SLURM_ARRAY_TASK_ID",
-    print_environment_info: bool = False,
-    python_command_prefix: str | None = None,
-    **kwargs: Unpack[SlurmJobKwargs],
-) -> SubmitOutput:
-    """
-    Create the batch script for the job.
-    """
-
-    from ...picklerunner import serialize_many
-
-    kwargs = _update_kwargs(kwargs, dest)
-
-    # Convert the command/callable to a string for the command
-    num_jobs = len(args_list)
-
-    destdir = dest / "fns"
-    destdir.mkdir(exist_ok=True)
-
-    serialized_command = serialize_many(
-        destdir,
-        callable,
-        [(args, {}) for args in args_list],
-        start_idx=1,  # Slurm job indices are 1-based
-    )
-    helper_path = write_helper_script(
-        destdir,
-        serialized_command.to_bash_command(
-            job_index_variable, print_environment_info=print_environment_info
-        ),
-        kwargs.get("environment", {}),
-        kwargs.get("setup_commands", []),
-        command_prefix=python_command_prefix,
-    )
-    command = helper_script_to_command(helper_path, kwargs.get("command_template"))
-
-    script_path = _write_batch_script_to_file(
-        dest / "launch.sh",
-        kwargs,
-        command,
-        job_array_n_jobs=num_jobs,
-    )
-    script_path = script_path.resolve().absolute()
-    return SubmitOutput(
-        command_parts=["sbatch", f"{script_path}"],
-        script_path=script_path,
-    )