Repository: opendatacube/datacube-core

Issue:
load_ard() task graph creation limited by GeoBox/CRS hash function performance
### Expected behaviour
When I run load_ard() on a large region (the Murray Darling Basin) using distributed dask, I would like to be able to create up to 500,000 tasks for my workflow. That is, fully utilise the dask scheduler’s capacity before I parallelise the problem across multiple dask clusters/tiles/time slices etc.

### Actual behaviour
On my simplified problem, using dask chunking (1x10240x10240) the production of the task graph (6144 tasks) in the distributed dask cluster takes around 16 seconds. On a real problem, I haven’t the patience to wait for it to finish. 

As suggested in the dask manual, I have run the %prun over the persist function to see what is consuming the resources. The step that is taking most of the resources is the cull step of the task graph optimisation. The cull step identifies any tasks that don’t need to be executed and removes them from the task graph. When culling, dask stores the tasks that it identifies in a set. Python sets make use of object hashes to efficiently calculate which objects are in the intersection of those objects used vs those not used. However in the case of the simplified problem, calculating the hash of the CRS consumes most of the processing time (96%). 

Why are the CRS hashes part of the task graph? As child of task dc_load_fmask-to_float, two lazy functions (datacube.api.core.fuse_lazy) include a GeoBox as an argument. The GeoBox is a composite object including a CRS. When GeoBox is hashed, then the hash CRS is required. 

Why is CRS hashing slow? Open data cube CRS is a wrapper around PyProj CRS, which is a wrapper around Proj. ODC CRS considers CRS an immutable object defined by the value of its contents. When calculating the hash, it uses the hash of the WKT projection text. The underlying Proj library can calculate the text very quickly. However, PyProj function call is much slower. I think it might be due to the overhead of transferring data from C++ to Python. Or it could be creating new proj_context for each new CRS object (to handle multi-threading issues). 

The implementation seems well motivated, so I have not been able to bring a solution. Perhaps during the step that generates multiple GeoBox objects, the CRS should be eagerly converted to a WKT for the purposes of the hash. 

