automatically restart kernel in colab

Author: ziw-liu

examples/virtual_staining/vcp_tutorials/quick_start.ipynb

@@ -124,17 +124,18 @@
    ]
   },
   {
-   "cell_type": "markdown",
-   "id": "2ae23986",
-   "metadata": {
-    "cell_marker": "\"\"\""
-   },
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "f0f20186",
+   "metadata": {},
+   "outputs": [],
    "source": [
-    "<div class=\"alert alert-warning\">\n",
-    "On Google Colab, please restart the session after installing dependencies.\n",
-    "Click 'Runtime -> Restart session' in the header bar\n",
-    "or press <kbd>⌘</kbd>/<kbd>Ctrl</kbd> + <kbd>m</kbd> <kbd>.</kbd>.\n",
-    "</div>"
+    "# restart kernel if running in Google Colab\n",
+    "if \"get_ipython\" in globals():\n",
+    "    session = get_ipython()\n",
+    "    if \"google.colab\" in str(session):\n",
+    "        print(\"Shutting down colab session.\")\n",
+    "        session.kernel.do_shutdown(restart=True)"
    ]
   },
   {

examples/virtual_staining/vcp_tutorials/quick_start.py

@@ -90,14 +90,13 @@
 # See the [repository](https://github.com/mehta-lab/VisCy) for more details
 # !pip install "viscy[metrics,visual]==0.3.0rc3"
 
-# %% [markdown]
-"""
-<div class="alert alert-warning">
-On Google Colab, please restart the session after installing dependencies.
-Click 'Runtime -> Restart session' in the header bar
-or press <kbd>⌘</kbd>/<kbd>Ctrl</kbd> + <kbd>m</kbd> <kbd>.</kbd>.
-</div>
-"""
+# %%
+# restart kernel if running in Google Colab
+if "get_ipython" in globals():
+    session = get_ipython()
+    if "google.colab" in str(session):
+        print("Shutting down colab session.")
+        session.kernel.do_shutdown(restart=True)
 
 # %%
 # Validate installation