chaseking / slurmexec

Lightweight Python library for creating scripts easily executable on Slurm.

Geek Repo:Geek Repo

Github PK Tool:Github PK Tool

SlurmExec

Lightweight Python library to easily deploy and execute Slurm tasks. Fast and easy; no dependencies needed.

Features

  • Define a job function in Python and call slurm_exec in the python file __main__
  • Dynamically-created slurm scripts avoid the headache of writing and manually keeping track of messy slurm job files
  • Control function arguments from command line

Install

Option 1 (preferred): Directly:

# Install
pip install git+https://github.com/chaseking/slurmexec.git

Note: If you use this method you will need to run

# Upgrade
pip uninstall -y slurmexec && pip install git+https://github.com/chaseking/slurmexec

to upgrade this package. Eventually I will make this available on PyPI to make this process easier.

Option 2: Clone and install

git clone git@github.com:chaseking/slurmexec.git
cd slurmexec
pip install -e .

Example

from slurmexec import *

# set_slurm_debug()  # to run locally

if is_this_a_slurm_job():
    # Conditional imports to avoid memory
    # crashing (OOM) Slurm login node
    import biglibrary

@slurm_job
def hello(message: str = "World"):
    print(f"Hello {message}!")
    
    biglibrary.expensive_function()
    
    print(f"Finished slurm job {get_slurm_id()}")

if __name__ == "__main__":
    slurm_exec(
        func = hello,
        slurm_args = {
            "--gres": "gpu:1",
        }
        pre_run_commands = [
            "conda activate myenv"
        ]
    )
# Executed on a slurm login node
python myfile.py --message 'big beautiful world'

About

Lightweight Python library for creating scripts easily executable on Slurm.

License:MIT License


Languages

Language:Python 100.0%