Slurm config

class SlurmConfig(BaseModel):
    """
    Abstraction for SLURM parameters

    **NOTE**: `SlurmConfig` objects are created internally in `fractal-server`,
    and they are not meant to be initialized by the user; the same holds for
    `SlurmConfig` attributes (e.g. `mem_per_task_MB`), which are not meant to
    be part of the superuser-defined `resource.jobs_runner_config` JSON field.

    Part of the attributes map directly to some of the SLURM attributes (see
    https://slurm.schedmd.com/sbatch.html), e.g. `partition`. Other attributes
    are metaparameters which are needed in fractal-server to combine multiple
    tasks in the same SLURM job (e.g. `parallel_tasks_per_job` or
    `max_num_jobs`).

    Attributes:
        partition: Corresponds to SLURM option.
        cpus_per_task: Corresponds to SLURM option.
        mem_per_task_MB: Corresponds to `mem` SLURM option.
        job_name: Corresponds to `name` SLURM option.
        constraint: Corresponds to SLURM option.
        gres: Corresponds to SLURM option.
        account: Corresponds to SLURM option.
        gpus: Corresponds to SLURM option.
        time: Corresponds to SLURM option (WARNING: not fully supported).
        nodelist: Corresponds to SLURM option.
        nodes: Corresponds to SLURM option.
        exclude: Corresponds to SLURM option.
        prefix: Prefix of configuration lines in SLURM submission scripts.
        shebang_line: Shebang line for SLURM submission scripts.
        use_mem_per_cpu:
            If `True`, use `--mem-per-cpu` rather than `--mem`, both at the job
            level and for `srun` statements.
        extra_lines: Additional lines to include in SLURM submission scripts.
        tasks_per_job: Number of tasks for each SLURM job.
        parallel_tasks_per_job: Number of tasks to run in parallel for
                                each SLURM job.
        target_cpus_per_job: Optimal number of CPUs to be requested in each
                             SLURM job.
        max_cpus_per_job: Maximum number of CPUs that can be requested in each
                          SLURM job.
        target_mem_per_job: Optimal amount of memory (in MB) to be requested in
                            each SLURM job.
        max_mem_per_job: Maximum amount of memory (in MB) that can be requested
                         in each SLURM job.
        target_num_jobs: Optimal number of SLURM jobs for a given WorkflowTask.
        max_num_jobs: Maximum number of SLURM jobs for a given WorkflowTask.
        user_local_exports:
            Key-value pairs to be included as `export`-ed variables in SLURM
            submission script, after prepending values with the user's cache
            directory.

    """

    model_config = ConfigDict(extra="forbid")

    # Required SLURM parameters (note that the integer attributes are those
    # that will need to scale up with the number of parallel tasks per job)
    partition: str
    cpus_per_task: int
    mem_per_task_MB: int

    prefix: str = "#SBATCH"
    shebang_line: str = "#!/bin/sh"

    use_mem_per_cpu: bool = False

    # Optional SLURM parameters
    job_name: str | None = None
    constraint: str | None = None
    gres: str | None = None
    gpus: str | None = None
    time: str | None = None
    account: str | None = None
    nodelist: str | None = None
    nodes: int | None = None
    exclude: str | None = None

    # Free-field attribute for extra lines to be added to the SLURM job
    # preamble
    extra_lines: list[str] = Field(default_factory=list)

    # Variables that will be `export`ed in the SLURM submission script
    user_local_exports: dict[str, str] = Field(default_factory=dict)

    # Metaparameters needed to combine multiple tasks in each SLURM job
    tasks_per_job: int | None = None
    parallel_tasks_per_job: int | None = None
    target_cpus_per_job: int
    max_cpus_per_job: int
    target_mem_per_job: int
    max_mem_per_job: int
    target_num_jobs: int
    max_num_jobs: int

    def _sorted_extra_lines(self: Self) -> list[str]:
        """
        Return a copy of `self.extra_lines`, where lines starting with
        `self.prefix` are listed first.
        """

        def _no_prefix(_line) -> int:
            if _line.startswith(self.prefix):
                return 0
            else:
                return 1

        return sorted(self.extra_lines, key=_no_prefix)

    def sort_script_lines(self: Self, script_lines: list[str]) -> list[str]:
        """
        Return a copy of `script_lines`, where lines are sorted as in:

        1. `self.shebang_line` (if present);
        2. Lines starting with `self.prefix`;
        3. Other lines.

        Args:
            script_lines:
        """

        def _sorting_function(_line) -> int:
            if _line == self.shebang_line:
                return 0
            elif _line.startswith(self.prefix):
                return 1
            else:
                return 2

        return sorted(script_lines, key=_sorting_function)

    def to_sbatch_preamble(
        self: Self,
        remote_export_dir: str,
        use_mem_per_cpu: bool = False,
    ) -> list[str]:
        """
        Compile `SlurmConfig` object into the preamble of a SLURM submission
        script.

        Args:
            remote_export_dir:
                Base directory for exports defined in
                `self.user_local_exports`.
            use_mem_per_cpu:
                If `True`, use `--mem-per-cpu` rather than `--mem`.
        """
        if self.parallel_tasks_per_job is None:
            raise ValueError(
                "SlurmConfig.sbatch_preamble requires that "
                f"{self.parallel_tasks_per_job=} is not None."
            )
        if len(self.extra_lines) != len(set(self.extra_lines)):
            raise ValueError(f"{self.extra_lines=} contains repetitions")

        if use_mem_per_cpu:
            memory_line = f"{self.prefix} --mem-per-cpu={self.mem_per_cpu_MB}M"
        else:
            mem_per_job_MB = self.parallel_tasks_per_job * self.mem_per_task_MB
            memory_line = f"{self.prefix} --mem={mem_per_job_MB}M"

        lines = [
            self.shebang_line,
            f"{self.prefix} --partition={self.partition}",
            f"{self.prefix} --ntasks={self.parallel_tasks_per_job}",
            f"{self.prefix} --cpus-per-task={self.cpus_per_task}",
            memory_line,
        ]
        for key in [
            "job_name",
            "constraint",
            "gres",
            "gpus",
            "time",
            "account",
            "exclude",
            "nodelist",
            "nodes",
        ]:
            value = getattr(self, key)
            if value is not None:
                # Handle the `time` parameter
                if key == "time" and self.parallel_tasks_per_job > 1:
                    # NOTE: see issue #1632
                    logger.warning(
                        f"`time` SLURM parameter is set to {self.time}, "
                        "but this does not take into account the number of "
                        f"SLURM tasks ({self.parallel_tasks_per_job})."
                    )
                option = key.replace("_", "-")
                lines.append(f"{self.prefix} --{option}={value}")

        for line in self._sorted_extra_lines():
            lines.append(line)

        if self.user_local_exports:
            for key, value in self.user_local_exports.items():
                tmp_value = str(Path(remote_export_dir) / value)
                lines.append(f"export {key}={tmp_value}")

        """
        FIXME export SRUN_CPUS_PER_TASK
        # From https://slurm.schedmd.com/sbatch.html: Beginning with 22.05,
        # srun will not inherit the --cpus-per-task value requested by salloc
        # or sbatch.  It must be requested again with the call to srun or set
        # with the SRUN_CPUS_PER_TASK environment variable if desired for the
        # task(s).
        if config.cpus_per_task:
            #additional_setup_lines.append(
                f"export SRUN_CPUS_PER_TASK={config.cpus_per_task}"
            )
        """

        return lines

    @property
    def batch_size_or_zero(self: Self) -> int:
        """
        Equivalent to `JobRunnerConfigLocal.batch_size_or_zero`.

        A `0` value is replaced with `tot_tasks` within
        `enrich_task_files_multisubmit`.
        """
        return self.tasks_per_job or 0

    @property
    def batch_size_or_one(self: Self) -> int:
        """
        Safe to use e.g. in `range(0, tot, batch_size_or_one)`.
        """
        return self.tasks_per_job or 1

    @property
    def mem_per_cpu_MB(self: Self) -> int:
        return int(self.mem_per_task_MB / self.cpus_per_task)