basketballcantho/remote-sensing
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity
Files
add7d16ecfb705e55ee1105de806432979afd137
remote-sensing/notebooks/dask/02-ODC_with_Dask.ipynb
T

1202 lines
44 KiB
Plaintext
Raw Blame History

{
"cells": [
{
"cell_type": "markdown",
"id": "d7e59cba-16e4-49b9-928e-e41d1c0196f6",
"metadata": {},
"source": [
"# ODC and Dask (LocalCluster) <img align=\"right\" src=\"../../resources/csiro_easi_logo.png\">\n",
"\n",
"- [Setup](#Setup)\n",
"- [Load data without dask](#Load-data-without-dask)\n",
"- [Exploring dask concepts with ODC](#Exploring-dask-concepts-with-ODC)\n",
"- [The impact of dask on ODC](#The-impact-of-dask-on-ODC)\n",
"- [Exploiting delayed tasks](#Exploiting-delayed-tasks)\n",
"- [Data and computational locality](#Data-and-computational-locality)\n",
" - [With compute on the algorithm](#With-compute-on-the-algorithm)\n",
" - [With selected measurements](#With-selected-measurements)\n",
"- [A quick check on the task graph](#A-quick-check-on-the-task-graph)\n",
"\n",
"This notebook explores the use of ODC with Dask LocalCluster. The goal is to introduce fundamental concepts and the role Dask can serve with `datacube` and subsequent computation using `xarray`.\n",
"\n",
"The example computation is fairly typical of an EO data processing pipeline. We'll be using a small area and time period to start with and progressively scaling this example. EO scientists may find some aspects of these examples unrealistic, but this isn't an EO science course &#9786;. \n",
"\n",
"The basic workflow is:\n",
" 1. Specify Region of Interest, Satellite product, EO satellite bands, Time range, desired CRS for the `datacube` query\n",
" 1. Load data using `datacube.load()`\n",
" 1. Mask valid data\n",
" 1. Visualisation of the ROI\n",
" 1. Compute NDVI\n",
" 1. Visualise NDVI\n",
" \n",
"__NOTE__: Some cells in this notebook will take minutes to run so please be patient. Also, some cells can load large datasets into Jupyter's memory (based on the defaults), which can exhaust the available memory and cause the kernel to crash. If this occurs: restart the kernal, run the [Setup](#Setup) cells, and then jump to the next section. The (text) outputs from previous sections can be retained between kernel restarts. Example times for each exercise:\n",
"\n",
"| Process | N times | Load |\n",
"|--|--|--|\n",
"| [Load data without dask](#Load-data-without-dask) | 7 | 3-3.5 mins |\n",
"| [Exploring dask concepts with ODC](#Exploring-dask-concepts-with-ODC) | 2 | 1 min |\n",
"| [The impact of dask on ODC](#The-impact-of-dask-on-ODC) | 22 | 1-2 mins |\n",
"| [Exploiting delayed tasks](#Exploiting-delayed-tasks) | 7 | 1 min |\n",
"| [Data and computational locality](#Data-and-computational-locality) | 7 | 30-50 secs |\n",
"| [With compute on the algorithm](#With-compute-on-the-algorithm) | 7 | 10 secs |\n",
"| [With selected measurements](#With-selected-measurements) | 7 | 10 secs |\n",
"\n",
"### Setup"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "31d9ee40-ddbe-4f60-b935-8598b332a75f",
"metadata": {
"tags": []
},
"outputs": [],
"source": [
"import git\n",
"import sys, os\n",
"from dateutil.parser import parse\n",
"from dateutil.relativedelta import relativedelta\n",
"from dask.distributed import Client, LocalCluster\n",
"import datacube\n",
"from datacube.utils import masking\n",
"from datacube.utils.aws import configure_s3_access\n",
"\n",
"# EASI defaults\n",
"os.environ['USE_PYGEOS'] = '0'\n",
"repo = git.Repo('.', search_parent_directories=True).working_tree_dir\n",
"if repo not in sys.path: sys.path.append(repo)\n",
"from easi_tools import EasiDefaults, notebook_utils\n",
"easi = EasiDefaults()\n",
"client = None"
]
},
{
"cell_type": "markdown",
"id": "47aca1bd-ff2c-4d33-8a3b-6ca520f017d6",
"metadata": {},
"source": [
"The next cell sets out all the query parameters used in our `datacube.load()`.\n",
"For this run we keep the ROI quite small."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "346b679f-4927-4cef-9af8-ae94ed5c0227",
"metadata": {
"tags": []
},
"outputs": [],
"source": [
"# Get the default latitude & longitude extents\n",
"study_area_lat = easi.latitude\n",
"study_area_lon = easi.longitude\n",
"\n",
"# Or choose your own by uncommenting and modifying this section\n",
"###############################################################\n",
"# # Central Tasmania (near Little Pine Lagoon)\n",
"# central_lat = -42.019\n",
"# central_lon = 146.615\n",
"\n",
"# # Set the buffer to load around the central coordinates\n",
"# # This is a radial distance for the bbox to actual area so bbox 2x buffer in both dimensions\n",
"# buffer = 0.05\n",
"\n",
"# # Compute the bounding box for the study area\n",
"# study_area_lat = (central_lat - buffer, central_lat + buffer)\n",
"# study_area_lon = (central_lon - buffer, central_lon + buffer)\n",
"###############################################################\n",
"\n",
"# Data product\n",
"product = easi.product('landsat')\n",
"# product = 'landsat8_c2l2_sr'\n",
"\n",
"# Set the date range to load data over\n",
"set_time = easi.time\n",
"# set_time = (\"2021-01-01\", \"2021-01-31\")\n",
"\n",
"# Selected measurement names (used in this notebook). None` will load all of them\n",
"alias = easi.aliases('landsat')\n",
"measurements = None\n",
"# measurements = [alias[x] for x in ['qa_band', 'red', 'nir']]\n",
"\n",
"# Set the QA band name and mask values\n",
"qa_band = alias['qa_band']\n",
"qa_mask = easi.qa_mask('landsat')\n",
"\n",
"# Set the resampling method for the bands\n",
"resampling = {qa_band: \"nearest\", \"*\": \"average\"}\n",
"\n",
"# Set the coordinate reference system and output resolution\n",
"set_crs = easi.crs('landsat') # If defined, else None\n",
"set_resolution = easi.resolution('landsat') # If defined, else None\n",
"# set_crs = \"epsg:3577\"\n",
"# set_resolution = (-30, 30)\n",
"\n",
"# Set the scene group_by method\n",
"group_by = \"solar_day\""
]
},
{
"cell_type": "markdown",
"id": "fbd89d20-7532-45e8-85bb-6a29bd078971",
"metadata": {},
"source": [
"Now initialise the `datacube`."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "4e72951b-e95c-4901-84df-b8c28bf9fbed",
"metadata": {
"tags": []
},
"outputs": [],
"source": [
"dc = datacube.Datacube()\n",
"\n",
"# Access AWS \"requester-pays\" buckets\n",
"# This is necessary for reading data from most third-party AWS S3 buckets such as for Landsat and Sentinel-2\n",
"from datacube.utils.aws import configure_s3_access\n",
"configure_s3_access(aws_unsigned=False, requester_pays=True);"
]
},
{
"cell_type": "markdown",
"id": "499ec0c7-84d8-4a66-8d89-6c54e291e2e7",
"metadata": {},
"source": [
"## Load data without dask\n",
"\n",
"Now load the data. This first `dc.load()` does not use Dask so it will take a few minutes.\n",
"\n",
"We use `%%time` to keep track of how long things take to complete."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "0f8bdfa2-7422-41e2-91e0-3ca7ca630ece",
"metadata": {
"tags": []
},
"outputs": [],
"source": [
"%%time\n",
"dataset = None # clear results from any previous runs\n",
"dataset = dc.load(\n",
" product=product,\n",
" x=study_area_lon,\n",
" y=study_area_lat,\n",
" time=set_time,\n",
" measurements=measurements,\n",
" resampling=resampling,\n",
" output_crs=set_crs,\n",
" resolution=set_resolution,\n",
" group_by=group_by,\n",
" )"
]
},
{
"cell_type": "markdown",
"id": "9b1ec342-d654-41c1-bce7-c91b84269019",
"metadata": {},
"source": [
"The result of the `datacube.load()` function is an `xarray.Dataset`.\n",
"\n",
"Jupyter notebooks can render a description of the xarray `dataset` variable with a _lot of useful information_ about the structure of data."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "e62b096d-0659-43d1-91fb-e4499e8522f6",
"metadata": {
"tags": []
},
"outputs": [],
"source": [
"dataset"
]
},
{
"cell_type": "markdown",
"id": "f800f9f2-ef28-4bf0-94bb-0587c2c0b181",
"metadata": {},
"source": [
"Open the `Data variables` (click \"&#8227; Data variables\") and click on the stacked cylinders for one of them. You will see the actual data array is available and shown in summary form.\n",
"\n",
"> __NOTE__ that you can see real numbers in the array when you do this. This will change when we start using Dask.\n",
"\n",
"This graphical summary of an xarray variable will become increasingly importantly when dask is enabled and as scale out occurs so take a moment now to just poke around the interface. Depending on your area of interest set above, you should have a relatively small area: perhaps around 300 to 400 pixels in each of the `x` and `y` dimensions and perhaps up to 10 time slices. This is a relatively small size and fine to do without using Dask.\n",
"\n",
"Next up filter out pixels that are affect by clouds and other issues and compute the NDVI. Since we aren't specifying a time range this will be performed for all images."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "e4ccfa8f-6ed5-470c-9802-998ca7839d78",
"metadata": {
"tags": []
},
"outputs": [],
"source": [
"%%time\n",
"# Identify pixels that don't have cloud, cloud shadow or water\n",
"from datacube.utils import masking\n",
"\n",
"cloud_free_mask = masking.make_mask(dataset[qa_band], **qa_mask)\n",
"\n",
"# Apply the mask\n",
"cloud_free = dataset.where(cloud_free_mask)\n",
"\n",
"# Calculate the components that make up the NDVI calculation\n",
"band_diff = cloud_free[alias['nir']] - cloud_free[alias['red']]\n",
"band_sum = cloud_free[alias['nir']] + cloud_free[alias['red']]\n",
"# Calculate NDVI\n",
"ndvi = None\n",
"ndvi = band_diff / band_sum"
]
},
{
"cell_type": "markdown",
"id": "94ce1fad-4110-4585-9efd-acf271c70fc5",
"metadata": {},
"source": [
"The result `ndvi` is an `xarray.DataArray`. Let's take a look at it. Again the notebook will render an html version of the data in summary form.\n",
"Notice again the actual data values are being shown and that there are the same number of time slices as above and the x and y dimensions are identical."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "5fa68d03-b2a5-4124-88f6-3f8a9e8360cb",
"metadata": {
"tags": []
},
"outputs": [],
"source": [
"ndvi"
]
},
{
"cell_type": "markdown",
"id": "1a7596db-6a59-42a8-a8fd-411c909d4244",
"metadata": {},
"source": [
"Raw numbers aren't nice to look at so let's draw a time slice. We'll select just one of them to draw and pick one that didn't get masked out by cloud completely. You can see that all clouds and water has been masked out so that we are just looking at the NDVI of the land area."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "2b2905d2-042b-4751-83bd-55318aea96ea",
"metadata": {
"tags": []
},
"outputs": [],
"source": [
"ndvi.isel(time=1).plot()"
]
},
{
"cell_type": "markdown",
"id": "7cb7e93a-c72d-4d85-9d84-249beac08553",
"metadata": {},
"source": [
"## Exploring dask concepts with ODC\n",
"\n",
"Let's set our time range to a couple of weeks, or approximately two passes of Landsat 8 for this ROI. Less data will allow us to explore how dask works with the `datacube` and `xarray` libraries."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "81b5270b-2c88-4d6b-96ca-4a1b12c99adf",
"metadata": {
"tags": []
},
"outputs": [],
"source": [
"set_time = (set_time[0], parse(set_time[0]) + relativedelta(weeks=3))\n",
"# set_time = (\"2021-01-01\", \"2021-01-14\")\n",
"set_time"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "464f1e62-1ab2-459e-b6fa-3f564ac13209",
"metadata": {
"tags": []
},
"outputs": [],
"source": [
"%%time\n",
"dataset = None # clear results from any previous runs\n",
"dataset = dc.load(\n",
" product=product,\n",
" x=study_area_lon,\n",
" y=study_area_lat,\n",
" time=set_time,\n",
" measurements=measurements,\n",
" resampling=resampling,\n",
" output_crs=set_crs,\n",
" resolution=set_resolution,\n",
" group_by=group_by,\n",
" )\n",
"dataset"
]
},
{
"cell_type": "markdown",
"id": "d58d1061-f60b-4ead-98c3-52d852731716",
"metadata": {},
"source": [
"As before you can see the actual data in the results but this time there should only be 1 or 2 observation times"
]
},
{
"cell_type": "markdown",
"id": "6e0e104c-9427-4de1-9146-9a6067c55dcf",
"metadata": {},
"source": [
"### Start a dask LocalCluster\n",
"\n",
"Now let's create a `LocalCluster` as we did in the earlier notebook."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "0351ebf0-5c3b-4fbd-ae3e-a27f008b6b90",
"metadata": {
"tags": []
},
"outputs": [],
"source": [
"from dask.distributed import Client, LocalCluster\n",
"\n",
"cluster = LocalCluster()\n",
"client = Client(cluster)\n",
"client"
]
},
{
"cell_type": "markdown",
"id": "07d10fbc-951b-40ad-a86e-c2e1ccc812ee",
"metadata": {},
"source": [
"You may like to open up the dashboard for the cluster, although for this notebook we won't be talking about the dashboard (that's for a later discussion)."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "c265cddb-1371-4080-9eb2-4c0962b4242b",
"metadata": {
"tags": []
},
"outputs": [],
"source": [
"notebook_utils.localcluster_dashboard(client=client, server=easi.hub)"
]
},
{
"cell_type": "markdown",
"id": "cee3f697-57db-498a-92d8-3150368c7580",
"metadata": {},
"source": [
"Now that we are using a cluster, even though it is local, we need to make sure that our cluster has the right configuration to use __Requester Pays__ buckets in AWS S3. To do this, we need to re-run the `configure_s3_access()` function that we ran earlier, but we need to pass the `client` to the function as well."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "1b4dc0ba-93c4-4b09-804c-a1e4295d90aa",
"metadata": {
"tags": []
},
"outputs": [],
"source": [
"from datacube.utils.aws import configure_s3_access\n",
"configure_s3_access(aws_unsigned=False, requester_pays=True, client=client);"
]
},
{
"cell_type": "markdown",
"id": "6e39f921-db05-4466-bf51-70c3e19eb602",
"metadata": {},
"source": [
"`datacube.load()` will use the default `dask` cluster (the one we just created) __if the `dask_chunks` parameter is specified__.\n",
"\n",
"The chunk shape and memory size is a critial parameter in tuning `dask` and we will be discussing it in great detail as scale increases. For now we're simply going to specify that the `time` dimension should individually chunked (`1` slice of time) and by not specifying any chunking for the other dimensions they will be form a single contiguous block.\n",
"\n",
"If that made no sense what's so ever, that's fine because we will look at an example."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "2abfacab-bd47-4e78-b841-233d054309ba",
"metadata": {
"tags": []
},
"outputs": [],
"source": [
"chunks = {\"time\":1}"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "d9603c34-b948-4e29-b0a8-0d4ae1831243",
"metadata": {
"tags": []
},
"outputs": [],
"source": [
"%%time\n",
"dataset = None # clear results from any previous runs\n",
"dataset = dc.load(\n",
" product=product,\n",
" x=study_area_lon,\n",
" y=study_area_lat,\n",
" time=set_time,\n",
" measurements=measurements,\n",
" resampling=resampling,\n",
" output_crs=set_crs,\n",
" resolution=set_resolution,\n",
" dask_chunks = chunks, ###### THIS IS THE ONLY LINE CHANGED. #####\n",
" group_by=group_by,\n",
" )\n",
"dataset"
]
},
{
"cell_type": "markdown",
"id": "27aa0f4d-0b20-4b0d-b3ed-7ea46d33497b",
"metadata": {
"tags": []
},
"source": [
"First thing you probably noticed is that whilst only one line changed the load time dropped to sub-seconds!\n",
"The second thing you probably noticed is if you look at one of the `data variables` by clicking on the database icon as before, there is no data but instead there is a diagram which shows you the __Dask Chunks__ for each measurement. It's really fast because it didn't actually load any data!\n",
"\n",
"When `datatcube` has `dask_chunks` specified it switches from creating `xarrays` to instead use `dask.arrays` in the backend and `lazy loads` them - this means that __no data is loaded until used__. If you look at one of the data variables you will see it now has `dask.array<chunksize=(....)>` rather than values and the cylinder icon will show the Array _and_ Chunk parameters along with some statistics, not actual data.\n",
"\n",
"The `datacube.load()` has used the `dask.Delayed` interface which will not perform any `tasks` (Dask's name for `calculations`) until the _result_ of the `task` is actually required. We'll load the data in a moment but first let's take a look at the parameters in that pretty visualisation. Click on the cylinder for the `red` Data variables and look at the table and the figure. It should look similar to the image below.\n",
"\n",
"<img src=\"../../resources/dask_array_example_small.png\">\n",
"\n",
"Looking at this image (yours may be different), you can see that:\n",
" 1. The Array is `221.92 kiB` in total size and is broken into Chunks which have size `110.96 kiB`\n",
" 2. The Array shape is `(2, 375, 303) (time, y, x)` but each chunk is `(1,375,303)` because we specified the `time` dimension should have chunks of length `1`.\n",
" 3. There are `2` chunk tasks, one for each time slice and in this instance, and one graph layer. More complex calculations will have more layers in the graph.\n",
" 4. The Array type is `uint16` and is split up into chunks which are `numpy.ndarrays`.\n",
" \n",
"The chunking has split the array loading into two Chunks. __Dask can execute these in parallel.__\n",
"\n",
"We can look at the delayed tasks and how they will be executed by visualising the task graph for one of the variables. We'll use the red band measurement."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "e7751aa4-ef91-437d-afcf-9efebfd0fdcb",
"metadata": {
"tags": []
},
"outputs": [],
"source": [
"dataset[alias['red']].data.visualize()"
]
},
{
"cell_type": "markdown",
"id": "39f79fb7-b570-4db0-b000-25d73b296413",
"metadata": {},
"source": [
"Details on the task graph can be found in the dask user guide but what's clear is you have two independent paths of execution which produce one time slice each (0,0,0) and (1,0,0). These are the two chunks that that full array has been split into."
]
},
{
"cell_type": "markdown",
"id": "8fcb46d3-a854-460e-bfd3-1d2681dab13c",
"metadata": {},
"source": [
"To retrieve the actual data we need to `compute()` the result, this will cause all the delayed tasks to be executed for the variable we are computing. Let's `compute()` the red variable."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "d29ecc2f-7a93-4006-9c42-bc1222e39d69",
"metadata": {
"tags": []
},
"outputs": [],
"source": [
"%%time\n",
"actual_red = dataset[alias['red']].compute()\n",
"actual_red"
]
},
{
"cell_type": "markdown",
"id": "197675fa-14b0-4c63-a1b1-fb83ef12df72",
"metadata": {},
"source": [
"As you can see we now have actual data (there are real numbers, not just Dask arrays). You can do the same thing for all arrays in the dataset in one go by computing the dataset itself."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "9547f92d-0546-4c03-ae43-56b0e54b2cc1",
"metadata": {
"tags": []
},
"outputs": [],
"source": [
"%%time\n",
"actual_dataset = dataset.compute()\n",
"actual_dataset"
]
},
{
"cell_type": "markdown",
"id": "e212dd20-baa1-4e62-a7b0-fe4cd410f60b",
"metadata": {},
"source": [
"## The impact of dask on ODC\n",
"\n",
"From the above we can see that specifying `dask_chunks` in `datacube.load()` splits up the `load()` operation into a set of `chunk` shaped arrays and `delayed` _tasks_. Dask can now perform those tasks in _parallel_. Dask will only _compute_ the results for those parts of the data we are using but we can force the computation of all the `delayed` _tasks_ using `compute()`.\n",
"\n",
"There is a _lot_ more opportunity than described in this simple example but let's just focus on the impact of dask on ODC for this simple case.\n",
"\n",
"The time period and ROI are far too small to be interesting so let's change our time range to a few months of data."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "854dacbe-a68a-4b93-820f-2330da2b2510",
"metadata": {
"tags": []
},
"outputs": [],
"source": [
"set_time = (set_time[0], parse(set_time[0]) + relativedelta(months=6))\n",
"# set_time = (\"2021-01-01\", \"2021-06-30\")\n",
"set_time"
]
},
{
"cell_type": "markdown",
"id": "a277a369-9d56-4396-bcd9-b968efc0779c",
"metadata": {},
"source": [
"We skip loading this longer time range (larger data selection) without dask because it can take many minutes and may use more than the available memory in the Jupyter node.\n",
"\n",
"Let's enable dask and then do the load. We're chunking by time (length one) so dask will be able to load each time slice in parallel. The data variables are also independent so will be done in parallel as well."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "2d57bd8a-84a2-4ec6-a7d7-974c84c2524f",
"metadata": {
"tags": []
},
"outputs": [],
"source": [
"if client is None:\n",
" cluster = LocalCluster()\n",
" client = Client(cluster)\n",
" configure_s3_access(aws_unsigned=False, requester_pays=True, client=client);\n",
" display(notebook_utils.localcluster_dashboard(client=client, server=easi.hub))\n",
"else:\n",
" client.restart()"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "6a2aa626-f2ae-49af-84bb-bcc189ed78db",
"metadata": {
"tags": []
},
"outputs": [],
"source": [
"%%time\n",
"\n",
"chunks = {\"time\":1}\n",
"\n",
"dataset = None # clear results from any previous runs\n",
"dataset = dc.load(\n",
" product=product,\n",
" x=study_area_lon,\n",
" y=study_area_lat,\n",
" time=set_time,\n",
" measurements=measurements,\n",
" resampling=resampling,\n",
" output_crs=set_crs,\n",
" resolution=set_resolution,\n",
" dask_chunks = chunks, ###### THIS IS THE ONLY LINE CHANGED. #####\n",
" group_by=group_by,\n",
" )\n",
"dataset"
]
},
{
"cell_type": "markdown",
"id": "3844e75c-2f25-41cb-bcc5-b3f1dc0c01f6",
"metadata": {},
"source": [
"Woah!! that was fast - but we didn't actually compute anything so no load has occurred and all tasks are pending.\n",
"Open up the Data Variables, click the stacked cylinders and take a look at the delayed task counts. These exist for every variable.\n",
"\n",
"Let's visualise the _task graph_ for the `red` band."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "5584957c-10b1-470c-b5d3-135ed773c404",
"metadata": {
"tags": []
},
"outputs": [],
"source": [
"dataset[alias['red']].data.visualize()"
]
},
{
"cell_type": "markdown",
"id": "3444fb54-be18-4732-b58e-d51101f3b2bc",
"metadata": {},
"source": [
"Well that's not as useful, is it!\n",
"\n",
"You should just be able to make out that each of the _chunks_ are able to independently `load()`. `time` _chunk_ is length 1 so these are individual times. This holds true for all the bands so dask can spread these out across multiple threads.\n",
"\n",
"> __Tip__: Visualising task graphs is less effective as your task graph complexity increases. You may need to use simpler examples to see what is going on."
]
},
{
"cell_type": "markdown",
"id": "9800e8fd-21dc-4c3b-ba58-a09ed309cf3b",
"metadata": {},
"source": [
"Let's get the actual data"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "3b3e0a7b-1147-4cdc-901a-8b247ff01864",
"metadata": {
"tags": []
},
"outputs": [],
"source": [
"%%time\n",
"actual_dataset = dataset.compute()\n",
"actual_dataset"
]
},
{
"cell_type": "markdown",
"id": "977ce13e-3ba3-45ba-8122-97031a3871b3",
"metadata": {},
"source": [
"How fast this step is will depend on how many cores are in your Jupyter notebook's local cluster. In real world scenarios, an 8-core cluster the `datacube.load()` this may take between 1/4 or 1/6 of the time compared to without `dask` (not shown) depending on many factors. This is great!\n",
"\n",
"Why not 1/8 of the time?\n",
"\n",
"Dask has overheads, and `datacube.load()` itself is IO limited. There are all sorts of things that result in limits and part of the art of parallel computing is tuning your algorithm to reduce the impact of these and achieve greater performance. As we scale up this example we'll explore some of these.\n",
"\n",
"> __Tip__: recent updates to Dask have greatly improved performance and we are now seeing more substantial performance gains, more in line with the increase in cores.\n",
">\n",
"> Do not always expect 8x as many cores to produce 8x the speed up. Algorithms can be tuned to perform better (or worse) as scale increases. This is part of the art of parallel programming. Dask does it's best, and you can often do better."
]
},
{
"cell_type": "markdown",
"id": "73b9c436-50cd-4768-9ad9-0a769c805777",
"metadata": {},
"source": [
"## Exploiting delayed tasks\n",
"\n",
"Now let's repeat the full example, with NDVI calculation and masking, in a single cell with `dask` and `compute` to load the data in. We get the total time for later comparison.\n",
"\n",
"Most of the time (not shown) is in the data load and the NDVI calculation is < 1 second.\n",
"\n",
"To ensure comparable timings, we will `.restart()` the Dask cluster. This makes sure that we aren't just seeing performance gains for data caching.\n",
"\n",
"> __Note__ that this may show some `Restarting worker` warnings. That is ok and it is just telling you that each of the four workers in the cluster are restarting."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "022ffe6f-149d-4738-a489-1ebe65d686ed",
"metadata": {
"tags": []
},
"outputs": [],
"source": [
"if client is None:\n",
" cluster = LocalCluster()\n",
" client = Client(cluster)\n",
" configure_s3_access(aws_unsigned=False, requester_pays=True, client=client);\n",
" display(notebook_utils.localcluster_dashboard(client=client, server=easi.hub))\n",
"else:\n",
" client.restart()"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "28df278f-c09c-43ed-b6dd-e6c333a45a9c",
"metadata": {
"tags": []
},
"outputs": [],
"source": [
"%%time\n",
"\n",
"chunks = {\"time\":1}\n",
"\n",
"dataset = None # clear results from any previous runs\n",
"dataset = dc.load(\n",
" product=product,\n",
" x=study_area_lon,\n",
" y=study_area_lat,\n",
" time=set_time,\n",
" measurements=measurements,\n",
" resampling=resampling,\n",
" output_crs=set_crs,\n",
" resolution=set_resolution,\n",
" dask_chunks = chunks, \n",
" group_by=group_by,\n",
" )\n",
"actual_dataset = dataset.compute() ### Compute the dataset ###\n",
"\n",
"# Identify pixels that don't have cloud, cloud shadow or water\n",
"cloud_free_mask = masking.make_mask(actual_dataset[qa_band], **qa_mask)\n",
"\n",
"# Apply the mask\n",
"cloud_free = actual_dataset.where(cloud_free_mask)\n",
"\n",
"# Calculate the components that make up the NDVI calculation\n",
"band_diff = cloud_free[alias['nir']] - cloud_free[alias['red']]\n",
"band_sum = cloud_free[alias['nir']] + cloud_free[alias['red']]\n",
"# Calculate NDVI and store it as a measurement in the original dataset ta da\n",
"ndvi = None\n",
"ndvi = band_diff / band_sum"
]
},
{
"cell_type": "markdown",
"id": "f1c86f18-2573-447e-95e1-c03c36064744",
"metadata": {},
"source": [
"Quicker but can do better..."
]
},
{
"cell_type": "markdown",
"id": "6c37f9bb-b82f-49f8-8d29-b45b42398821",
"metadata": {},
"source": [
"## Data and computational locality\n",
"\n",
"When `compute()` is called `dask` not only executes all the tasks but it consolidates all the distributed chunks back into a normal array on the client machine - in this case the notebook's kernel. In the previous cell we have two variables that both refer to the data we are loading:\n",
"1. _dataset_ refers to the `delayed` version of the data. The `delayed` _tasks_ and the _chunks_ that make it up will be __on the cluster__\n",
"2. _actual_dataset_ refers to the actual array in the notebook kernel memory after execution of the _tasks_. The _actual_dataset_ is a complete array in memory in the notebook kernel (__on the _client___).\n",
"\n",
"So in the previous cell everything _after_ the `actual_dataset = dataset.compute()` line is computed in the Jupyter kernel and doesn't use the dask cluster at all for computation.\n",
"\n",
"If we shift the location of this `compute()` call we can perform more _tasks_ in parallel on the dask cluster. \n",
"\n",
"> __Tip__: Locality is an important concept and applies to both data and computation\n",
"\n",
"Now let's repeat the load and NDVI calculation but this time rather than `compute()` on the full `dataset` we'll run the compute at the cloud masking step (`cloud_free = dataset.where(cloud__free_mask).compute()`) so the masking operation can be performed in parallel. Let's see what the impact is...\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "de47a1f3-f6e3-40eb-ac45-69fbeb4dbd2c",
"metadata": {
"tags": []
},
"outputs": [],
"source": [
"if client is None:\n",
" cluster = LocalCluster()\n",
" client = Client(cluster)\n",
" configure_s3_access(aws_unsigned=False, requester_pays=True, client=client);\n",
" display(notebook_utils.localcluster_dashboard(client=client, server=easi.hub))\n",
"else:\n",
" client.restart()"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "dbe23d9a-6db9-476e-a1a6-39da48eadb32",
"metadata": {
"tags": []
},
"outputs": [],
"source": [
"%%time\n",
"\n",
"chunks = {\"time\":1}\n",
"\n",
"dataset = None # clear results from any previous runs\n",
"dataset = dc.load(\n",
" product=product,\n",
" x=study_area_lon,\n",
" y=study_area_lat,\n",
" time=set_time,\n",
" measurements=measurements,\n",
" resampling=resampling,\n",
" output_crs=set_crs,\n",
" resolution=set_resolution,\n",
" dask_chunks = chunks, \n",
" group_by=group_by,\n",
" )\n",
"\n",
"# Identify pixels that are either \"valid\", \"water\" or \"snow\"\n",
"cloud_free_mask = masking.make_mask(dataset[qa_band], **qa_mask)\n",
"\n",
"# Apply the mask\n",
"cloud_free = dataset.where(cloud_free_mask).compute() ### COMPUTE MOVED HERE ###\n",
"\n",
"# Calculate the components that make up the NDVI calculation\n",
"band_diff = cloud_free[alias['nir']] - cloud_free[alias['red']]\n",
"band_sum = cloud_free[alias['nir']] + cloud_free[alias['red']]\n",
"# Calculate NDVI and store it as a measurement in the original dataset ta da\n",
"ndvi = None\n",
"ndvi = band_diff / band_sum\n",
"actual_ndvi = ndvi"
]
},
{
"cell_type": "markdown",
"id": "a9b429a6-3cce-41f5-9e59-8455b2740ea8",
"metadata": {},
"source": [
"A few seconds quicker but not that different. This isn't too surprising since the masking operation is pretty quick (it's all numpy) and the data load is the bulk of the processing.\n",
"\n",
"Dask can see the entire task graph for both load and mask computation. As a result _some_ of the computation can be performed concurrently with file IO, and CPUs are busier as a result, so it will be slightly faster in practice but with IO dominating we won't see much overall improvement.\n",
"\n",
"### With compute on the algorithm\n",
"\n",
"Perhaps doing more of the calculation on the cluster will help. Let's also move `ndvi.compute()` so the entire calculation is done on the cluster and only the final result returned to the client."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "2c1fe1c6-59cd-40ee-942f-983e1e2b3b4d",
"metadata": {
"tags": []
},
"outputs": [],
"source": [
"if client is None:\n",
" cluster = LocalCluster()\n",
" client = Client(cluster)\n",
" configure_s3_access(aws_unsigned=False, requester_pays=True, client=client);\n",
" display(notebook_utils.localcluster_dashboard(client=client, server=easi.hub))\n",
"else:\n",
" client.restart()"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "9b8f58e6-5b88-4fa1-9ed9-e0fa302b0c2f",
"metadata": {
"tags": []
},
"outputs": [],
"source": [
"%%time\n",
"\n",
"chunks = {\"time\":1}\n",
"\n",
"dataset = None # clear results from any previous runs\n",
"dataset = dc.load(\n",
" product=product,\n",
" x=study_area_lon,\n",
" y=study_area_lat,\n",
" time=set_time,\n",
" measurements=measurements,\n",
" resampling=resampling,\n",
" output_crs=set_crs,\n",
" resolution=set_resolution,\n",
" dask_chunks = chunks, \n",
" group_by=group_by,\n",
" )\n",
"\n",
"# Identify pixels that don't have cloud, cloud shadow or water\n",
"cloud_free_mask = masking.make_mask(dataset[qa_band], **qa_mask)\n",
"\n",
"# Apply the mask\n",
"cloud_free = dataset.where(cloud_free_mask)\n",
"\n",
"# Calculate the components that make up the NDVI calculation\n",
"band_diff = cloud_free[alias['nir']] - cloud_free[alias['red']]\n",
"band_sum = cloud_free[alias['nir']] + cloud_free[alias['red']]\n",
"# Calculate NDVI and store it as a measurement in the original dataset ta da\n",
"ndvi = None\n",
"ndvi = band_diff / band_sum\n",
"actual_ndvi = ndvi.compute() ### COMPUTE MOVED HERE ###"
]
},
{
"cell_type": "markdown",
"id": "7bbd24f8-d3e5-4abd-bc59-067269af1284",
"metadata": {},
"source": [
"Now we are seeing a huge difference!\n",
"\n",
"You may be thinking \"Hold on a sec, the NDVI calculation is pretty quick in this example with such a small dataset, why such a big difference?\" - and you'd be right. There is more going on.\n",
"\n",
"Remember that `dataset` is a _task graph_ with `delayed` tasks waiting to be executed __when the result is required__. In the example `dataset`, there are many data variables available but _only 3 are used_ to produce the `ndvi` (`qa_band`, `red` and `nir`). As a result _`dask` doesn't load the other variables_ and because computation time in this case is mostly IO related the execution time is a *lot* faster.\n",
"\n",
"### With selected measurements\n",
"\n",
"Of course we can save `dask` the trouble of figuring this out on our behalf and only `load()` the `measurements` we need in the first place. Let's check that now, we should see a similar performance figure."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "6ee76d53-9b81-4956-b5ea-cde3f9cd18b7",
"metadata": {
"tags": []
},
"outputs": [],
"source": [
"if client is None:\n",
" cluster = LocalCluster()\n",
" client = Client(cluster)\n",
" configure_s3_access(aws_unsigned=False, requester_pays=True, client=client);\n",
" display(notebook_utils.localcluster_dashboard(client=client, server=easi.hub))\n",
"else:\n",
" client.restart()"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "b5c8ff2b-d2a7-4c2a-ada4-da56ee1b71bf",
"metadata": {},
"outputs": [],
"source": [
"%%time\n",
"\n",
"chunks = {\"time\":1}\n",
"measurements = [alias[x] for x in ['qa_band', 'red', 'nir']]\n",
"\n",
"dataset = None # clear results from any previous runs\n",
"dataset = dc.load(\n",
" product=product,\n",
" x=study_area_lon,\n",
" y=study_area_lat,\n",
" time=set_time,\n",
" measurements=measurements,\n",
" resampling=resampling,\n",
" output_crs=set_crs,\n",
" resolution=set_resolution,\n",
" dask_chunks = chunks, \n",
" group_by=group_by,\n",
" )\n",
"\n",
"# Identify pixels that don't have cloud, cloud shadow or water\n",
"cloud_free_mask = masking.make_mask(dataset[qa_band], **qa_mask)\n",
"# Apply the mask\n",
"cloud_free = dataset.where(cloud_free_mask)\n",
"\n",
"# Calculate the components that make up the NDVI calculation\n",
"band_diff = cloud_free[alias['nir']] - cloud_free[alias['red']]\n",
"band_sum = cloud_free[alias['nir']] + cloud_free[alias['red']]\n",
"# Calculate NDVI and store it as a measurement in the original dataset ta da\n",
"ndvi = None\n",
"ndvi = band_diff / band_sum\n",
"actual_ndvi = ndvi.compute()"
]
},
{
"cell_type": "markdown",
"id": "e6669d68-5d23-4769-83e1-4f1a06ada5fe",
"metadata": {},
"source": [
"Pretty similar as expected, but again, a slight improvement because now there are less overheads and a smaller task graph.\n",
"Now it can pay to give `dask` a hand and not have the _task graph_ cluttered with tasks you are not going to use. Still it's nice to see that `dask` can save you some time by only computing what is required when you need it."
]
},
{
"cell_type": "markdown",
"id": "217f4d21-1e94-4891-95fe-0200f940224f",
"metadata": {},
"source": [
"# A quick check on the task graph\n",
"\n",
"For completeness we will take a look at the _task graph_ for the full calculation, all the way to the NDVI result. Given the complexity of the full graph we'll simplify it to 2 time observations like we did when the task graph was introduced previously.\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "e722aa50-a7c9-46c7-9bd1-ee8ebea54f10",
"metadata": {},
"outputs": [],
"source": [
"set_time = (set_time[0], parse(set_time[0]) + relativedelta(weeks=3))\n",
"# set_time = (\"2021-01-01\", \"2021-01-14\")\n",
"set_time"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "a7d70e13-6613-4492-b81c-c620024b485f",
"metadata": {
"tags": []
},
"outputs": [],
"source": [
"if client is None:\n",
" cluster = LocalCluster()\n",
" client = Client(cluster)\n",
" configure_s3_access(aws_unsigned=False, requester_pays=True, client=client);\n",
" display(notebook_utils.localcluster_dashboard(client=client, server=easi.hub))\n",
"else:\n",
" client.restart()"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "6b9b4922-71f3-418b-81af-f3e0c0d15fd9",
"metadata": {},
"outputs": [],
"source": [
"%%time\n",
"dataset = None # clear results from any previous runs\n",
"measurements = [alias[x] for x in ['qa_band', 'red', 'nir']]\n",
"dataset = dc.load(\n",
" product=product,\n",
" x=study_area_lon,\n",
" y=study_area_lat,\n",
" time=set_time,\n",
" measurements=measurements,\n",
" resampling=resampling,\n",
" output_crs=set_crs,\n",
" resolution=set_resolution,\n",
" dask_chunks = chunks, \n",
" group_by=group_by,\n",
" )\n",
"\n",
"# Identify pixels that don't have cloud, cloud shadow or water\n",
"cloud_free_mask = masking.make_mask(dataset[qa_band], **qa_mask)\n",
"# Apply the mask\n",
"cloud_free = dataset.where(cloud_free_mask)\n",
"\n",
"# Calculate the components that make up the NDVI calculation\n",
"band_diff = cloud_free[alias['nir']] - cloud_free[alias['red']]\n",
"band_sum = cloud_free[alias['nir']] + cloud_free[alias['red']]\n",
"# Calculate NDVI and store it as a measurement in the original dataset ta da\n",
"ndvi = None\n",
"ndvi = band_diff / band_sum"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "9dabc0f2-5899-4900-b73c-6a46a3286ad0",
"metadata": {},
"outputs": [],
"source": [
"ndvi.data.visualize()"
]
},
{
"cell_type": "markdown",
"id": "0c27734b-602e-4a3d-9758-d6eca5e7cb85",
"metadata": {},
"source": [
"The computation flows from __bottom to top__ in the _task graph_. You can see there are two main paths, one for each time (since the time chunk is length 1). You can also see the three data sources are loaded independently. After that it gets a little more difficult to follow but you can see `qa_band` being used to produce the mask (_and\\__, _eq_). Then combined via the `where` function with other two datasets. Then finally the NDVI calculation - a sub, add and divide (truediv).\n",
"\n",
"Dask has lots of internal optimizations that it uses to help identify the dependencies and parallel components of a task graph. Sometimes it will reorder or prune operations where possible to further optimise (for example, not loading _data variables_ that aren't used in the NDVI calculation).\n",
"\n",
"> __Tip__: The _task graph_ can be complex but it is a useful tool in understanding your algorithm and how it scales."
]
},
{
"cell_type": "markdown",
"id": "41b3d8d9-1a40-4d17-a5a1-601ba5dda563",
"metadata": {
"tags": []
},
"source": [
"## Be a good dask user - Clean up the cluster resources"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "48eeb0ee-6888-4362-8615-19cf217e80c0",
"metadata": {
"tags": []
},
"outputs": [],
"source": [
"client.close()\n",
"\n",
"cluster.close()"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "556ea4f3-e6e5-4dd7-9326-8d690e068093",
"metadata": {},
"outputs": [],
"source": []
}
],
"metadata": {
"kernelspec": {
"display_name": "Python 3 (ipykernel)",
"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.6"
}
},
"nbformat": 4,
"nbformat_minor": 5
}
Reference in New Issue View Git Blame Copy Permalink