Propagate env to dask workers with a LayeredEnv

[gjoseph92]

Because dask workers run in separate processes, setting os.environ isn't a reliable way to set GDAL configuration options. stackstac.LayeredEnv is basically the same as rasterio.Env, and lets you embed that configuration into the dask graph, so every worker will use it.

Note that LayeredEnv is not pickleable yet, so this requires gjoseph92/stackstac@d01da76... I'll release a new stackstac version shortly with these changes UPDATE: try with stackstac==0.1.1.

[gjoseph92]

Sorry, I hoped GitHub would be able to render a nicer diff of the notebook... the change is basically doing something like:

base_env = stackstac.DEFAULT_GDAL_ENV.updated(dict(
    GDAL_MAX_RAW_BLOCK_CACHE_SIZE='200000000',
    GDAL_SWATH_SIZE='200000000',
    VSI_CURL_CACHE_SIZE='200000000',
    GDAL_HTTP_UNSAFESSL='YES',
    GDAL_HTTP_COOKIEFILE=os.path.expanduser('~/cookies.txt'),
    GDAL_HTTP_COOKIEJAR=os.path.expanduser('~/cookies.txt'),
))

if s3 and data=='hls':
    #Get NASA Temp AWS Credentials
    s3_cred = requests.get('https://lpdaac.earthdata.nasa.gov/s3credentials').json()
    s3_cred
    
    env = base_env.updated(
        dict(
            AWS_REGION='us-west-2',
            AWS_NO_SIGN_REQUEST='NO',
            AWS_REQUEST_PAYER='REQUESTER',
            AWS_SECRET_ACCESS_KEY=s3_cred['secretAccessKey'],
            AWS_ACCESS_KEY_ID=s3_cred['accessKeyId'],
            AWS_SESSION_TOKEN=s3_cred['sessionToken']
        )
    )

# Default to the StackStac ~.LayeredEnv
if s3==False and data=='hls':
    env = base_env.updated(dict(AWS_NO_SIGN_REQUEST='YES'))

if s3 and (data == 'landsat8' or data == 'sentinel_2a'):
    env = base_env.updated(
        dict(
            AWS_REGION='us-west-2',
            AWS_REQUEST_PAYER='REQUESTER',
            region_name='us-west-2',
            AWS_NO_SIGN_REQUEST='YES',
        )
    )

(Where we use LayeredEnv.updated to build up a new environment, adding in necessary options depending on the data.)

Then the only other things you need to do are:

with env.open, rasterio.open(...) as src:
    ....

(so credentials from the env are applied during the open)

stackstac.stack(..., gdal_env=env)

(actually pass the env into stackstac)

[rmg55]

Awesome - thanks @gjoseph92!