working with rasterio.Env / gdal configuration options

[artttt]

Im trying to understand how (or if) rioxarray works with rasterio.Env aka gdal configuration.

For demo use im using a vrt that points to a remote tif file
I then set CPL_VSIL_CURL_ALLOWED_EXTENSIONS=".vrt" so that the vrt can be read but the tif cant. The idea is that open_rasterio can work on the vrt but when i try to compute its opening the tif.

If I get an error then it is clear that the gdal configuration is working.

It seemed like a definitive way to test what what environment is seen on the workers separately to locally. Hopefully that makes sense.

I was expecting (hoping) that using rioxarray.open inside a rasterio.Env context would mean that context would magically be used where ever rasterio ended up being called by rioxarray to open that dataset.

Im not sure if im doing something wrong or just expecting too much. Any pointers?

Below is a set of examples. They need to each be run separately to avoid gdal caching.

Thanks

url = '/vsicurl/https://github.com/rasterio/rasterio/blob/main/tests/data/world.rgb.tif?raw=true'
vrt = 'demo.vrt'

from osgeo import gdal
ds = gdal.BuildVRT(vrt,url)
ds.FlushCache()
ds = None

#rasterio - testing normal behaviour
#this works - as expected
vrt = 'demo.vrt'
import rasterio
with rasterio.Env(CPL_VSIL_CURL_ALLOWED_EXTENSIONS=".vrt .tif" ):
    with rasterio.open(vrt) as src:
        print (src.profile)
        w = src.read(1)
        print (w)

#rasterio - testing normal behaviour
#this reads the profile from the vrt and then errors on the src.read line - as expected
vrt = 'demo.vrt'
import rasterio
with rasterio.Env(CPL_VSIL_CURL_ALLOWED_EXTENSIONS=".vrt" ):
    with rasterio.open(vrt) as src:
        print (src.profile)
        w = src.read(1)
        print (w)

#rioxarray - compute outside of env context
#this doesnt error
vrt = 'demo.vrt'
import rasterio
import rioxarray
with rasterio.Env(CPL_VSIL_CURL_ALLOWED_EXTENSIONS=".vrt" ):
    da = rioxarray.open_rasterio(vrt, chunks='auto')
print (da)
print (da.compute())

#rioxarray - compute inside of env context
#this errors on the compute line
vrt = 'demo.vrt'
import rasterio
import rioxarray
with rasterio.Env(CPL_VSIL_CURL_ALLOWED_EXTENSIONS=".vrt" ):
    da = rioxarray.open_rasterio(vrt, chunks='auto')
    print (da)
    print (da.compute())

#rioxarray - using distributed workers compute inside of env context
#this doesnt error
from dask.distributed import Client,LocalCluster
cluster = LocalCluster(n_workers=3, threads_per_worker=1)
client = Client(cluster)
vrt = 'demo.vrt'
import rasterio
import rioxarray
with rasterio.Env(CPL_VSIL_CURL_ALLOWED_EXTENSIONS=".vrt" ):
    da = rioxarray.open_rasterio(vrt, chunks='auto')
    print (da)
    print (da.compute())
client.shutdown()

#rioxarray - using distributed workers compute outside of env context
#this doesnt error
from dask.distributed import Client,LocalCluster
cluster = LocalCluster(n_workers=3, threads_per_worker=1)
client = Client(cluster)
vrt = 'demo.vrt'
import rasterio
import rioxarray
with rasterio.Env(CPL_VSIL_CURL_ALLOWED_EXTENSIONS=".vrt" ):
    da = rioxarray.open_rasterio(vrt, chunks='auto')
print (da)
print (da.compute())
client.shutdown()

#rioxarray - setting the distributed workers environment with environment variables 
#this errors on the compute line - as expected
from dask.distributed import Client,LocalCluster,Environ
cluster = LocalCluster(n_workers=3, threads_per_worker=1)
client = Client(cluster)
client.register_worker_plugin(Environ({'CPL_VSIL_CURL_ALLOWED_EXTENSIONS':".vrt" }), nanny=True)
vrt = 'demo.vrt'
import rasterio
import rioxarray
da = rioxarray.open_rasterio(vrt, chunks='auto')
print (da)
print (da.compute())
client.shutdown()

