{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "86aff193",
   "metadata": {},
   "source": [
    "# Dask Gateway cluster setup"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bea9c781-da57-4ad2-99fa-b0502465c4d3",
   "metadata": {
    "tags": []
   },
   "source": [
    "**Initialize `gateway` object.** It will be used to interact with your Dask clusters."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dd1e08c7-1606-4995-8ae9-750e45bfefe1",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import os\n",
    "from dask_gateway import Gateway\n",
    "\n",
    "# To submit jobs via SLURM (Purdue users only!)\n",
    "gateway = Gateway()\n",
    "\n",
    "# To submit jobs via Kubernetes (all users)\n",
    "# gateway = Gateway(\n",
    "#     \"http://dask-gateway-k8s.geddes.rcac.purdue.edu/\",\n",
    "#     proxy_address=\"traefik-dask-gateway-k8s.cms.geddes.rcac.purdue.edu:8786\",\n",
    "# )"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d07c6819-8987-437e-ad44-845eaa3241d0",
   "metadata": {},
   "source": [
    "**Create a new cluster.**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1f5f2107-344a-4340-8dd0-fa2758f93e40",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# You may need to update some environment variables before creating a cluster.\n",
    "# For example:\n",
    "os.environ[\"X509_USER_PROXY\"] = \"/path-to-voms-proxy/\"\n",
    "\n",
    "# Create the cluster\n",
    "cluster = gateway.new_cluster(\n",
    "conda_env = \"/depot/cms/kernels/python3\", # path to conda env\n",
    "worker_cores = 1,    # cores per worker\n",
    "worker_memory = 4,   # memory per worker in GB\n",
    "env = dict(os.environ), # pass environment as a dictionary\n",
    ")\n",
    "\n",
    "cluster"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "71c9092f",
   "metadata": {},
   "source": [
    "*This is how the widget for Dask Gateway cluster will look like, if it gets created successfully:*\n",
    "<div>\n",
    "<img src=\"images/dask-gateway-widget-cluster.png\" width=\"600\"/>\n",
    "</div>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6dde425e-c018-459f-90dc-e1bf332ba7c3",
   "metadata": {},
   "source": [
    "- Use adaptive (recommended) or manual scaling to create Dask workers.\n",
    "- Click on the dashboard link to open the Dask dashboard\n",
    "- To access worker logs, click on \"Info\" tab in the Dask dashboard"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e9ab3b24-215c-494c-85e7-a4d734075cea",
   "metadata": {},
   "source": [
    "**Check if you already have clusters running:**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "61d458a9-e2cd-4036-86ee-1a78462682ab",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# List available clusters\n",
    "clusters = gateway.list_clusters()\n",
    "print(clusters)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3e60ca38-a8a2-404e-9502-bc4d883fea39",
   "metadata": {},
   "source": [
    "**Shut down cluster.**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "22c9c5f6-29b7-4e44-8ed8-8b97bf1f808f",
   "metadata": {},
   "outputs": [],
   "source": [
    "cluster.shutdown()\n",
    "\n",
    "# Or shut down a specific cluster by name:\n",
    "# cluster_name = \"b2aec555e2f844d68a5febd6c5d1414e\"   # paste cluster name here\n",
    "# client = gateway.connect(cluster_name).shutdown()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ced19b61-6706-4861-8390-804261092e71",
   "metadata": {},
   "source": [
    "**Shut down all clusters:**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0b20a397-fe18-446b-a9f1-f164507c2524",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "for cluster_info in gateway.list_clusters():\n",
    "    gateway.connect(cluster_info.name).shutdown()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python3 kernel (default)",
   "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.10.10"
  },
  "widgets": {
   "application/vnd.jupyter.widget-state+json": {
    "state": {},
    "version_major": 2,
    "version_minor": 0
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}