### Steps to reproduce the behaviour
I have attached a Jupyter notebook. 
[geobox_hash_performance.zip](https://github.com/opendatacube/datacube-core/files/8077431/geobox_hash_performance.zip)
A collection of major lines are here:
```
query = {
# Site details
'latitude': (-24.5856075286866, -37.6786722912282),
'longitude': (138.566245999366, 152.491485595703),  
'time': ('1988-01-01', '1988-02-29'), 
'output_crs': 'EPSG:4326',
'resolution': (-0.0002698, 0.0002698)
}

dc= datacube.Datacube(app='dc-LS-extract')

# Read in Landsat images and bands
bands=['nbart_blue','nbart_green','nbart_red','nbart_nir','nbart_swir_1','nbart_swir_2']

chunks = {'time': 1, 'latitude': 512*2*10, 'longitude': 512*2*10}

ds = load_ard(dc=dc,
            products=['ga_ls5t_ard_3','ga_ls7e_ard_3', 'ga_ls8c_ard_3'],
            measurements=bands, **query, group_by='solar_day',
             dask_chunks= chunks
             )   
%prun ds.nbart_blue.persist()
773424 function calls (646389 primitive calls) in 15.597 seconds

Ordered by: internal time

ncalls  tottime  percall  cumtime  percall filename:lineno(function)
2664   15.096    0.006   15.128    0.006 {method 'to_wkt' of 'pyproj._crs.Base' objects}
52282/17594    0.056    0.000    3.474    0.000 utils.py:1669(stringify)
10315    0.042    0.000   10.197    0.001 core.py:159(keys_in_tasks)
3    0.035    0.012    5.145    1.715 optimization.py:429(fuse)
```
A small snippet to reproduce the pyproj performance (takes about 0.5 seconds):
```
from pyproj import CRS
from pyproj.enums import WktVersion
import timeit
crs = CRS.from_epsg(4326)
timeit.timeit(lambda:crs.to_wkt(),number=100)
```
A small snippet to eliminate the proj performance (takes 6 milliseconds)
```
#include <iostream>
#include <vector>
#include <string>
#include <chrono>
#include <proj.h>
#include <sqlite3.h>

using namespace std;

int main(void)
{
    const char *wkt;

    std::chrono::steady_clock::time_point begin = std::chrono::steady_clock::now();
    PJ_CONTEXT *m_ctxt = proj_context_create();
    auto crs = proj_create_from_database(m_ctxt, "EPSG", "4326",
                                         PJ_CATEGORY_CRS, false, nullptr);
    const char *const options[] = {"MULTILINE=NO", nullptr};
    for (int i = 0; i < 100; i++)
    { 
        wkt = proj_as_wkt(m_ctxt, crs, PJ_WKT1_GDAL, options);
    }

    std::chrono::steady_clock::time_point end = std::chrono::steady_clock::now();

    printf("%s", wkt);
    std::cout << "Time difference = " << std::chrono::duration_cast<std::chrono::milliseconds>(end - begin).count() << " millisecond" << std::endl;
}
```
### Environment information
Open Data Cube core, version 1.8.6
pyproj info:
    pyproj: 3.2.1
      PROJ: 7.2.1
  data dir: /usr/local/share/proj
user_data_dir: /home/jovyan/.local/share/proj

System:
    python: 3.8.10 (default, Jun  2 2021, 10:49:15)  [GCC 9.4.0]
executable: /env/bin/python
   machine: Linux-4.14.256-197.484.amzn2.x86_64-x86_64-with-glibc2.29
Python deps:
   certifi: 2021.10.08
       pip: 21.3.1
setuptools: 58.5.3
    Cython: 0.29.24



Discussion:
codecov[bot]: # [Codecov](https://codecov.io/gh/opendatacube/datacube-core/pull/1242?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=opendatacube) Report
> Merging [#1242](https://codecov.io/gh/opendatacube/datacube-core/pull/1242?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=opendatacube) (426e411) into [develop](https://codecov.io/gh/opendatacube/datacube-core/commit/2aee3af6da9e4843154efc695c22e35a1d66e704?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=opendatacube) (2aee3af) will **decrease** coverage by `0.00%`.
> The diff coverage is `92.10%`.

[![Impacted file tree graph](https://codecov.io/gh/opendatacube/datacube-core/pull/1242/graphs/tree.svg?width=650&height=150&src=pr&token=TqpPsC1owf&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=opendatacube)](https://codecov.io/gh/opendatacube/datacube-core/pull/1242?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=opendatacube)

```diff
@@             Coverage Diff             @@
##           develop    #1242      +/-   ##
===========================================
- Coverage    93.86%   93.86%   -0.01%     
===========================================
  Files          109      109              
  Lines        10760    10772      +12     
===========================================
+ Hits         10100    10111      +11     
- Misses         660      661       +1     
```


| [Impacted Files](https://codecov.io/gh/opendatacube/datacube-core/pull/1242?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=opendatacube) | Coverage Δ | |
|---|---|---|
| [datacube/utils/geometry/\_base.py](https://codecov.io/gh/opendatacube/datacube-core/pull/1242/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=opendatacube#diff-ZGF0YWN1YmUvdXRpbHMvZ2VvbWV0cnkvX2Jhc2UucHk=) | `99.33% <92.10%> (-0.13%)` | :arrow_down: |

------

[Continue to review full report at Codecov](https://codecov.io/gh/opendatacube/datacube-core/pull/1242?src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=opendatacube).
> **Legend** - [Click here to learn more](https://docs.codecov.io/docs/codecov-delta?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=opendatacube)
> `Δ = absolute <relative> (impact)`, `ø = not affected`, `? = missing data`
> Powered by [Codecov](https://codecov.io/gh/opendatacube/datacube-core/pull/1242?src=pr&el=footer&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=opendatacube). Last update [465c44e...426e411](https://codecov.io/gh/opendatacube/datacube-core/pull/1242?src=pr&el=lastupdated&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=opendatacube). Read the [comment docs](https://docs.codecov.io/docs/pull-request-comments?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=opendatacube).


Patch:
diff --git a/datacube/utils/geometry/_base.py b/datacube/utils/geometry/_base.py
index 43454159..41b906ef 100644
--- a/datacube/utils/geometry/_base.py
+++ b/datacube/utils/geometry/_base.py
@@ -116,10 +116,25 @@ class BoundingBox(_BoundingBox):
                                    (p1[1], p2[1]))
 
 
-@cachetools.cached({})
-def _make_crs(crs_str: str) -> Tuple[_CRS, Optional[int]]:
-    crs = _CRS.from_user_input(crs_str)
-    return (crs, crs.to_epsg())
+def _make_crs_key(crs_spec: Union[str, _CRS]) -> str:
+    if isinstance(crs_spec, str):
+        normed_epsg = crs_spec.upper()
+        if normed_epsg.startswith("EPSG:"):
+            return normed_epsg
+        return crs_spec
+    return crs_spec.to_wkt()
+
+
+@cachetools.cached({}, key=_make_crs_key)
+def _make_crs(crs: Union[str, _CRS]) -> Tuple[_CRS, str, Optional[int]]:
+    if isinstance(crs, str):
+        crs = _CRS.from_user_input(crs)
+    epsg = crs.to_epsg()
+    if epsg is not None:
+        crs_str = f"EPSG:{epsg}"
+    else:
+        crs_str = crs.to_wkt()
+    return (crs, crs_str, crs.to_epsg())
 
 
 def _make_crs_transform_key(from_crs, to_crs, always_xy):
@@ -131,26 +146,6 @@ def _make_crs_transform(from_crs, to_crs, always_xy):
     return Transformer.from_crs(from_crs, to_crs, always_xy=always_xy).transform
 
 
-def _guess_crs_str(crs_spec: Any) -> Optional[str]:
-    """
-    Returns a string representation of the crs spec.
-    Returns `None` if it does not understand the spec.
-    """
-    if isinstance(crs_spec, str):
-        return crs_spec
-    if isinstance(crs_spec, dict):
-        crs_spec = _CRS.from_dict(crs_spec)
-
-    if hasattr(crs_spec, 'to_wkt'):
-        return crs_spec.to_wkt()
-    if hasattr(crs_spec, 'to_epsg'):
-        epsg = crs_spec.to_epsg()
-        if epsg is not None:
-            return 'EPSG:{}'.format(crs_spec.to_epsg())
-
-    return None
-
-
 class CRS:
     """
     Wrapper around `pyproj.CRS` for backwards compatibility.
@@ -160,20 +155,34 @@ class CRS:
 
     __slots__ = ('_crs', '_epsg', '_str')
 
-    def __init__(self, crs_str: Any):
+    def __init__(self, crs_spec: Any):
         """
         :param crs_str: string representation of a CRS, often an EPSG code like 'EPSG:4326'
         :raises: `pyproj.exceptions.CRSError`
         """
-        crs_str = _guess_crs_str(crs_str)
-        if crs_str is None:
-            raise CRSError("Expect string or any object with `.to_epsg()` or `.to_wkt()` method")
-
-        _crs, _epsg = _make_crs(crs_str)
-
-        self._crs = _crs
-        self._epsg = _epsg
-        self._str = crs_str
+        if isinstance(crs_spec, str):
+            self._crs, self._str, self._epsg = _make_crs(crs_spec)
+        elif isinstance(crs_spec, CRS):
+            self._crs = crs_spec._crs
+            self._epsg = crs_spec._epsg
+            self._str = crs_spec._str
+        elif isinstance(crs_spec, _CRS):
+            self._crs, self._str, self._epsg = _make_crs(crs_spec)
+        elif isinstance(crs_spec, dict):
+            self._crs, self._str, self._epsg = _make_crs(_CRS.from_dict(crs_spec))
+        else:
+            _to_epsg = getattr(crs_spec, "to_epsg", None)
+            if _to_epsg is not None:
+                self._crs, self._str, self._epsg = _make_crs(f"EPSG:{_to_epsg()}")
+                return
+            _to_wkt = getattr(crs_spec, "to_wkt", None)
+            if _to_wkt is not None:
+                self._crs, self._str, self._epsg = _make_crs(_to_wkt())
+                return
+
+            raise CRSError(
+                "Expect string or any object with `.to_epsg()` or `.to_wkt()` methods"
+            )
 
     def __getstate__(self):
         return {'crs_str': self._str}
@@ -257,7 +266,7 @@ class CRS:
         return self._str
 
     def __hash__(self) -> int:
-        return hash(self.to_wkt())
+        return hash(self._str)
 
     def __repr__(self) -> str:
         return "CRS('%s')" % self._str
@@ -987,7 +996,6 @@ class GeoBox:
     """
 
     def __init__(self, width: int, height: int, affine: Affine, crs: MaybeCRS):
-        assert is_affine_st(affine), "Only axis-aligned geoboxes are currently supported"
         self.width = width
         self.height = height
         self.affine = affine
@@ -1110,6 +1118,7 @@ class GeoBox:
         """
         dict of coordinate labels
         """
+        assert is_affine_st(self.affine), "Only axis-aligned geoboxes are currently supported"
         yres, xres = self.resolution
         yoff, xoff = self.affine.yoff, self.affine.xoff
 


Test patch:
diff --git a/tests/test_geometry.py b/tests/test_geometry.py
index 1c989e20..235ba464 100644
--- a/tests/test_geometry.py
+++ b/tests/test_geometry.py
@@ -43,12 +43,12 @@ from datacube.utils.geometry._base import (
     bounding_box_in_pixel_domain,
     geobox_intersection_conservative,
     geobox_union_conservative,
-    _guess_crs_str,
     force_2d,
     _align_pix,
     _round_to_res,
     _norm_crs,
     _norm_crs_or_error,
+    _make_crs_key,
 )
 from datacube.testutils.geom import (
     epsg4326,
@@ -1475,9 +1475,9 @@ def test_crs_hash():
 
 
 def test_base_internals():
-    assert _guess_crs_str(CRS("epsg:3577")) == epsg3577.to_wkt()
+    assert _make_crs_key("epsg:3577") == "EPSG:3577"
     no_epsg_crs = CRS(SAMPLE_WKT_WITHOUT_AUTHORITY)
-    assert _guess_crs_str(no_epsg_crs) == no_epsg_crs.to_wkt()
+    assert _make_crs_key(no_epsg_crs.proj) == no_epsg_crs.proj.to_wkt()
 
     gjson_bad = {'type': 'a', 'coordinates': [1, [2, 3, 4]]}
     assert force_2d(gjson_bad) == {'type': 'a', 'coordinates': [1, [2, 3]]}