3-dask-gatewaycluster.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Best-practices for Cloud-Optimized Geotiffs\n",
    "\n",
    "**Part 4. Dask GatewayCluster**\n",
    "\n",
    "Unlike LocalCluster, a Dask GatewayCluster gives us the ability to dynamically increase our CPU and RAM across many machines! This is extremely powerful, because now we can load very big datasets into RAM for efficient calculations. There is a complication in that now we are running computations on many physical machines instead of just one, so network communication is more challenging and the dask machines likely don't have access to your local files. When COGS are stored on S3 though, we can access them from any machine!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import xarray as xr\n",
    "import s3fs\n",
    "import pandas as pd\n",
    "import os \n",
    "\n",
    "import dask\n",
    "from dask.distributed import Client, progress\n",
    "from dask_gateway import Gateway"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# use the same GDAL environment settings as we did for the single COG case\n",
    "env = dict(GDAL_DISABLE_READDIR_ON_OPEN='EMPTY_DIR', \n",
    "           AWS_NO_SIGN_REQUEST='YES',\n",
    "           GDAL_MAX_RAW_BLOCK_CACHE_SIZE='200000000',\n",
    "           GDAL_SWATH_SIZE='200000000',\n",
    "           VSI_CURL_CACHE_SIZE='200000000')\n",
    "os.environ.update(env)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Dask GatewayCluster\n",
    "\n",
    "dask gateway allows us to connect to a Kubernetes Cluster so that we can go beyond the RAM and CPU of a single machine. It can take several minutes for these machines to initialize on the Cloud, so be patient when starting a cluster."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# dask gateway allow us to connect to a Kubernetes Cluster so that we can go beyond the RAM and CPU of a single machine\n",
    "\n",
    "# NOTE: we have to explicitly pass local environment variables to the cluster now\n",
    "# By default each worker has 2 cores and 4GB memory and effectively runs as a separate process\n",
    "# NOTE: how to deal with cores vs threads in a gateway cluster?\n",
    "\n",
    "gateway = Gateway()\n",
    "options = gateway.cluster_options()\n",
    "options.environment = env \n",
    "cluster = gateway.new_cluster(options)\n",
    "cluster.scale(4) # let's get the same number of \"workers\" as our previous LocalCluster examples"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# The dashboard link can also be pasted into the dask lab-extension\n",
    "cluster"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# NOTE: just like with a LocalCluster, it's good to explicitly connect to our GatewayCluster\n",
    "client = Client(cluster) "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# the dashboard link works just like a localcluster\n",
    "client"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "options"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Make sure that your dask workers see GDAL environment variables\n",
    "def get_env(env):\n",
    "    import os\n",
    "    return os.environ.get(env)\n",
    "\n",
    "print(client.run(get_env, 'GDAL_DISABLE_READDIR_ON_OPEN'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time \n",
    "\n",
    "s3 = s3fs.S3FileSystem(anon=True)\n",
    "objects = s3.glob('sentinel-s1-rtc-indigo/tiles/RTC/1/IW/10/T/ET/**Gamma0_VV.tif')\n",
    "images = ['s3://' + obj for obj in objects]\n",
    "print(len(images))\n",
    "images.sort(key=lambda x: x[-32:-24]) #sort list in place by date in filename\n",
    "# Let's use first 100 images for simplicity\n",
    "images = images[:100]\n",
    "dates = [pd.to_datetime(x[-32:-24]) for x in images]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "@dask.delayed\n",
    "def lazy_open(href):\n",
    "    chunks=dict(band=1, x=2745, y=2745)\n",
    "    return xr.open_rasterio(href, chunks=chunks) "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time \n",
    "\n",
    "# ~6.5 s\n",
    "\n",
    "dataArrays = dask.compute(*[lazy_open(href) for href in images])\n",
    "da = xr.concat(dataArrays, dim='band', join='override', combine_attrs='drop').rename(band='time')\n",
    "da['time'] = dates\n",
    "da"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "\n",
    "# 41 s\n",
    "\n",
    "da.mean(dim=['x','y']).compute()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "\n",
    "# 44 s\n",
    "# just like with a LocalCluster this workflow requires pulling (nCOGS x chunk size) into worker RAM to get mean through time for each chunk (3GB)\n",
    "# Now we have 4GB per worker (and we can adjust this via cluster settings)\n",
    "\n",
    "da.mean(dim='time').compute()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Visualization\n",
    "\n",
    "Using hvplot like we've done before will utilize the dask cluster as you request plotting of each image"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import hvplot.xarray\n",
    "da.hvplot.image(rasterize=True, \n",
    "                aspect='equal', frame_width=500,\n",
    "                cmap='gray', clim=(0,0.4))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python [conda env:notebook] *",
   "language": "python",
   "name": "conda-env-notebook-py"
  },
  "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.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}