Example_2 Intake+Dask parallel download.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import xarray as xr"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "from dask_jobqueue import SLURMCluster\n",
    "cluster = SLURMCluster()\n",
    "cluster"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "# from distributed import LocalCluster\n",
    "# cluster = LocalCluster()\n",
    "# cluster"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<table style=\"border: 2px solid white;\">\n",
       "<tr>\n",
       "<td style=\"vertical-align: top; border: 0px solid white\">\n",
       "<h3>Client</h3>\n",
       "<ul>\n",
       "  <li><b>Scheduler: </b>tcp://localhost:8789\n",
       "  <li><b>Dashboard: </b><a href='http://localhost:8787/status' target='_blank'>http://localhost:8787/status</a>\n",
       "</ul>\n",
       "</td>\n",
       "<td style=\"vertical-align: top; border: 0px solid white\">\n",
       "<h3>Cluster</h3>\n",
       "<ul>\n",
       "  <li><b>Workers: </b>40</li>\n",
       "  <li><b>Cores: </b>80</li>\n",
       "  <li><b>Memory: </b>600.00 GB</li>\n",
       "</ul>\n",
       "</td>\n",
       "</tr>\n",
       "</table>"
      ],
      "text/plain": [
       "<Client: scheduler='tcp://146.118.38.56:8789' processes=40 cores=80>"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from distributed import Client\n",
    "# client=Client(cluster)\n",
    "client=Client('tcp://localhost:8789')\n",
    "client"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "# If a number of workers get killed or things just seem a bit haywire, or you want to clear the temp memory, \n",
    "# it can be useful to run this\n",
    "# it kills and restarts the python kernels on each worker\n",
    "\n",
    "# client.restart()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Demonstrate the usage of intake-thredds"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "import intake\n",
    "cat = intake.open_thredds_cat('http://thredds.aodn.org.au/thredds/catalog/IMOS/SRS/OC/gridded/aqua/P1D/catalog.xml')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Pick out a month - the \\[''\\]() syntax is a little ugly, but it works to return the catalogue objects"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'http://thredds.aodn.org.au/thredds/catalog/IMOS/SRS/OC/gridded/aqua/P1D/2007/02/catalog.xml'"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "month=cat['2007']()['02']()\n",
    "month.url"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Use Dask to walk through the the THREDDS catalog in parallel\n",
    "\n",
    "Use dask to traverse the catalog in parallel. This involves creating some delayed functions that are chained together"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "from dask import delayed\n",
    "from dask.distributed import get_client\n",
    "\n",
    "@delayed\n",
    "def obtain_day_urls(url):\n",
    "    day_urls = []\n",
    "    cat = intake.open_thredds_cat(url)\n",
    "    for filename in cat:\n",
    "        if 'K_490' in filename: # Filter the filenames - just the light attenuation files ('*K_490*')\n",
    "            day_urls.append(cat[filename].urlpath)\n",
    "    return day_urls\n",
    "\n",
    "@delayed\n",
    "def obtain_month_urls(url):\n",
    "    month_urls = []\n",
    "    cat = intake.open_thredds_cat(url)\n",
    "    for month in cat:\n",
    "        day_urls=obtain_day_urls(cat[month].url)\n",
    "        month_urls = month_urls+day_urls\n",
    "    return month_urls\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This will return a delayed function call (yet to be evaluated) to request the catalogs for the months of each year"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2002\n",
      "2003\n",
      "2004\n",
      "2005\n",
      "2006\n",
      "2007\n",
      "2008\n",
      "2009\n",
      "2010\n",
      "2011\n",
      "2012\n",
      "2013\n",
      "2014\n",
      "2015\n",
      "2016\n",
      "2017\n",
      "2018\n",
      "2019\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "[Delayed('obtain_month_urls-541f2418-97b7-40c9-bf62-e37025e2b494'),\n",
       " Delayed('obtain_month_urls-9b6cc6fc-8327-47f2-bb32-1b7623f5f236'),\n",
       " Delayed('obtain_month_urls-750ff324-ebc5-44ad-a639-528436c67fed'),\n",
       " Delayed('obtain_month_urls-c65720c9-23cb-4ea2-98e1-98fd97fc4b64'),\n",
       " Delayed('obtain_month_urls-b74b3e7b-90e3-4f10-8560-8745d103cafb'),\n",
       " Delayed('obtain_month_urls-b6c42bd6-a559-492e-a706-84ddf770d75b'),\n",
       " Delayed('obtain_month_urls-1ff14604-a98a-42d9-9acf-6a212230bea2'),\n",
       " Delayed('obtain_month_urls-7bda8d8d-eee9-4c6d-8afd-bb9c196017d0'),\n",
       " Delayed('obtain_month_urls-71e80985-8159-424f-9929-1c533f5b95d9'),\n",
       " Delayed('obtain_month_urls-e43841ea-da48-4093-957a-1d422e369dcc'),\n",
       " Delayed('obtain_month_urls-40b29213-ae88-4353-ac6b-e6941c3f6a1f'),\n",
       " Delayed('obtain_month_urls-a0e7e3be-04f9-45a7-a624-afe6afdb5dc8'),\n",
       " Delayed('obtain_month_urls-af1cc414-0d84-4e06-bbb8-78a7875c858a'),\n",
       " Delayed('obtain_month_urls-500f6b5a-45ef-441e-a510-086b662ac07b'),\n",
       " Delayed('obtain_month_urls-b9b92d9c-a173-49cc-8726-eb5d71569bcd'),\n",
       " Delayed('obtain_month_urls-cb8ee651-7f53-409e-8ae5-2205d46a4fdf'),\n",
       " Delayed('obtain_month_urls-5b7060f9-b4ca-43b3-9ba7-c0a0d5ed3a71'),\n",
       " Delayed('obtain_month_urls-44afd256-20cf-4613-bb9e-177119f8039a')]"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "futures=[]\n",
    "for year in cat:\n",
    "    # if year in ['2014',]: #could filter here\n",
    "    print(year)\n",
    "    futures.append(obtain_month_urls(cat[year].url))\n",
    "futures"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now send these delayed function calls to the dask scheduler to be computed. First to return the months of each year and then the all the files for each month"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "from dask import compute\n",
    "month_urls=compute(futures)[0]\n",
    "day_urls=compute(month_urls)[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "6102"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# day_urls contains a list of lists for each month, concatentate them all together\n",
    "all_urls = []\n",
    "for m in day_urls:\n",
    "    all_urls=all_urls+m\n",
    "len(all_urls)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Note, now we have the thredds 'dodsC' OpenDAP URLs for each netCDF file, this is whats needed to open the file directly using Xarray"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['http://thredds.aodn.org.au/thredds/dodsC/IMOS/SRS/OC/gridded/aqua/P1D/2002/07/A.P1D.20020706T000000Z.aust.K_490.nc',\n",
       " 'http://thredds.aodn.org.au/thredds/dodsC/IMOS/SRS/OC/gridded/aqua/P1D/2002/07/A.P1D.20020707T000000Z.aust.K_490.nc',\n",
       " 'http://thredds.aodn.org.au/thredds/dodsC/IMOS/SRS/OC/gridded/aqua/P1D/2002/07/A.P1D.20020708T000000Z.aust.K_490.nc',\n",
       " 'http://thredds.aodn.org.au/thredds/dodsC/IMOS/SRS/OC/gridded/aqua/P1D/2002/07/A.P1D.20020709T000000Z.aust.K_490.nc',\n",
       " 'http://thredds.aodn.org.au/thredds/dodsC/IMOS/SRS/OC/gridded/aqua/P1D/2002/07/A.P1D.20020710T000000Z.aust.K_490.nc']"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "all_urls[0:5]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Probably good to archive all the URLs somewhere, quicker than hitting the thredds server again"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "import yaml\n",
    "# save/load a yaml file with the URLs for later....\n",
    "with open('K490_url.yml', 'w') as outfile:\n",
    "    yaml.dump(all_urls, outfile, default_flow_style=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "# with open('K490_url.yml', 'r') as infile:\n",
    "#     all_urls=yaml.safe_load(infile)\n",
    "# all_urls[0:10]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### And download all the files using the available workers"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And we can use dask a bit like multiprocessing 'map' and call a delayed function for each URL. Here I am substituting 'fileServer' for dodsC as I want to download the whole file rather than make an OpenDAP request. We can monitor the progress from the dask dashboard link provided above"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import urllib.request, os \n",
    "\n",
    "@delayed\n",
    "def download_file_tds(url):\n",
    "    dest=os.environ['MYSCRATCH'] + '/aqua_files/' + url.split('/')[-1]\n",
    "    urllib.request.urlretrieve(url.replace('dodsC','fileServer'), dest)\n",
    "    return dest\n",
    "\n",
    "# url format to go directly to S3: https://s3-ap-southeast-2.amazonaws.com/imos-data/IMOS/SRS/OC/gridded/aqua/P1D/2002/07/A.P1D.20020705T000000Z.aust.dt.nc\n",
    "# with 40 workers on Pawsey, the files downloaded at approximately 2GB/s (!)\n",
    "@delayed\n",
    "def download_file_s3(url):\n",
    "    dest=os.environ['MYSCRATCH'] + '/aqua_files/' + url.split('/')[-1]\n",
    "    urllib.request.urlretrieve(url.replace('http://thredds.aodn.org.au/thredds/dodsC/','https://s3-ap-southeast-2.amazonaws.com/imos-data/'), dest)\n",
    "    return dest\n",
    "\n",
    "futures = [download_file_s3(url) for url in all_urls]\n",
    "\n",
    "files = compute(futures)[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python [conda env:c3dis2] *",
   "language": "python",
   "name": "conda-env-c3dis2-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.6.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}