[snowman2]

Im trying to understand how (or if) rioxarray works with rasterio.Env aka gdal configuration.

This hasn't been looked into too deeply. It may or may not work as you suggest. However, if it were to work, I imagine all code has to be executed inside of the Env block.

[artttt]

Thanks Alan,
Ive seen a few example of rasterio.Env being used with rioxarray.open_rasterio (including in the rioxarray tests) but havent found it mentioned that it wont work in all contexts. Maybe worth a comment in the docs?

I think its quite a common pattern with rasterio to do:

    env = rasterio.Env(
        GDAL_HTTP_MERGE_CONSECUTIVE_RANGES="YES",
        GDAL_DISABLE_READDIR_ON_OPEN="EMPTY_DIR",
        CPL_VSIL_CURL_USE_HEAD=False,
        CPL_VSIL_CURL_ALLOWED_EXTENSIONS="TIF",
    )
    with env:
        with rasterio.open(url) as src:
.................

It would be nice if this pattern worked consistently in rioxarray.

My 3rd from last example is the one that is most counter intuitive in the way it currently functions. I think there is also an argument the the second from last should also work differently.

I had a dig around in the code to see what it would take to make it work.

if rioxarray.open_rasterio kept a copy of the environment by either grabbing rasterio_env_options = rasterio.env.getenv() or having an optional argument rasterio_env_options={}

Then it could wrap the rasterio.open call in a rasterio.Env(**rasterio_env_options) context.
I think that happens in the ThreadURIManager (but i'm not sure if that's the right level to do the wrapping)
https://github.com/corteva/rioxarray/blob/master/rioxarray/_io.py#L184

I guess a similar approach would also be appropriate in RasterioWriter.

[Kirill888]

In odc-stac I had to add explicit capture of configured rasterio environment when constructing Dask graph for data loading

https://github.com/opendatacube/odc-stac/blob/77f8aa117770abfc7cef4478a39302091f1d0bba/odc/stac/_load.py#L193-L206

I think it is "least surprising" way of handling it, but it does require quite a bit of code, especially for handling AWS credentials

https://github.com/opendatacube/odc-stac/blob/develop/odc/stac/_rio.py

and there are some gotchas: like you probably don't want to copy GDAL_DATA over to some remote.

For things like AWS credentials, sometimes you want to get credentials in the main process and then copy them over to the cluster (static credentials/assume role), other times you can ONLY obtain credentials when running on the cluster (machine credentials). And I'm sure there are plenty of tuning settings you might want to configure at the "fleet" level rather than worry about them in code. Configuration space can get pretty complex quick.

[artttt]

Thanks Kirill and Alan,
Thanks for the examples Kirill - As you say quite a bit of code. Im wondering if there is some middle ground that could be implemented usefully in rioxarray?

I sunk a bit of time into it and got something sort of working but ive given up for now.

I will stick with setting environment variables on the 'fleet' of workers for now as it seems to be the only reliable way to get a gdal configuration to work with rioxarray.

[snowman2]

I am not sure what a good solution would look like for rioxarray, but sounds like it will take some investigation to get it right. If you do come up with a solution, we would be happy to discuss it with you to see if it is a good fit.

[artttt]

Thanks Alan,
and thanks for discussing this with me.

For reference in case anyone else wants to have a look at this.

It was quite simple to make a change that works for RasterioWriter. I passed in a dict of options and wrapped the 2 calls to rasterio.open like would be done with rasterio anywhere else.

with rasterio.Env(**self.env_options):
            with rasterio.open(self.raster_path.........................

open_rasterio is what got me stuck. Because the rasterio.open calls are not done with a context manager and the file handles are kept open so i guess the rasterio.Env needs to be kept active. I managed that but it ends up in a situation where there are potentially many active environments, potentially conflicting, potentially exiting out of order, so it gets difficult to guarantee what environment is actually in place when making calls to a rasterio dataset.

I also dont know enough about rasterio and gdal to know what happens if the environment changes while a dataset is open.

The workaround im using for now is to forgo more dynamic use of gdal configuration options and set them up more globally like so:

client.register_worker_plugin(Environ({'CPL_VSIL_CURL_ALLOWED_EXTENSIONS':".vrt" }), nanny=True)