2-dask-localcluster.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Best-practices for Cloud-Optimized Geotiffs\n",
    "\n",
    "\n",
    "**Part 3. Dask LocalCluster**\n",
    "\n",
    "As the number of COGs starts to grow you can quickly excede the amount of RAM on your system. This is where a Dask Cluster can be extremely useful. A LocalCluster is able to utlize all your CPUs and will manage your RAM such that you shouldn't get 'out of memory' errors when running computations. Often this amount of parallelism is all you need for efficient data exploration and analysis.\n",
    "\n",
    "In this notebook we'll focus on computing the temporal mean for a stack of COGS that excede our notebook memory (8GB)"
   ]
  },
  {
   "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, LocalCluster, progress"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "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": "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",
    "images[:6] #january 2020 scenes"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 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": "markdown",
   "metadata": {},
   "source": [
    "### Load in series (no dask)\n",
    "\n",
    "skip these cells and go to Dask if you want to avoid local caching from previous steps"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "# ~8.5s\n",
    "\n",
    "# 100 images, 12 GB uncompressed\n",
    "# All the image metadata = ~275 images, 33GB uncompressed\n",
    "\n",
    "\n",
    "dataArrays = [xr.open_rasterio(url, chunks={}) for url in images]\n",
    "# note use of join='override' b/c we know these COGS have the same coordinates\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": [
    "# as a benchmark, let's say we want to calculate the mean of each of these COGs.\n",
    "# we can just loop over all 100 images if each calculation is ~1.5s (based on single-cog notebook) that should take ~2.5 min:\n",
    "100 * 1.5 / 60\n",
    "\n",
    "# we should be able to do better than this though since each COG can be operated on independently..."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Load in parallel (dask w/ threads)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cluster = LocalCluster(processes=False, local_directory='/tmp') \n",
    "client = Client(cluster) \n",
    "client\n",
    "# NOTES: \n",
    "# dask workers write to SSD (/tmp) rather than home directory NFS mount\n",
    "# 1 worker, 4 cores --> 1 process w/ 4 threads\n",
    "# Open 'Dask 'Graph', 'Task Steam', and 'Workers' from labextension to see computation progress"
   ]
  },
  {
   "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",
    "# ~10s ... basically loading in series (file locks?)\n",
    "# picks up cache if run again (300ms)\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": [
    "# lets say we want the spatial mean of each COG. We can operate on each of these 278 files simultaneously (\"embarrassingly parallel\")\n",
    "# workers should be able to operate on each COG in isolation and just return a single result\n",
    "\n",
    "# It can be helpful to look at the task graph for a single COG like so:\n",
    "da.isel(time=0).mean(dim=['x','y']).data.visualize(optimize_graph=True, rankdir='LR')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "\n",
    "# 2min 25s\n",
    "# spatial mean of each COG (output = 278x1 vector)\n",
    "# task stream shows that this actually goes in series (due to xr.open_rasterio file lock?)\n",
    "\n",
    "da.mean(dim=['x','y']).compute()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "\n",
    "# ~ 2 min 32s\n",
    "# temporal mean of all COGs (output = 5490x5490 array)\n",
    "# this workflow requires pulling (nCOGS x chunk size) into worker RAM to get mean through time for each chunk (3GB)\n",
    "\n",
    "da.mean(dim='time').compute()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# GOTCHAS: The following is not a good idea because the output is the full uncompressed DataArray in local memory, \n",
    "# so we eventually hit RAM limits and start writing bytes to disk instead of RAM or the computation fails\n",
    "\n",
    "#scaled = da + 100\n",
    "#scaled.compute() "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### recap\n",
    "\n",
    "* The initial load of this dataset is slow b/c each thread is reading metadata sequentially \n",
    "* subsequent calls to da are an order of magnitude faster b/c the file handles and metadata are cached locally\n",
    "* computations can be slow (maybe due to file read locks preventing simultaneous operations)\n",
    "* threads are good for computations where memory needs to be shared by tasks (e.g. temporal mean for many COGs)\n",
    "* might want to experiment with chunk sizes (30--> 100MB), there should be less network requests that way..."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Load in parallel (dask w/ processes)\n",
    "\n",
    "Restart the kernel before running this section to avoid cache in timing"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#processes=True allows us to open COGs in parallel, circumventing locks. should be faster by a factor of 'nCores'\n",
    "# we have 4 by default on this machine\n",
    "cluster = LocalCluster(local_directory='/tmp') #processes=True by default\n",
    "client = Client(cluster) \n",
    "client"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time \n",
    "\n",
    "# ~3.7s first run, 1.6s subsequent run (caching but maybe cache is separate per process?),\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",
    "# 42.8s\n",
    "\n",
    "da.mean(dim=['x','y']).compute()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "\n",
    "# Kernel dies :(\n",
    "\n",
    "# temporal mean of all COGs (output = 5490x5490 array)\n",
    "# this workflow requires pulling (nCOGS x chunk size) into worker RAM to get mean through time for each chunk (3GB)\n",
    "# because each processe uses it's own RAM with a max of 2GB, we are forced to do some writing to disk and this is super slow.\n",
    "\n",
    "da.mean(dim='time').compute() # task stream is very inefficient here with high memory use!"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### recap\n",
    "\n",
    "* processes=True is great for dask delayed opening a bunch of datasets\n",
    "* it's also great for tasks where workers don't need to communicate information\n",
    "* it is really bad if tasks need to store a lot of intermediate results in memory (workers start writing to disk instead of RAM) "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Best of both worlds?\n",
    "\n",
    "Turns out you can mix and match dask cluster operations in a workflow. For example:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time \n",
    "\n",
    "with LocalCluster(local_directory='/tmp') as cluster, Client(cluster) as client:\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",
    "    spatial_means = da.mean(dim=['x','y']).compute()\n",
    "\n",
    "with LocalCluster(processes=False, local_directory='/tmp') as cluster, Client(cluster) as client:\n",
    "    temporal_mean = da.mean(dim='time').compute()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### recap\n",
    "\n",
    "* if performance is what you're going for, you might need to mix and match dask settings with processes and threads\n",
    "* while COGs are loaded as Dask Arrays via xarray, references to files and file locks can complicated parallelization"
   ]
  }
 ],
 "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
}