Add google Application Default Credentials to download

[fgerzer]

Description of changes:

#315 introduced GCS authentication for service accounts as an alternative to HMAC credentials. However, google.auth provides the google.auth.default function that uses Application Default Credentials to obtain credentials. See https://cloud.google.com/docs/authentication/provide-credentials-adc for more details. This does the following (from their docs):

    1. If the environment variable ``GOOGLE_APPLICATION_CREDENTIALS`` is set
       to the path of a valid service account JSON private key file, then it is
       loaded and returned. [...]
    2. If the `Google Cloud SDK`_ is installed and has application default
       credentials set they are loaded and returned. [...]
    3. If the application is running in the `App Engine standard environment`_
       (first generation) then the credentials and project ID from the
       `App Identity Service`_ are used.
    4. If the application is running in `Compute Engine`_ or `Cloud Run`_ or
       the `App Engine flexible environment`_ or the `App Engine standard
       environment`_ (second generation) then the credentials and project ID
       are obtained from the `Metadata Service`_.
    5. If no credentials are found,
       :class:`~google.auth.exceptions.DefaultCredentialsError` will be raised.

Point 1 was introduces in #315, so this is essentially a superset of the service account authentication introduced in that PR.

This is important for us because we use the Compute Engine credentials to avoid mounting secrets in the container which have unrestricted lifetime. I assume others are in the same situation as this is the recommended and more secure way.

Thanks to @b-chu for #315!

Backwards-Incompatible Change

There is one backwards-incompatible change in the order of HMAC/Service Account.

Previously, streaming prioritized a service account with an explicitly set GOOGLE_APPLICATION_CREDENTIALS environment variable over the HMAC credentials, i.e.

1. GOOGLE_APPLICATION_CREDENTIALS
2. HMAC

Now, this prioritizes HMAC over GOOGLE_APPLICATION_CREDENTIALS (the test has been changed correspondingly). This is because the default authentication includes explicit GOOGLE_APPLICATION_CREDENTIALS and we'd have to either (a) special-case outside of default authentication or (b) could not use HMAC if any of the default authentication methods in 2.-4. above were set. That means the order is

1. HMAC
2. GOOGLE_APPLICATION_CREDENTIALS
3. App Engine
4. Compute Engine
5. Raise an error

with 2.-4. provided by google.auth.default. This order allows us to explicitly set either HMAC or explicit service account credentials through environment variables, and use App Engine or Compute Engine if these are available.

I feel this backwards-incompatible change is acceptable because it only applies if both HMAC and GOOGLE_APPLICATION_CREDENTIALS are set at the same time. It also hasn't been part of any release yet, so we can expect adoption of GOOGLE_APPLICATION_CREDENTIALS isn't widespread.

If you disagree, I'll special-case it.

Issue #, if available:

Merge Checklist:

Put an x without space in the boxes that apply. If you are unsure about any checklist, please don't hesitate to ask. We are here to help! This is simply a reminder of what we are going to look for before merging your pull request.

General

• I have read the contributor guidelines
• This is a documentation change or typo fix. If so, skip the rest of this checklist.
• I certify that the changes I am introducing will be backward compatible, and I have discussed concerns about this, if any, with the MosaicML team.
  • See above.
• I have updated any necessary documentation, including README and API docs (if appropriate).
  • Note that the diff in the docs below is difficult to read because I moved the sections to have the new order. It looks like this:

Docs Screenshot

Tests

• I ran pre-commit on my change. (check out the pre-commit section of prerequisites)
• I have added tests that prove my fix is effective or that my feature works (if appropriate).
• I ran the tests locally to make sure it pass. (check out testing)
• I have added unit and/or integration tests as appropriate to ensure backward compatibility of the changes.

cc @fellhorn

[karan6181]

Thanks for the detailed PR description. It was super useful. And thanks for creating a PR to improve the user experience. The approach looks ok to me. Posted some minor comments.

Can you also share a snapshot that things works when user set HMAC and Service credentials (one at a time) ?

[fgerzer]

Can you also share a snapshot that things works when user set HMAC and Service credentials (one at a time) ?

I've used the script hidden below to demonstrate that; I'm mocking away gce credentials for demonstrating HMAC and only set HMAC credentials for HMAC. Note also the error if none of them are available. I'm running this on a google cloud instance, which has gce credentials available.

demo script

import os
import tempfile
import unittest.mock

import google.auth

from streaming import StreamingDataset

GCS_KEY = "YOURKEY"
GCS_SECRET = "YOUR_SECRET"


def available_methods() -> dict:
    available, unavailable = set(), set()
    if 'GCS_KEY' in os.environ and 'GCS_SECRET' in os.environ:
        available.add("HMAC")
    else:
        unavailable.add("HMAC")
    # Using the default google checkers here.
    checkers = {
        "GOOGLE_APPLICATION_CREDENTIALS": google.auth._default._get_explicit_environ_credentials(),
        "gcloud sdk credentials": google.auth._default._get_gcloud_sdk_credentials(),
        "gae credentials": google.auth._default._get_gae_credentials(),
        "gce credentials": google.auth._default._get_gce_credentials()
    }
    for name, result in checkers.items():
        if result[0] is None:
            unavailable.add(name)
        else:
            available.add(name)
    method_results = {"available": available, "unavailable": unavailable}
    return method_results


def read_one_sample():
    remote_dir = "YOUR_DIR"
    tmpdir = tempfile.mkdtemp()
    dataset = StreamingDataset(remote=remote_dir, local=tmpdir, split=None, shuffle=True)
    dl = dataset
    sample = next(iter(dl))
    print(f"{type(sample)=}")
    del dataset


def main():
    with unittest.mock.patch("google.auth._default._get_gce_credentials", lambda **kwargs: (None, None)):
        print("1. Using HMAC")
        os.environ['GCS_KEY'] = GCS_KEY
        os.environ['GCS_SECRET'] = GCS_SECRET
        print(available_methods())
        read_one_sample()
        del os.environ['GCS_KEY']
        del os.environ['GCS_SECRET']

    print("---------------------------\n\n")
    print("2. Using GCE credentials")
    print(available_methods())
    read_one_sample()
    print("---------------------------\n\n")

    print("3. No method available")
    with unittest.mock.patch("google.auth._default._get_gce_credentials", lambda *args, **kwargs: (None, None)):
        print(available_methods())
        read_one_sample()


if __name__ == "__main__":
    main()

[b-chu]

Thanks!