flow_simulation/flow_simulation.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Rainwater Flow Simulation\n",
    "\n",
    "In this notebook, we're going to use MPI to evaluate a range of parameters for\n",
    "a linear rainfall-runoff model in parallel on a Banyan computing session. Banyan\n",
    "automatically sets up cloud computing resources in your own Virtual Private Cloud,\n",
    "sets up MPI and Python, manages your software environment, lets you run with any\n",
    "number of workers, and scales down when you're finished. All your estimated costs,\n",
    "running clusters, sessions, and logs can be viewed on\n",
    "the [Banyan dashboard](https://www.banyancomputing.com/dashboard).\n",
    "\n",
    "To run this notebook, you will need to have set up an account with Banyan. Check\n",
    "out the documentation [here](https://www.banyancomputing.com/getting-started/)\n",
    "for step by step instructions."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Configuring\n",
    "\n",
    "To run this example, please ensure that you have set up your Banyan account.\n",
    "\n",
    "Run the first cell below to import `banyan`.\n",
    "\n",
    "To configure your AWS credentials, run the second cell below and provide your\n",
    "AWS credentials when prompted. Banyan does not save your AWS credentials, but\n",
    "they are needed so that you can run your computation in your AWS account.\n",
    "Finally, run the third cell below to set your Banyan credentials and configure\n",
    "Banyan.\n",
    "\n",
    "You must pass your User ID and API Key to the `configure` function in order\n",
    "to authenticate. You can find this information on the Account page of the\n",
    "Banyan Dashboard. After running this cell, your credentials will be saved\n",
    "in `$HOME/.banyan/banyanconfig.toml` and will be read from that file in the\n",
    "future. This means that you only need to run this cell once."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Import packages\n",
    "import banyan\n",
    "from banyan import configure\n",
    "from banyan import create_cluster, destroy_cluster, get_cluster, get_clusters\n",
    "from banyan import run_session, end_session\n",
    "\n",
    "import getpass\n",
    "import urllib.request"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Run this cell to configure the AWS CLI. When prompted, specify your AWS\n",
    "# credentials for the AWS account that you connected with Banyan. If you have\n",
    "# already configured the AWS CLI with the credentials for the account you have\n",
    "# configured with your Banyan account, you can skip this step.\n",
    "\n",
    "import os\n",
    "\n",
    "os.environ[\"AWS_ACCESS_KEY_ID\"] = getpass.getpass(prompt=\"Enter AWS_ACCESS_KEY_ID\\n\")\n",
    "os.environ[\"AWS_SECRET_ACCESS_KEY\"] = getpass.getpass(prompt=\"Enter AWS_SECRET_ACCESS_KEY\\n\")\n",
    "os.environ[\"AWS_DEFAULT_REGION\"] = getpass.getpass(prompt=\"Enter AWS_DEFAULT_REGION\\n\")\n",
    "\n",
    "print(\"AWS is now configured.\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Run this cell to configure Banyan. When prompted, provide your user ID and API\n",
    "# key. You can find these on the Account page of your Banyan dashboard.\n",
    "# If you have already configured Banyan, you can skip this step.\n",
    "\n",
    "user_id = getpass.getpass(prompt=\"Please enter your User ID\\n\")\n",
    "api_key = getpass.getpass(prompt=\"Please enter your API Key\\n\")\n",
    "\n",
    "# Configures Banyan client library with your Banyan credentials\n",
    "configure(user_id=user_id, api_key=api_key)\n",
    "print(\"Banyan is now configured.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Creating a cluster\n",
    "\n",
    "For this example, you will need a Banyan cluster. You can either use an existing\n",
    "cluster or create a new cluster. Run the following code block and enter in either\n",
    "the name of an existing cluster or the name you would like to use for a new cluster.\n",
    "\n",
    "If you already have a cluster, you should specify its name, when prompted.\n",
    "\n",
    "If you would like to instead create a new cluster, provide a name and the name\n",
    "of the [Amazon EC2 key pair](https://docs.aws.amazon.com/AWSEC2/latest/UserGuide/ec2-key-pairs.html#having-ec2-create-your-key-pair) that you created during [Banyan setup](https://www.banyancomputing.com/creating-clusters).\n",
    "\n",
    "In the cell below, you can change `instance_type` to create a cluster with a\n",
    "different EC2 instance type that may have a larger amount of memory or workers.\n",
    "See the documentation [here](https://www.banyancomputing.com/banyan-py-docs/create-cluster/) for the other parameters for creating a cluster."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cluster_name = input(\"Cluster name for existing cluster or new cluster \")\n",
    "print(f\"Checking if cluster {cluster_name} exists...\")\n",
    "\n",
    "clusters = get_clusters()\n",
    "print(f\"You have {len(clusters)} clusters\")\n",
    "if not ((cluster_name in clusters) and (clusters[cluster_name].status == \"running\")):\n",
    "    print(f\"Creating new cluster {cluster_name}\")\n",
    "    ec2_key_pair_name = getpass.getpass(prompt=\"Name of SSH EC2 Key Pair\")\n",
    "    print(f\"Using EC2 key pair {ec2_key_pair_name}\")\n",
    "    create_cluster(\n",
    "        name=cluster_name,\n",
    "        instance_type=\"t3.2xlarge\",\n",
    "        initial_num_workers=2,\n",
    "        ec2_key_pair_name=ec2_key_pair_name\n",
    "    )\n",
    "else:\n",
    "    print(f\"Using existing cluster {cluster_name}\")\n",
    "get_cluster(cluster_name)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Running the Simulation\n",
    "\n",
    "We have a script in `simulation.py` which uses mpi4py (the commonly-used Python package that provides MPI bindings). This script runs multiple linear rainfall-runoff models using different parameters in parallel on different workers. To run this script on a session of multiple workers running in parallel, we call `run_session` and pass in the names of one or more `.py` files into the `code_files` to run."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Download a file required for the simulation\n",
    "urllib.request.urlretrieve(\"https://gist.githubusercontent.com/jdherman/252be34ced79dc42dc2300f227c2af29/raw/5d58fff0c47a6c1f78f0523488e4b946af507338/leaf-river-data.txt\", \"leaf-river-data.txt\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "run_session(\n",
    "    cluster_name=cluster_name, \n",
    "    session_name=\"rainfall_model\", \n",
    "    nworkers=2,\n",
    "    files=[\n",
    "        \"file://stockflow.py\",\n",
    "        \"leaf-river-data.txt\"\n",
    "    ],\n",
    "    code_files=[\"file://run_simulation.py\"],\n",
    "    print_logs=True\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Clean up resources\n",
    "\n",
    "To destroy the cluster entirely (note that it will take 10-30 min to recreate the cluster), run the following:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Destroy the cluster\n",
    "destroy_cluster(cluster_name)"
   ]
  }
 ],
 "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.8.7"
  },
  "vscode": {
   "interpreter": {
    "hash": "30bc4c5ed26e9abad4968d61b1e09a864ca213f653605935e7191f624fc541fd"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}