[Gautzilla]

🍹Slurp the SLURM

This PR aims at adding the osekit.job.scheduler.slurm class, which can be used to submit jobs in a SLURM scheduler.

It'll be developed alongside #465, as some of the code in #465 is still PBS-locked. I'll do some edits in #465 when I feel like some features that I thought were PBS-specific can be moved up in the Scheduler class, and I'll rebase this branch on the Scheduler one from time to time.

In the end, the user should be able to run public transforms in a slurm workload manager easily:

from osekit.job.builder import JobBuilder
from osekit.job.config import JobConfig
from osekit.job.scheduler.slurm import Slurm
from osekit.public.project import Project

project = Project(...)  # See the Project documentation

job_config = JobConfig(...)  # See the JobConfig documentation

scheduler = Slurm(partition="cpu")  # Scheduler in which the job is submitted

project.job_builder = JobBuilder(
    config=job_config,
    scheduler=scheduler,
)

# Now the dataset has a non-None job_builder attribute,
# running a transform will write a SLURM file in the logs directory
# and submit it through the selected scheduler.

project.run(...)  # See the Transform documentation

🐕‍🦺 Shebang config and subprocess prefix

I had some issues when trying to test the Slurm scheduler in situ:

• The shebang of the slurm files needed to specify the login mode (with an appended -l parameter) of the bash interpreter, otherwise the module command wasn't found
• The terminal kernels spawned on jupyterhub didn't know the path of the slurm commands: an explicit ssh $HOSTNAME command was needed before running slurm commands

Shebang config

The shebang can now be edited in the Scheduler.shebang property. It is defaulted to #!/bin/bash, so the login mode can be specified by:

from osekit.job.scheduler.slurm import Slurm

scheduler = Slurm(...)
scheduler.shebang += r" -l"

Subprocess prefix

I added the Scheduler.cmd_prefix property that contains a list of commands that are run before running the actual commands.

To automatically run the ssh $HOSTNAME command before running any slurm command (that are run by OSEkit under the hood), we have to get the HOSTNAME through the environment variables and add the ssh command to the cmd_prefix property:

import os
from osekit.job.scheduler.slurm import Slurm

hostname = os.environ["HOSTNAME"]

scheduler = Slurm(...)
scheduler.cmd_prefix = ["ssh", hostname]

Full Slurm job config example

import os

from pandas import Timedelta

from osekit.job.builder import JobBuilder
from osekit.job.config import JobConfig
from osekit.job.scheduler.slurm import Slurm

hostname = os.environ["HOSTNAME"]

job_config = JobConfig(
    nb_nodes=1,  # Number of nodes on which the job runs
    ncpus=28,  # Number of total cores used per node
    mem="60gb",  # Maximum amount of physical memory used by the job
    walltime=Timedelta(
        hours=5,
    ),  # Maximum amount of real itime during which the job can be running
    venv_name=os.environ["CONDA_DEFAULT_ENV"],  # Works only for conda venvs
)

scheduler = Slurm()  # Scheduler in which the job is submitted
scheduler.cmd_prefix = ["ssh", hostname]
scheduler.shebang += r" -l"  # Foces login mode on slurm job files

job_builder = JobBuilder(
    config=job_config,
    scheduler=scheduler,
)

resolves #279

[coveralls]

Coverage is 98.957% — Gautzilla:slurm into Project-OSmOSE:main. No base build found for Project-OSmOSE:main.