D1_09_Slurm_and_python.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "8230ab11-bec2-45e6-ba89-e34b06e7e621",
   "metadata": {},
   "source": [
    "# Slurm and python environments"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "39c4e360-20bf-4c08-b918-d11a0f0a70e9",
   "metadata": {},
   "source": [
    "To run python from Slurm batch scripts we have to take into account what we already learned about VSC and the options we have for python environments (Module, venv & conda).\n",
    "\n",
    "<div class=\"alert alert-info rounded-pill rounded-5\" style=\"margin: auto auto 10px auto; width:60rem; text-align: center;\">\n",
    "    <strong>No matter what approach we take: always check your assumptions and make sure that you are using the right packages and versions.</strong>\n",
    "</div>\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e6399443-294a-4d49-882a-3ae9b6b89a7d",
   "metadata": {},
   "source": [
    "# Usage with `module`\n",
    "\n",
    "When we had a look at the module system earlier we selected python 3.9, numpy and mpi4py packages.\n",
    "\n",
    "Lets take this example and run a small test script via slurm that uses these packages."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3fa7025c-9149-4ac5-a1e0-05e408f5d94f",
   "metadata": {},
   "outputs": [],
   "source": [
    "!module avail py-numpy/*p3dg2gd"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dff9ef94-fdfa-4821-ac53-c802d30bb023",
   "metadata": {},
   "outputs": [],
   "source": [
    "!module avail py-mpi4py/*xvabib2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1f9a58cd-3869-4ae9-98bf-fe980537fc08",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%writefile temp/slurm/job-python-module.sh\n",
    "#!/bin/bash\n",
    "\n",
    "#SBATCH --job-name=python-module\n",
    "#SBATCH --output=temp/slurm/slurm-python-module-%j.out\n",
    "\n",
    "#SBATCH --account=p70824                          # set the account to use (for billing)\n",
    "#SBATCH --qos=zen3_0512                           # qos\n",
    "#SBATCH --reservation=jh_training_python4hpc_1    # during training we have a fixed reservation\n",
    "\n",
    "#SBATCH --partition=zen3_0512       # hardware to use\n",
    "#SBATCH --time=00:05:00             # maximum time of 5 min for testing\n",
    "\n",
    "#SBATCH --ntasks=4                  # use 4 cpus\n",
    "#SBATCH --mem=2G                    # and 2G of memory\n",
    "\n",
    "# abort on bash errors\n",
    "set -e\n",
    "\n",
    "# load the python modules we need\n",
    "# take care of the order of loading if we load multiple modules\n",
    "module purge\n",
    "module load --auto py-mpi4py/3.1.3-gcc-12.2.0-xvabib2\n",
    "module load --auto py-numpy/1.24.3-gcc-12.2.0-p3dg2gd\n",
    "module list\n",
    "\n",
    "# print out which python executable is active and its version\n",
    "echo -e \"\\n>> Using: $( python3 -V ) from $( which python3 )\"\n",
    "\n",
    "# run a program in the allocation\n",
    "export PYTHONPATH=./tooling/:$PYTHONPATH\n",
    "python3 examples/mpi_numpy_test.py"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "45c2ae73-4458-4da5-adda-8358cad0952d",
   "metadata": {},
   "outputs": [],
   "source": [
    "!source tooling/unload_jupyter_env.sh && sbatch temp/slurm/job-python-module.sh"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "84bf7972-cd36-46f4-ae36-d79c4293680c",
   "metadata": {},
   "outputs": [],
   "source": [
    "!squeue --me"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "489419b4-98d2-4a4b-beb8-fe8fbac1ab04",
   "metadata": {},
   "source": [
    "Please note that when we are using module to load python packages the `pip` binary is not automatically available.\n",
    "\n",
    "Spack's 'python' package only provides the basic python interpreter and its library."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2e38cde7-85bb-4567-8067-eb3c7b2735dc",
   "metadata": {},
   "source": [
    "# Using simple python venvs"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "32a826c5-e41d-40c4-9d47-fd7d4ec29925",
   "metadata": {},
   "source": [
    "Instead of loading all python packages via `module` we just have to load the libraries we used for building and activate the virtual environment. Then we simply execute the python code.\n",
    "\n",
    "Different to the usage with the module system `pip` comes installed as a default here."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fcaa8b13-589a-4e73-a866-552539f2a20b",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%writefile temp/slurm/job-python-venv.sh\n",
    "#!/bin/bash\n",
    "\n",
    "#SBATCH --job-name=python-venv\n",
    "#SBATCH --output=temp/slurm/slurm-python-venv-%j.out\n",
    "\n",
    "#SBATCH --account=p70824                          # set the account to use (for billing)\n",
    "#SBATCH --qos=zen3_0512                           # qos\n",
    "#SBATCH --reservation=jh_training_python4hpc_1    # during training we have a fixed reservation\n",
    "\n",
    "#SBATCH --partition=zen3_0512       # hardware to use\n",
    "#SBATCH --time=00:05:00             # maximum time of 5 min for testing\n",
    "\n",
    "#SBATCH --ntasks=4                  # use 4 cpus\n",
    "#SBATCH --mem=2G                    # and 2G of memory\n",
    "\n",
    "# load the dependencies\n",
    "module load --auto gcc/13.2.0-gcc-12.2.0-wmf5yxk\n",
    "module load --auto intel-oneapi-mkl/2024.0.0-gcc-12.2.0-tk3clqd\n",
    "module load openmpi/4.1.6-gcc-12.2.0-exh7lqk\n",
    "\n",
    "# source the venv we want to use\n",
    "source temp/env/venv_mkl/bin/activate\n",
    "\n",
    "# run some commands for introspection\n",
    "echo \">> Pip: $( which pip )\"\n",
    "echo \">> Installed packages\"\n",
    "pip list --verbose\n",
    "echo \">> Python: $( python3 --version ) from $( which python3 )\"\n",
    "echo \">> Check expected numpy lapack linkage\"\n",
    "ldd temp/env/venv_mkl/lib/python3.11/site-packages/numpy/linalg/lapack_lite.cpython-311-x86_64-linux-gnu.so | grep \"intel\"\n",
    "\n",
    "# run a program in the allocation\n",
    "export PYTHONPATH=tooling/:$PYTHONPATH\n",
    "python3 examples/mpi_numpy_test.py"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cfac5499-1376-4a16-aef0-e3192447a9e8",
   "metadata": {},
   "outputs": [],
   "source": [
    "!source tooling/unload_jupyter_env.sh && sbatch temp/slurm/job-python-venv.sh"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c581871e-cd30-4dcc-8a7d-eeeb208e4c9e",
   "metadata": {},
   "outputs": [],
   "source": [
    "!squeue --me"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fffbed25-bd3a-40b3-9bfe-513be0ed7cc7",
   "metadata": {},
   "source": [
    "# Slurm and conda environments"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "add3725d-1da1-4572-b052-38d9c0dbb1d9",
   "metadata": {},
   "source": [
    "Similar to using a venv we also just have to load the conda environment and related modules we used for building.\n",
    "\n",
    "The situation gets trickier if we want to use external packages like OpenMPI. This will also be covered in more detail when we look into `mpi4py`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "21635901-bb2e-4ccc-8236-c8d5d1a21f09",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "%%writefile temp/slurm/job-python-conda.sh\n",
    "#!/bin/bash\n",
    "\n",
    "#SBATCH --job-name=python-conda\n",
    "#SBATCH --output=temp/slurm/slurm-python-conda-%j.out\n",
    "\n",
    "#SBATCH --account=p70824                          # set the account to use (for billing)\n",
    "#SBATCH --qos=zen3_0512                           # qos\n",
    "##SBATCH --reservation=jh_training_python4hpc_1    # during training we have a fixed reservation\n",
    "\n",
    "#SBATCH --partition=zen3_0512       # hardware to use\n",
    "#SBATCH --time=00:05:00             # maximum time of 5 min for testing\n",
    "\n",
    "#SBATCH --ntasks=4                  # use 4 cpus\n",
    "#SBATCH --mem=2G                    # and 2G of memory\n",
    "\n",
    "# load miniconda package and source shell functions\n",
    "module load miniconda3/latest\n",
    "eval \"$(conda shell.bash hook)\"\n",
    "# activate the environment to use (either by name or path)\n",
    "# here we use the previously generated conda-vsc-openmpi environment\n",
    "conda activate temp/env/conda-vsc-openmpi-env\n",
    "\n",
    "# load a VSC openmpi module and export the library path so applications can find the shared objects\n",
    "module load openmpi/4.1.6-gcc-12.2.0-exh7lqk\n",
    "# depending on the package configuration the LD_LIBRARY_PATH is not set to the lib folder\n",
    "# so we need to set it explicitly from LIBRARY_PATH thats usually used for compling code\n",
    "export LD_LIBRARY_PATH=\"$LIBRARY_PATH:$LD_LIBRARY_PATH\"\n",
    "\n",
    "# run some commands for introspection\n",
    "echo \">> Pip: $( which pip )\"\n",
    "echo \">> Installed packages\"\n",
    "pip list --verbose\n",
    "echo \">> Python: $( python3 --version ) from $( which python3 )\"\n",
    "\n",
    "# run a program in the allocation\n",
    "export PYTHONPATH=tooling/:$PYTHONPATH\n",
    "python3 examples/mpi_numpy_test.py"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "101cd61a-3cbc-4ee5-b733-6a3e0793a0bc",
   "metadata": {},
   "outputs": [],
   "source": [
    "!source tooling/unload_jupyter_env.sh && sbatch temp/slurm/job-python-conda.sh"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "31258c31-756f-4dcd-9863-440be8165211",
   "metadata": {},
   "outputs": [],
   "source": [
    "!squeue --me"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4d89d6f1-b016-4d52-83b9-c6e92ff1011e",
   "metadata": {},
   "source": [
    "# Slurm and Apptainer"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cf5504f7-e9c7-4f39-8692-6c5b814c07c5",
   "metadata": {},
   "source": [
    "Running Apptainer images from a Slurm batch script is pretty similar to using other environment options but we have a couple of more options regarding e.g. bindmounts & passthrough of hardware etc.\n",
    "\n",
    "In a nutshell we simply have to load an apptainer module and then execute the application or script.\n",
    "\n",
    "The application we execute does not necessarily need to be contained in the apptainer image since apptainer automatically mounts the users home directory. We simply use the container as executing environment.\n",
    "\n",
    "If we need other external directories these can also be mounted to be accessible from within the container."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9176ad5d-75a7-48d4-bd29-241899d4150d",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%writefile temp/slurm/job-python-apptainer.sh\n",
    "#!/bin/bash\n",
    "\n",
    "#SBATCH --job-name=python-apptainer\n",
    "#SBATCH --output=temp/slurm/slurm-python-apptainer-%j.out\n",
    "\n",
    "#SBATCH --account=p70824                          # set the account to use (for billing)\n",
    "#SBATCH --qos=zen3_0512                           # qos\n",
    "#SBATCH --reservation=jh_training_python4hpc_1    # during training we have a fixed reservation\n",
    "\n",
    "#SBATCH --partition=zen3_0512       # hardware to use\n",
    "#SBATCH --time=00:05:00             # maximum time of 5 min for testing\n",
    "\n",
    "#SBATCH --ntasks=4                  # use 4 cpus\n",
    "#SBATCH --mem=2G                    # and 2G of memory\n",
    "\n",
    "# abort on bash errors\n",
    "set -e\n",
    "\n",
    "# load apptainer\n",
    "module purge\n",
    "module load --auto apptainer/1.1.6-gcc-12.2.0-xxfuqni\n",
    "\n",
    "# in order to make openmpi work with the container we have to load the module\n",
    "module load openmpi/4.1.6-gcc-12.2.0-exh7lqk\n",
    "# add the LIBRARY_PATH that was set by the modules to the LD_LIBRARY_PATH in the container\n",
    "export APPTAINERENV_LD_LIBRARY_PATH=\"$LIBRARY_PATH\"\n",
    "# and bind the software package tree as well as the system's gpfs library (!hacky!)\n",
    "export APPTAINER_BIND=\"/gpfs/opt/sw/:/gpfs/opt/sw/,/lib64/libgpfs.so:/lib/libgpfs.so\"\n",
    "\n",
    "# run a program using the resource allocation via apptainer image\n",
    "# the 'example' path does not need to be present in the image since apptainer\n",
    "# automatically mounts the home directory of the executing user\n",
    "export PYTHONPATH=tooling/:$PYTHONPATH\n",
    "apptainer exec temp/env/apptainer-mpi.sif python3 ./examples/mpi_numpy_test.py"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "61116363-d966-416d-be8d-228e548d71f2",
   "metadata": {},
   "outputs": [],
   "source": [
    "!source tooling/unload_jupyter_env.sh && sbatch temp/slurm/job-python-apptainer.sh"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "03b4376b-1301-4c01-80f1-cf0f549e2a22",
   "metadata": {},
   "outputs": [],
   "source": [
    "!squeue --me"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9fede467-289a-4035-b296-0a081b34b2d1",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}