00_core.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# core\n",
    "\n",
    "> Core functionality for distributing Earth Engine requests among Dask workers."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#| default_exp core"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#| hide\n",
    "from nbdev.showdoc import *"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#| export\n",
    "import coiled\n",
    "import dask.distributed\n",
    "import ee\n",
    "import google.auth"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#| export\n",
    "class InitEarthEngine(dask.distributed.WorkerPlugin):\n",
    "    def __init__(self, **kwargs):\n",
    "        print('InitEarthEngine.init: starting')  # This appears in the notebook output where the cluster is initiated.\n",
    "        self.kwargs = kwargs\n",
    "\n",
    "    def setup(self, worker):\n",
    "        # Print statements output to the dask cluster logs (viewable via Coiled dashboard)\n",
    "        print('InitEarthEngine.setup: starting')\n",
    "        print('InitEarthEngine.setup: default to using the high volume endpoint')\n",
    "        self.kwargs.setdefault('opt_url', 'https://earthengine-highvolume.googleapis.com')\n",
    "        import ee\n",
    "        print(f'InitEarthEngine.setup: ee.Initialize(**{self.kwargs})')\n",
    "        ee.Initialize(**self.kwargs)\n",
    "\n",
    "\n",
    "class ClusterGEE(coiled.Cluster):\n",
    "    def __init__(self, **kwargs):\n",
    "        print('ClusterGEE init')\n",
    "        super().__init__(**kwargs)\n",
    "        # Wait for the workers to start, then send the ADCs\n",
    "        self.wait_for_workers(kwargs['n_workers'])\n",
    "        coiled.credentials.google.send_application_default_credentials(self)\n",
    "\n",
    "    def get_client(self):\n",
    "        print('ClusterGEE get_client')\n",
    "        client = super().get_client()\n",
    "        client.register_plugin(InitEarthEngine())\n",
    "        return client\n",
    "\n",
    "\n",
    "# For local development\n",
    "class LocalClusterGEE(dask.distributed.LocalCluster):    \n",
    "    def __init__(self, **kwargs):\n",
    "        print('LocalClusterGEE init')\n",
    "        super().__init__(**kwargs)\n",
    "\n",
    "    def get_client(self):\n",
    "        print('LocalClusterGEE get_client')\n",
    "        client = super().get_client()\n",
    "        client.register_plugin(InitEarthEngine())\n",
    "        return client"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Try it out\n",
    "\n",
    "Create a cluster and run a few jobs."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Authenticate & Initialize Earth Engine"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Get credentials and the GCP project ID, authenticating if necessary."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "try:\n",
    "    credentials, project_id = google.auth.default()\n",
    "except google.auth.exceptions.DefaultCredentialsError:\n",
    "    print('Unable to get auth credentials. Initiating auth flow.')\n",
    "    !gcloud auth application-default login\n",
    "    credentials, project_id = google.auth.default()\n",
    "try:\n",
    "    ee.Initialize(credentials=credentials, project=project_id)\n",
    "except google.auth.exceptions.RefreshError:\n",
    "    print('Credentials refresh error. Initiating auth flow.')\n",
    "    !gcloud auth application-default login\n",
    "    credentials, project_id = google.auth.default()\n",
    "ee.Initialize(credentials=credentials, project=project_id)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Start Dask Cluster\n",
    "\n",
    "Start up a Earth Engine enabled cluster. This may take a few minutes to complete."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "ClusterGEE init\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "[2024-06-25 07:27:45,250][INFO    ][coiled] Resolving your local default Python environment...\n",
      "[2024-06-25 07:27:56,526][INFO    ][coiled] Creating Cluster (name: test-class-cluster, https://cloud.coiled.io/clusters/505936?account=tylere ). This usually takes 1-2 minutes...\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Google Application Default Credentials have been written to a file on your Coiled VM(s).\n",
      "These credentials will potentially be valid until explicitly revoked by running\n",
      "gcloud auth application-default revoke\n"
     ]
    }
   ],
   "source": [
    "cluster = ClusterGEE(\n",
    "    name='test-class-cluster',\n",
    "    n_workers=1,\n",
    "    worker_cpu=4,\n",
    "    spot_policy=\"spot_with_fallback\",\n",
    "    region='us-central1',\n",
    "    idle_timeout=\"1 hours\",\n",
    ")\n",
    "# cluster = LocalClusterGEE()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Retrieve a client for the cluster, and display it."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "ClusterGEE get_client\n",
      "InitEarthEngine.init: starting\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<div>\n",
       "    <div style=\"width: 24px; height: 24px; background-color: #e1e1e1; border: 3px solid #9D9D9D; border-radius: 5px; position: absolute;\"> </div>\n",
       "    <div style=\"margin-left: 48px;\">\n",
       "        <h3 style=\"margin-bottom: 0px;\">Client</h3>\n",
       "        <p style=\"color: #9D9D9D; margin-bottom: 0px;\">Client-638c87d8-32ff-11ef-84e0-fe11494405b5</p>\n",
       "        <table style=\"width: 100%; text-align: left;\">\n",
       "\n",
       "        <tr>\n",
       "        \n",
       "            <td style=\"text-align: left;\"><strong>Connection method:</strong> Cluster object</td>\n",
       "            <td style=\"text-align: left;\"><strong>Cluster type:</strong> __main__.ClusterGEE</td>\n",
       "        \n",
       "        </tr>\n",
       "\n",
       "        \n",
       "            <tr>\n",
       "                <td style=\"text-align: left;\">\n",
       "                    <strong>Dashboard: </strong> <a href=\"https://cluster-ysonu.dask.host/z8jKyCbSXIVx2pAk/status\" target=\"_blank\">https://cluster-ysonu.dask.host/z8jKyCbSXIVx2pAk/status</a>\n",
       "                </td>\n",
       "                <td style=\"text-align: left;\"></td>\n",
       "            </tr>\n",
       "        \n",
       "\n",
       "        </table>\n",
       "\n",
       "        \n",
       "\n",
       "        \n",
       "            <details>\n",
       "            <summary style=\"margin-bottom: 20px;\"><h3 style=\"display: inline;\">Cluster Info</h3></summary>\n",
       "            <div class=\"jp-RenderedHTMLCommon jp-RenderedHTML jp-mod-trusted jp-OutputArea-output\">\n",
       "    <div style=\"width: 24px; height: 24px; background-color: #e1e1e1; border: 3px solid #9D9D9D; border-radius: 5px; position: absolute;\">\n",
       "    </div>\n",
       "    <div style=\"margin-left: 48px;\">\n",
       "        <h3 style=\"margin-bottom: 0px; margin-top: 0px;\">ClusterGEE</h3>\n",
       "        <p style=\"color: #9D9D9D; margin-bottom: 0px;\">test-class-cluster</p>\n",
       "        <table style=\"width: 100%; text-align: left;\">\n",
       "            <tr>\n",
       "                <td style=\"text-align: left;\">\n",
       "                    <strong>Dashboard:</strong> <a href=\"https://cluster-ysonu.dask.host/z8jKyCbSXIVx2pAk/status\" target=\"_blank\">https://cluster-ysonu.dask.host/z8jKyCbSXIVx2pAk/status</a>\n",
       "                </td>\n",
       "                <td style=\"text-align: left;\">\n",
       "                    <strong>Workers:</strong> 1\n",
       "                </td>\n",
       "            </tr>\n",
       "            <tr>\n",
       "                <td style=\"text-align: left;\">\n",
       "                    <strong>Total threads:</strong> 4\n",
       "                </td>\n",
       "                <td style=\"text-align: left;\">\n",
       "                    <strong>Total memory:</strong> 15.02 GiB\n",
       "                </td>\n",
       "            </tr>\n",
       "            \n",
       "        </table>\n",
       "\n",
       "        <details>\n",
       "            <summary style=\"margin-bottom: 20px;\">\n",
       "                <h3 style=\"display: inline;\">Scheduler Info</h3>\n",
       "            </summary>\n",
       "\n",
       "            <div style=\"\">\n",
       "    <div>\n",
       "        <div style=\"width: 24px; height: 24px; background-color: #FFF7E5; border: 3px solid #FF6132; border-radius: 5px; position: absolute;\"> </div>\n",
       "        <div style=\"margin-left: 48px;\">\n",
       "            <h3 style=\"margin-bottom: 0px;\">Scheduler</h3>\n",
       "            <p style=\"color: #9D9D9D; margin-bottom: 0px;\">Scheduler-0612b84c-f388-4d1e-9495-0b25ef47c7a3</p>\n",
       "            <table style=\"width: 100%; text-align: left;\">\n",
       "                <tr>\n",
       "                    <td style=\"text-align: left;\">\n",
       "                        <strong>Comm:</strong> tls://10.2.0.85:8786\n",
       "                    </td>\n",
       "                    <td style=\"text-align: left;\">\n",
       "                        <strong>Workers:</strong> 1\n",
       "                    </td>\n",
       "                </tr>\n",
       "                <tr>\n",
       "                    <td style=\"text-align: left;\">\n",
       "                        <strong>Dashboard:</strong> <a href=\"http://10.2.0.85:8787/status\" target=\"_blank\">http://10.2.0.85:8787/status</a>\n",
       "                    </td>\n",
       "                    <td style=\"text-align: left;\">\n",
       "                        <strong>Total threads:</strong> 4\n",
       "                    </td>\n",
       "                </tr>\n",
       "                <tr>\n",
       "                    <td style=\"text-align: left;\">\n",
       "                        <strong>Started:</strong> Just now\n",
       "                    </td>\n",
       "                    <td style=\"text-align: left;\">\n",
       "                        <strong>Total memory:</strong> 15.02 GiB\n",
       "                    </td>\n",
       "                </tr>\n",
       "            </table>\n",
       "        </div>\n",
       "    </div>\n",
       "\n",
       "    <details style=\"margin-left: 48px;\">\n",
       "        <summary style=\"margin-bottom: 20px;\">\n",
       "            <h3 style=\"display: inline;\">Workers</h3>\n",
       "        </summary>\n",
       "\n",
       "        \n",
       "        <div style=\"margin-bottom: 20px;\">\n",
       "            <div style=\"width: 24px; height: 24px; background-color: #DBF5FF; border: 3px solid #4CC9FF; border-radius: 5px; position: absolute;\"> </div>\n",
       "            <div style=\"margin-left: 48px;\">\n",
       "            <details>\n",
       "                <summary>\n",
       "                    <h4 style=\"margin-bottom: 0px; display: inline;\">Worker: test-class-cluster-worker-31924985b5</h4>\n",
       "                </summary>\n",
       "                <table style=\"width: 100%; text-align: left;\">\n",
       "                    <tr>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>Comm: </strong> tls://10.2.0.84:38123\n",
       "                        </td>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>Total threads: </strong> 4\n",
       "                        </td>\n",
       "                    </tr>\n",
       "                    <tr>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>Dashboard: </strong> <a href=\"http://10.2.0.84:8787/status\" target=\"_blank\">http://10.2.0.84:8787/status</a>\n",
       "                        </td>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>Memory: </strong> 15.02 GiB\n",
       "                        </td>\n",
       "                    </tr>\n",
       "                    <tr>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>Nanny: </strong> tls://10.2.0.84:33847\n",
       "                        </td>\n",
       "                        <td style=\"text-align: left;\"></td>\n",
       "                    </tr>\n",
       "                    <tr>\n",
       "                        <td colspan=\"2\" style=\"text-align: left;\">\n",
       "                            <strong>Local directory: </strong> /scratch/dask-scratch-space/worker-ys5ui6sq\n",
       "                        </td>\n",
       "                    </tr>\n",
       "\n",
       "                    \n",
       "\n",
       "                    \n",
       "\n",
       "                </table>\n",
       "            </details>\n",
       "            </div>\n",
       "        </div>\n",
       "        \n",
       "\n",
       "    </details>\n",
       "</div>\n",
       "\n",
       "        </details>\n",
       "    </div>\n",
       "</div>\n",
       "            </details>\n",
       "        \n",
       "\n",
       "    </div>\n",
       "</div>"
      ],
      "text/plain": [
       "<Client: 'tls://10.2.0.85:8786' processes=1 threads=4, memory=15.02 GiB>"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "client = cluster.get_client()\n",
    "client"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Submit Jobs\n",
    "\n",
    "Test it out by:\n",
    "- Defining a function that can be distributed,\n",
    "- Submitting jobs running the function to workers, \n",
    "- Gathering the results locally, and\n",
    "- Displaying the results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Submitting jobs...\n",
      "...done\n",
      "Gathering results...\n",
      "...done\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "[{'country': 'Abyei Area', 'area_km2': 10460, 'mean_elev': 402.5921903247955},\n",
       " {'country': 'Zimbabwe', 'area_km2': 391916, 'mean_elev': 973.2955548809969}]"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Get a list of countries to analyze.\n",
    "country_fc = ee.FeatureCollection('USDOS/LSIB_SIMPLE/2017')\n",
    "country_list = country_fc.aggregate_array('country_na').distinct().sort().getInfo()\n",
    "\n",
    "import random\n",
    "\n",
    "# Write a function that can be run by the cluster workers. \n",
    "def get_country_stats(country_name):\n",
    "    country = country_fc.filter(ee.Filter.eq('country_na', country_name))\n",
    "    elev = ee.ImageCollection(\"COPERNICUS/DEM/GLO30\").select('DEM').mosaic()\n",
    "    return {\n",
    "        'country': country_name, \n",
    "        'area_km2': country.geometry().area().multiply(1e-6).round().getInfo(), \n",
    "        'mean_elev': elev.reduceRegion(reducer=ee.Reducer.mean(),\n",
    "                                       geometry=country.geometry(),\n",
    "                                       scale=10000,\n",
    "                                       ).get('DEM').getInfo(),\n",
    "    }\n",
    "\n",
    "\n",
    "# Create and submit jobs among the workers.\n",
    "print('Submitting jobs...')\n",
    "futures = [\n",
    "    client.submit(get_country_stats, country, retries=3)\n",
    "    for country in ['Abyei Area', 'Zimbabwe']\n",
    "]\n",
    "print('...done')\n",
    "\n",
    "# Gather up the results and display them.\n",
    "print('Gathering results...')\n",
    "results = client.gather(futures)\n",
    "print('...done')\n",
    "results"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Shut down the cluster."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "[2024-06-25 07:30:01,603][INFO    ][coiled] Cluster 505936 deleted successfully.\n"
     ]
    }
   ],
   "source": [
    "cluster.shutdown()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#| hide\n",
    "import nbdev; nbdev.nbdev_export()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "python3",
   "language": "python",
   "name": "python3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}