languages/python/Python_on_NREL_HPC.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "![NREL Logo](https://www.nrel.gov/_resources/images/nrel-logo.png)\n",
    "\n",
    "## **Note: _This notebook was designed to be executed from Eagle_**"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Python on NREL HPC\n",
    "By design, the HPC is a time-shared multi-machine system which necessarily warrants some nuanced consideration for how environments are managed relative to a single machine with a single user. Sometimes, the default workflow for environment creation and usage is not the most optimal for some use-cases.\n",
    "\n",
    "Below is a list of common pitfalls that users have encountered historically while using Python and Anaconda on NREL HPC. \n",
    "\n",
    "* Running a SLURM job that uses a `conda` environment which is stored in `$HOME`.\n",
    "* Exhausting the `$HOME` storage quota (50GB on the current HPC system) usually because of conda's package cache combined with their user environments.\n",
    "* Trying to share a `conda` environment from another user's `/home` directory.\n",
    "* Forgetting to install `jupyter` in a new conda environment, resulting in using the `base` installation's version which doesn't have your dependencies installed.\n",
    "\n",
    "Let's discuss strategies to mitigate or avoid these kinds of problems"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Installing Conda Environments in Different Directories\n",
    "By default, `conda` will install new environments in `$HOME/.conda`. Generally speaking, this a sensible default&mdash;it just happens to be the starting point to frequent issues that users have experienced historically. Something to consider is that `conda` has a `--prefix` flag which allows one to arbitrate where a conda environment gets installed to, notably allowing you to place environments on other file-systems and block devices besides the `/home` network-storage that is mounted on NREL HPC systems.\n",
    "\n",
    "For example, here is how one might create a project in their `/scratch` directory:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 144,
   "metadata": {},
   "outputs": [],
   "source": [
    "ENV_PREFIX=\"/scratch/$USER/demo_scratch_env\"\n",
    "\n",
    "import os ; os.environ['ENV_PREFIX']=ENV_PREFIX  # Export this variable for cells below "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 148,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Collecting package metadata: ...working... done\n",
      "Solving environment: ...working... done\n",
      "\n",
      "## Package Plan ##\n",
      "\n",
      "  environment location: /scratch/mbartlet/demo_scratch_env\n",
      "\n",
      "\n",
      "\n",
      "Preparing transaction: ...working... done\n",
      "Verifying transaction: ...working... done\n",
      "Executing transaction: ...working... done\n"
     ]
    }
   ],
   "source": [
    "!conda create --quiet --use-local --yes \\\n",
    "    --prefix $ENV_PREFIX   # `--prefix` in action \\\n",
    "    python=3.7"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 146,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "drwxr-xr-x. 3 mbartlet mbartlet 4096 Dec  3 11:10 /scratch/mbartlet/demo_scratch_env\r\n"
     ]
    }
   ],
   "source": [
    "!ls -ld $ENV_PREFIX"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 147,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Delete the demo environment for cleanliness\n",
    "!conda-env remove --yes --quiet --prefix $ENV_PREFIX &>/dev/null"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Below is a table which discusses the pros and cons of each block-device mount on NREL HPC as a location for storing your software environments.\n",
    "\n",
    "| Block-device mounts | Situations where you would want to use this block device for your conda environments                                                                                                                                                  | Caveats to consider when using this mount                                                                                                                                                                                                                                                                                                                                           |\n",
    "|---------------------------|---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------|--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------|\n",
    "| /home                     | `$HOME/.conda` is the default location for environments. For one-off environments,  or if you don't create environments often, this is a reasonable location for your environments and doesn't require any extra flags or parameters. | Files in <span>\\$</span>HOME will not be purged so long as you have an active NREL HPC account. However, <span>\\$</span>HOME is limited to a 50GB storage quota so you may have to take care to monitor your storage footprint.                                                                                                                                                                                        |\n",
    "| /scratch                  | `/scratch` or `/projects` is ultimately where you want your environment to end up if your jobs have more than 1 node&mdash;if your environment is in `/home` then every node in your job will be competing for read-access over a non-parallel network fabric to the source files of your environment. `/scratch` provides simultaneous access to all the nodes. A sensible approach is copying your environments from `/home` to `/scratch` as part of your job's initialization. | `/scratch` storage is unlimited. `/scratch` is a parallel filesystem, meaning simultaneous filesystem operations by several nodes is possible and performant. However, the contents of `/scratch` are subject to purge after 28 days of inactivity.                                                                                                                                        |\n",
    "| /projects                 | This is a great place to put a conda environment that you anticipate sharing with your colleagues who are also working on the project. You can structure the permissions such that others in the project have read-only, write-only, or no access (we also encourage restoring these permissions at a later date so others on the project can manage your files without a hassle). | `/projects` is also a parallel filesystem which reaps the same benefits as mentioned above. However, access to projects is contingent on having access to an HPC project allocation. Moreover, the storage quota allotted to each project is relative to the reasonableness of its requested needs, although a conda environment is very unlikely to have a significant storage footprint. |"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As mentioned above, let's demonstrate one might go about copying an environment from `/home` to `/scratch` in a SLURM job. The below cell will generate a nice code block based on variables used earlier in this notebook, as well as environment variables within your user account:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 149,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Acquire a default project handle to procedurally generate a SLURM job\n",
    "import subprocess\n",
    "\n",
    "command = \"/nopt/nrel/utils/bin/hours_report | tail -1 | awk '{print $1}'\" # Grab a valid project handle\n",
    "\n",
    "command_array = [\n",
    "    '/bin/bash',\n",
    "    '-c',\n",
    "    command\n",
    "]\n",
    "\n",
    "project_handle = subprocess.run(command_array, stdout=subprocess.PIPE).stdout.decode('utf-8')[:-1]\n",
    "\n",
    "import os ; os.environ['DEFAULT_HANDLE'] = project_handle  # Export handle for cells below"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 150,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "wks\r\n"
     ]
    }
   ],
   "source": [
    "!echo $DEFAULT_HANDLE"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 152,
   "metadata": {},
   "outputs": [],
   "source": [
    "conda_home_env=\"py3\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 153,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Acquire info about the default conda environment\n",
    "import subprocess\n",
    "\n",
    "command = f\"module load conda && . activate {conda_home_env} && echo $CONDA_PREFIX\"\n",
    "\n",
    "command_array = [\n",
    "    '/bin/bash',\n",
    "    '-lc',         # Have to run this from a login-shell\n",
    "    command\n",
    "]\n",
    "\n",
    "conda_home_env_prefix = subprocess.run(command_array, stdout=subprocess.PIPE).stdout.decode('utf-8')[:-1]\n",
    "\n",
    "import os ; os.environ['CONDA_HOME_ENV_PREFIX'] = conda_home_env_prefix  # Export handle for cells below"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 154,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "/home/mbartlet/.conda/envs/py3\r\n"
     ]
    }
   ],
   "source": [
    "!echo $CONDA_HOME_ENV_PREFIX"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 155,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/markdown": [
       "\n",
       "```bash\n",
       "#!/usr/bin/env bash\n",
       "#SBATCH --account wks\n",
       "#SBATCH --time 5\n",
       "#SBATCH --partition debug\n",
       "#SBATCH --nodes 2\n",
       "\n",
       "export SCRATCH_ENV=\"/scratch/mbartlet/home_conda_clone\"\n",
       "rsync -avz --ignore-existing \"/home/mbartlet/.conda/envs/py3\" \"$SCRATCH_ENV\" &>/dev/null\n",
       "\n",
       "srun bash -l <<EOF\n",
       "module purge\n",
       "module load conda\n",
       ". activate \"$SCRATCH_ENV\"\n",
       "which python\n",
       "EOF\n",
       "\n",
       "rm -rf \"$SCRATCH_ENV\"  # Optional clean-up\n",
       "```\n"
      ],
      "text/plain": [
       "<IPython.core.display.Markdown object>"
      ]
     },
     "execution_count": 155,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from IPython.display import Markdown as md\n",
    "from os import environ as env\n",
    "\n",
    "SCRATCH_ENV=f\"/scratch/{env['USER']}/home_conda_clone\"\n",
    "\n",
    "body=f\"\"\"\n",
    "```bash\n",
    "#!/usr/bin/env bash\n",
    "#SBATCH --account {env['DEFAULT_HANDLE']}\n",
    "#SBATCH --time 5\n",
    "#SBATCH --partition debug\n",
    "#SBATCH --nodes 2\n",
    "\n",
    "export SCRATCH_ENV=\"{SCRATCH_ENV}\"\n",
    "rsync -avz --ignore-existing \"{env['CONDA_HOME_ENV_PREFIX']}\" \"$SCRATCH_ENV\" &>/dev/null\n",
    "\n",
    "srun bash -l <<EOF\n",
    "module purge\n",
    "module load conda\n",
    ". activate \"$SCRATCH_ENV\"\n",
    "which python\n",
    "EOF\n",
    "\n",
    "rm -rf \"$SCRATCH_ENV\"  # Optional clean-up\n",
    "```\n",
    "\"\"\"\n",
    "\n",
    "md(body)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And after running what was generated above:\n",
    "```bash\n",
    "[mbartlet@el1 ~] $ cat slurm-1845968.out\n",
    "/scratch/mbartlet/home_conda_clone/bin/python\n",
    "/scratch/mbartlet/home_conda_clone/bin/python\n",
    "```\n",
    "Which shows both nodes sourced the environment from `/scratch`"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "py3",
   "language": "python",
   "name": "py3"
  },
  "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.7.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}