Passing AWS profile credentials to Dask Gateway workers

[rsignell-usgs]

We've had several people show up on our Qhub from different agencies, bringing with them their own cloud credentials, for example, an AWS keypair that allows that IAM role to write to a specified S3 bucket.

They can run aws configure --profile my-profile and specify their keys, which allows the notebook to write to the bucket, but when they try to add a Dask Gateway cluster they get killed workers, and the worker logs show that the workes can't find the specified profile.

How best to communicate these credentials set in the ~/.aws/{credentials file to the Dask Gateway workers?

[rsignell-usgs]

I see that @jkellndorfer uses this approach:

import configparser
def set_aws_credentials(cfile=os.path.join(os.environ['HOME'],'shared','.aws','credentials'),
           profile_name='default', region_name='us-east-1',
           endpoint='s3.amazonaws.com',verbose=False):
    '''Sets the aws credentials if not set already and profilename is default'''
    cp = configparser.ConfigParser()
    cp.read(cfile)
    os.environ['aws_access_key_id'.upper()]=cp[profile_name]['aws_access_key_id']	
    os.environ['aws_secret_access_key'.upper()]=cp[profile_name]['aws_secret_access_key']	
    os.environ['aws_profile'.upper()]=profile_name
    os.environ['aws_default_profile'.upper()]=profile_name
    os.environ['aws_s3_region'.upper()]=region_name
    os.environ['aws_s3_endpoint'.upper()]=endpoint
    os.environ['aws_default_region'.upper()]=region_name
    if verbose:
        print('export {}={}'.format('aws_access_key_id'.upper(),cp[profile_name]['aws_access_key_id']	))
        print('export {}={}'.format('aws_secret_access_key'.upper(),cp[profile_name]['aws_secret_access_key']	))

Is this the best method?

[jkellndorfer]

That's the one part. I also propagate the variables to the cluster with

def set_env(k,v):
    import os
    os.environ[k]=v
    
client.run(set_env,'AWS_ACCESS_KEY_ID',os.environ['AWS_ACCESS_KEY_ID'])
client.run(set_env,'AWS_SECRET_ACCESS_KEY',os.environ['AWS_SECRET_ACCESS_KEY'])
client.run(set_env,'AWS_DEFAULT_REGION',os.environ['AWS_DEFAULT_REGION'])
client.run(set_env,'AWS_S3_ENDPOINT',os.environ['AWS_S3_ENDPOINT'])
client.run(set_env,'AWS_PROFILE',os.environ['AWS_PROFILE'])

With a loop over all AWS variables this could be

for key in [x for x in os.environ if x.find('AWS')>-1]:
     client.run(set_env,key,os.environ[key])

[rsignell-usgs]

Thanks @jkellndorfer! I guess without that nothing gets propagated to the workers!

So if you add workers to the cluster you need to run that client.run() command again, right?

[jkellndorfer]

@rsignell-usgs Yes, I would think so, but have not tried that scenario adding new workers to a running cluster, and am not quite sure how that would work in autoscaling mode now that I think about it.

[rsignell-usgs]

Hopefully @tylerpotts @laisbsc or @costrouc will know how to pass env vars to the Dask cluster via the Gateway and we can avoid using thhe client.run() approach!

[dharhas]

So on an older version of qhub we had multiple sets of aws credentials and we had a couple of ways of handling this. But that was using dask_kubernetes and not dask_gateway so I'm not sure if those are relevant here. basically the storage_options kwarg could either be passed in the dask kubernetes cluster call or individually in the read_csv/read_parquet calls with the correct credentials.

[dharhas]

In another case we had several sets of named credentials in our file and we called s3fs with the correct name etc.

[dharhas]

@kcpevey how are we handling credentials on our client projects?

[kcpevey]

We are using a config file that the client package codebase reads on init and adds as env vars (using dotenv). Then we have some helper credentials functions that read from env vars:

def _load_aws_credentials():
    secret = os.environ.get("AWS_SECRET_ACCESS_KEY")
    key = os.environ.get("AWS_ACCESS_KEY_ID")

    if secret is None or key is None:
        logger.warning("Could not load AWS Secrets from environment!")
    return secret, key


def get_s3fs_client(**kwargs):
    secret, key = _load_aws_credentials()
    return s3fs.S3FileSystem(secret=secret, key=key, **kwargs)

# Different packages require different labeling so we duplicate them 
def get_aws_cluster_credentials():
    secret, key = _load_aws_credentials()
    return {
        "secret": secret,
        "key": key,
        "AWS_SECRET_ACCESS_KEY": secret,
        "AWS_ACCESS_KEY_ID": key,
    }

Then we pass these as a kwarg to the cluster:

from dask_kubernetes import KubeCluster
cluster = KubeCluster(silence_logs='error', env=get_aws_cluster_credentials())

I don't know what the differences would be between dask_kubernetes vs dask_gateway.

[rsignell-usgs]

Okay, this is embarrassing, as I totally forgot I already raised this as an issue a few months ago.

From that issue, it looks like it's possible to just pass the environment variables via cluster_options.environment for recent versions of Dask Gateway.

Since our Dask Gateway version doesn't yet allow that on our Qhub, we can use the approach by @jkellndorfer , or here's another approach, which copies the credentials from the ~/.aws directory to the workers via the cluster_workers using dask WorkerPlugin. This approach has the advantage that you can scale your cluster up without running anything again. @jsignell turned me onto this:

import os
from dask.distributed import WorkerPlugin

class UploadFile(WorkerPlugin):
    """A WorkerPlugin to upload a local file to workers.
    Parameters
    ----------
    filepath: str
        A path to the file to upload
    Examples
    --------
    >>> client.register_worker_plugin(UploadFile(".env"))
    """
    def __init__(self, filepath):
        """
        Initialize the plugin by reading in the data from the given file.
        """

        self.filename = os.path.basename(filepath)
        self.dirname = os.path.dirname(filepath)
        with open(filepath, "rb") as f:
            self.data = f.read()

    async def setup(self, worker):
        if not os.path.exists(self.dirname):
            os.mkdir(self.dirname)
        os.chdir(self.dirname)
        with open(self.filename, "wb+") as f:
            f.write(self.data)
        return os.listdir()

and then you can copy the credentials by doing:

client.register_worker_plugin(UploadFile('/home/jovyan/.aws/credentials'))
client.register_worker_plugin(UploadFile('/home/jovyan/.aws/config'))

[dharhas]

@aktech and @brl0 are working on a dask_gateway 0.9.0 update here: #315

hopefully we can merge it in soon.

[rsignell-usgs]

@dharhas, great! This WorkerPlugin could be useful in other ways too, of course -- it's a good tool to have in your Dask toolbox, and one I didn't have until this morning!

[brl0]

@rsignell-usgs,
We are also currently using the WorkerPlugin approach.
Here is the plugin:

class InitWorker(WorkerPlugin):
    name = "init_worker"

    def __init__(self, filepath=None, script=None):
        self.data = {}
        if filepath:
            if isinstance(filepath, str):
                filepath = [filepath]
            for file_ in filepath:
                with open(file_, "rb") as f:
                    filename = os.path.basename(file_)
                    self.data[filename] = f.read()
        if script:
            filename = f"{uuid.uuid1()}.py"
            self.data[filename] = script

    async def setup(self, worker):
        responses = await asyncio.gather(
            *[
                worker.upload_file(
                    comm=None, filename=filename, data=data, load=True
                )
                for filename, data in self.data.items()
            ]
        )
        assert all(
            len(data) == r["nbytes"]
            for r, data in zip(responses, self.data.values())
        )

The magic happens with load=True, which makes it possible to execute any script on the workers.

Here is a snippet of the script we are passing to the plugin:

script = f"""
import os
os.environ["AWS_ACCESS_KEY_ID"] = "{os.getenv("AWS_ACCESS_KEY_ID")}"
os.environ["AWS_SECRET_ACCESS_KEY"] = "{os.getenv("AWS_SECRET_ACCESS_KEY")}"
"""

And then to register the plugin:

plugin = InitWorker(script=script)
client.register_worker_plugin(plugin)

I am sure there are more elegant ways to do this and that this could be cleaned up and updated, but it shows that the approach works and provides flexibility.

[rsignell-usgs]

@brl0, very nice.

So now I first use the approach of @jkellndorfer to set the env vars from the config file, then set them for the workers using your approach.

So the whole thing looks like this:

import configparser
import os
def set_aws_credentials(cfile=os.path.join(os.environ['HOME'],'.aws','credentials'),profile_name='default',region_name='us-east-1',endpoint='s3.amazonaws.com',verbose=False):
    '''Sets the aws credentials if not set already and profilename is default'''
    cp = configparser.ConfigParser()
    cp.read(cfile)
    os.environ['aws_access_key_id'.upper()]=cp[profile_name]['aws_access_key_id']	
    os.environ['aws_secret_access_key'.upper()]=cp[profile_name]['aws_secret_access_key']	
    os.environ['aws_profile'.upper()]=profile_name
    os.environ['aws_default_profile'.upper()]=profile_name
    os.environ['aws_s3_region'.upper()]=region_name
    os.environ['aws_s3_endpoint'.upper()]=endpoint
    os.environ['aws_default_region'.upper()]=region_name
    if verbose:
        print('export {}={}'.format('aws_access_key_id'.upper(),cp[profile_name]['aws_access_key_id']	))
        print('export {}={}'.format('aws_secret_access_key'.upper(),cp[profile_name]['aws_secret_access_key']	))

set_aws_credentials(profile_name='esip-qhub')

from dask.distributed import WorkerPlugin
import os
import uuid
import asyncio

class InitWorker(WorkerPlugin):
    name = "init_worker"

    def __init__(self, filepath=None, script=None):
        self.data = {}
        if filepath:
            if isinstance(filepath, str):
                filepath = [filepath]
            for file_ in filepath:
                with open(file_, "rb") as f:
                    filename = os.path.basename(file_)
                    self.data[filename] = f.read()
        if script:
            filename = f"{uuid.uuid1()}.py"
            self.data[filename] = script

    async def setup(self, worker):
        responses = await asyncio.gather(
            *[
                worker.upload_file(
                    comm=None, filename=filename, data=data, load=True
                )
                for filename, data in self.data.items()
            ]
        )
        assert all(
            len(data) == r["nbytes"]
            for r, data in zip(responses, self.data.values())
        )

script = f"""
import os
os.environ["AWS_ACCESS_KEY_ID"] = "{os.getenv("AWS_ACCESS_KEY_ID")}"
os.environ["AWS_SECRET_ACCESS_KEY"] = "{os.getenv("AWS_SECRET_ACCESS_KEY")}"
os.environ["AWS_DEFAULT_REGION"] = "{os.getenv("AWS_DEFAULT_REGION")}"
"""

plugin = InitWorker(script=script)
client.register_worker_plugin(plugin)

[rsignell-usgs]

@dharhas just tipped me off that with #600, we now have a much easier way to pass environment variables to Dask workers. Hurrah!

@jkellndorfer we need to upgrade our ESIP Qhub to v0.3.12 to get this.

Question: @costrouc , @tylerpotts, @aktech, @danlester or @brl0 , can we just update the part of our config containing

default_images:
  jupyterhub: quansight/qhub-jupyterhub:v0.3.11
  jupyterlab: quansight/qhub-jupyterlab:v0.3.11
  dask_worker: quansight/qhub-dask-worker:v0.3.11
  dask_gateway: quansight/qhub-dask-gateway:v0.3.11

to upgrade from 0.3.11 to 0.3.12 and let github actions do their magic?

Or do we need to backup user data, redeploy, recover user data, etc?

[rsignell-usgs]

Okay, thanks @tylerpotts ! I look forward to hearing the best way to upgrade!

[costrouc]

Hey @rsignell-usgs these docs have a good explanation on how to upgrade https://docs.qhub.dev/en/latest/source/admin_guide/upgrade.html?highlight=upgrade

[rsignell-usgs]

@costrouc , great! So if I understand correctly...

If you are deploying via CI commit the two changes to qhub-config.yaml and the CI at the same time.

1. We can just do everything via github w/ github actions
2. We shouldn't just click the "edit icon" in github like we normally do
3. We don't need to backup user data

Right?

[tylerpotts]

@rsignell-usgs Update from the team:

We have a minor release that we plan to complete this week. The 3 major updates with this release will be:

• Cluster Monitoring with Grafana
• ClearML
• Improvements to Prefect integration

If any of the above are of interest, it may be worth waiting for the release.

As for your questions:

We can just do everything via github w/ github actions

Yep, that's right

We shouldn't just click the "edit icon" in github like we normally do

If you're referring to the edit button to change the code, that shouldn't cause a problem. I'd suggest making a branch and making changes to both files and then merge it in a single PR.

We don't need to backup user data

Backing up data isn't a required step to upgrade the cluster, but I would highly recommend it. At this stage in the project we don't test upgrades from one version to the next and it's not possible to guarantee the upgrade will go smoothly.

[rsignell-usgs]

Got it. Thanks Tyler! We will wait for the new release this week!

[rsignell-usgs]

Tried upgrading to 0.3.12 on a development deployment and failed