collab update

Author: aterrel

notebooks/8__mpi.ipynb

@@ -9,22 +9,14 @@
    },
    "outputs": [],
    "source": [
-    "## Uncomment and Execute this cell if running on slurm based cluster\n",
-    "SRUN=\"srun\""
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "id": "0sW2XPMR-687",
-   "metadata": {
-    "id": "0sW2XPMR-687"
-   },
-   "outputs": [],
-   "source": [
-    "## Uncomment and Execute this cell if running on google collab\n",
-    "#!pip install mpi4py\n",
-    "#SRUN=\"mpirun --allow-run-as-root --oversubscribe\""
+    "## Setup on \n",
+    "MPIRUN=\"srun\"\n",
+    "\n",
+    "import os\n",
+    "\n",
+    "if os.getenv(\"COLAB_RELEASE_TAG\"): # If running in Google Colab\n",
+    "    !pip install mpi4py\n",
+    "    MPIRUN=\"mpirun --allow-run-as-root --oversubscribe\"\n"
    ]
   },
   {
@@ -52,7 +44,7 @@
    },
    "outputs": [],
    "source": [
-    "!$SRUN -n 2 -c 2 python hello_world.py"
+    "!$MPIRUN -n 2 -c 2 python hello_world.py"
    ]
   },
   {
@@ -64,7 +56,7 @@
    },
    "outputs": [],
    "source": [
-    "!$SRUN  -n 4 -c 4 python hello_world.py"
+    "!$MPIRUN  -n 4 -c 4 python hello_world.py"
    ]
   },
   {
@@ -97,7 +89,7 @@
    },
    "outputs": [],
    "source": [
-    "!$SRUN  -n 2 -c 2 python hello_world_with_ranks.py"
+    "!$MPIRUN  -n 2 -c 2 python hello_world_with_ranks.py"
    ]
   },
   {
@@ -128,7 +120,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "!$SRUN  -n 4 -c 4 python fizzbuzz.py"
+    "!$MPIRUN  -n 4 -c 4 python fizzbuzz.py"
    ]
   },
   {
@@ -164,7 +156,7 @@
    },
    "outputs": [],
    "source": [
-    "!$SRUN  -n 2 -c 2 python send_recv.py"
+    "!$MPIRUN  -n 2 -c 2 python send_recv.py"
    ]
   },
   {
@@ -200,7 +192,7 @@
    },
    "outputs": [],
    "source": [
-    "!$SRUN  -n 2 -c 2 python send_recv_eff.py"
+    "!$MPIRUN  -n 2 -c 2 python send_recv_eff.py"
    ]
   },
   {
@@ -238,7 +230,7 @@
    },
    "outputs": [],
    "source": [
-    "!$SRUN  -n 2 -c 2 python send_recv_numpy.py"
+    "!$MPIRUN  -n 2 -c 2 python send_recv_numpy.py"
    ]
   },
   {
@@ -271,7 +263,7 @@
    },
    "outputs": [],
    "source": [
-    "!$SRUN  -n 4 -c 4 python barrier.py"
+    "!$MPIRUN  -n 4 -c 4 python barrier.py"
    ]
   },
   {
@@ -309,7 +301,7 @@
    },
    "outputs": [],
    "source": [
-    "!$SRUN  -n 4 -c 4 python scatter.py"
+    "!$MPIRUN  -n 4 -c 4 python scatter.py"
    ]
   },
   {
@@ -350,7 +342,7 @@
    },
    "outputs": [],
    "source": [
-    "!$SRUN  -n 4 -c 4 python gather.py"
+    "!$MPIRUN  -n 4 -c 4 python gather.py"
    ]
   },
   {
@@ -384,7 +376,7 @@
    },
    "outputs": [],
    "source": [
-    "!$SRUN  -n 4 -c 4 python reduce.py"
+    "!$MPIRUN  -n 4 -c 4 python reduce.py"
    ]
   },
   {
@@ -448,7 +440,7 @@
    },
    "outputs": [],
    "source": [
-    "!$SRUN  -n 4 -c 4 python compute_pi.py"
+    "!$MPIRUN  -n 4 -c 4 python compute_pi.py"
    ]
   },
   {
@@ -547,7 +539,7 @@
    },
    "outputs": [],
    "source": [
-    "!$SRUN  -n 4 -c 4 python compute_mandelbot.py"
+    "!$MPIRUN  -n 4 -c 4 python compute_mandelbot.py"
    ]
   },
   {