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:
I hope this issue report is useful. Rob W said that Kirill888 was looking at refactoring the geometry codes & might have some insights.

I have done a little bit more tracing of the python/c++ calls. Calls to proj_to_wkt are slow because of overheads related to creating sqlite database contexts. The trace image is attached.
![proj_as_wkt_profile](https://user-images.githubusercontent.com/94149433/154227834-7174e625-f5a3-4624-a632-6cec8f03c723.svg)
 
Related: https://github.com/pyproj4/pyproj/issues/782
Related:
- https://pyproj4.github.io/pyproj/stable/advanced_examples.html#optimizing-single-threaded-applications
- https://pyproj4.github.io/pyproj/stable/api/global_context.html#global-context
> I hope this issue report is useful. Rob W said that Kirill888 was looking at refactoring the geometry codes & might have some insights.

@csiro-dave thank you for this detailed analysis. I'm currently working on extracting that part of datacube into a separate library (https://github.com/opendatacube/odc-geo). There has been some significant clean up and refactor in there. I will create an issue in that repo. And will address it there first.

We are already caching pyproj CRS object as well as epsg code, so adding wkt to that cache should be relatively straightforward. 

https://github.com/opendatacube/datacube-core/blob/ac9a466e4c817f1eeb3ca2c64bb2b352df3fa344/datacube/utils/geometry/_base.py#L119-L122

Another quick to implement option is to use `epsg` code for hashing purposes when available and only fallback to wkt when isn't.

https://github.com/opendatacube/datacube-core/blob/ac9a466e4c817f1eeb3ca2c64bb2b352df3fa344/datacube/utils/geometry/_base.py#L259-L260

If you can modify above to something like:

```python
    def __hash__(self) -> int:
        if self._epsg is not None:
            return self._epsg
        return hash(self.to_wkt())
```
and rerun your tests and see if that helps.

@snowman2 thanks for the links, I didn't realize that CRS objects and transformers had threading constraints. We should probably switch to thread-local caches then.






EPSG codes are likely worse than WKT. See #1223

I recommend using the `srs` attribute if you need performance: https://pyproj4.github.io/pyproj/stable/api/crs/crs.html#pyproj.crs.CRS.srs
> EPSG codes are likely worse than WKT. See #1223

thanks for the info, but we always have them extracted at construction and they are cached unlike wkt. In the Dask context there is really only one CRS object, so this can be solved with caching very well. I don't think scenarios with a large number of different CRS objects are common, so we should try improve caching first and then look at pyproj context.


> I didn't realize that CRS objects and transformers had threading constraints. We should probably switch to thread-local caches then.

They shouldn't have issues now: https://github.com/pyproj4/pyproj/issues/782. It should create a new CRS object for each thread for you.
> They shouldn't have issues now: [pyproj4/pyproj#782](https://github.com/pyproj4/pyproj/issues/782). It should create a new CRS object for each thread for you.

As in, one global `pyproj.CRS` object will create a thread-local internal object when accessed from a different thread than where it was originally created?

> As in, one global pyproj.CRS object will create a thread-local internal object when accessed from a different thread than where it was originally created?

Correct: https://github.com/pyproj4/pyproj/pull/793

EDIT: requires pyproj 3.1+
Thank you snowman2 and Kirill888! It looks like pyproj.set_use_global_context(True) makes a huge difference. There is an overhead to re-creating the database context at the proj level. At the pyproj level, the thread protections are perhaps a little aggressive if they need to create a new database context for every method call on an object. The example above that took 0.5 seconds is now 0.001 seconds. The 16 second persist call now takes 0.7 seconds. Hooray!

```python
from pyproj import CRS
from pyproj.enums import WktVersion
import timeit
pyproj.set_use_global_context(True)
crs = CRS.from_epsg(4326)
timeit.timeit(lambda:crs.to_wkt(),number=100)
```

Some thoughts on the hash/equals of the CRS. I don't think GeoBox needs to include CRS in its hash function. If it has the same shape and transform, it is most likely going to be the same. In the rare case they are different, it will call equals a few extra times. So it might make sense to change
```python
    def __hash__(self):
        return hash((*self._shape, self._crs, self._affine))
```
to
```python
    def __hash__(self):
        return hash((*self._shape, self._affine))
```

I don't currently have an environment set up to test the datacube source changes, but they look like they could work too.
regarding hash change on geobox class, good point, but maybe replace `self._crs` with `self._crs.epsg` then hash collision will only happen for non-epsg geoboxes of the same shape/transform. `crs.epsg` is either `None` or `int` and is just a plain python property lookup without calling into pyproj.
Thanks Alan, useful to know.

Resolving patch:
diff --git a/datacube/utils/geometry/_base.py b/datacube/utils/geometry/_base.py
--- a/datacube/utils/geometry/_base.py
+++ b/datacube/utils/geometry/_base.py
@@ -116,10 +116,25 @@ def from_points(p1: Tuple[float, float],
                                    (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 @@ def __str__(self) -> str:
         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 @@ def coordinates(self) -> Dict[str, Coordinate]:
         """
         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
--- a/tests/test_geometry.py
+++ b/tests/test_geometry.py
@@ -43,12 +43,12 @@
     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]]}